SageMaker モデルparallel ライブラリ v2 リファレンス - Amazon SageMaker

翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

SageMaker モデルparallel ライブラリ v2 リファレンス

以下は、 SageMaker モデルparallel ライブラリ v2 (SMP v2) のリファレンスです。

SMP v2 コア機能の構成パラメーター

をアクティブ化して設定するためのパラメータの完全なリストを以下に示します。 SageMaker モデル並列処理ライブラリ v2 のコア機能これらは JSON 形式で記述し、 SageMaker Python SDK PyTorch のエスティメーターに渡すか、または JSON ファイルとして保存する必要があります。 SageMaker HyperPod

{ "hybrid_shard_degree": Integer, "sm_activation_offloading": Boolean, "activation_loading_horizon": Integer, "fsdp_cache_flush_warnings": Boolean, "allow_empty_shards": Boolean, "tensor_parallel_degree": Integer, "expert_parallel_degree": Integer, "random_seed": Integer }
  • hybrid_shard_degree(Integer) — シャーディングされた並列度を指定します。値はとの間の整数でなければなりません。0 world_sizeデフォルト値は、0です。

    • に設定した場合0、が 1 PyTorch の場合はスクリプト内のネイティブ実装と API にフォールバックします。tensor_parallel_degreeそれ以外の場合は、hybrid_shard_degreetensor_parallel_degreeworld_sizeとに基づいて可能な限り最大の値を計算します。ネイティブの PyTorch FSDP ユースケースにフォールバックする場合、FULL_SHARD使用する戦略は GPU のクラスター全体に分散されます。HYBRID_SHARD_HYBRID_SHARD_ZERO2またはがストラテジーであれば、8 のストラテジーと同等です。hybrid_shard_degreeテンソル並列処理が有効になっている場合は、修正されたものに基づいて分割されます。hybrid_shard_degree

    • に設定すると1、が 1 PyTorch NO_SHARD のときはスクリプト内のネイティブ実装と API for にフォールバックします。tensor_parallel_degreeそれ以外の場合は、任意のテンソルparallel NO_SHARD グループ内と同等です。

    • 2 ~ の間の整数に設定するとworld_size、シャーディングは指定された数の GPU で行われます。FSDP sharding_strategy スクリプトで設定しなかった場合は、に上書きされます。HYBRID_SHARD設定すると_HYBRID_SHARD_ZERO2、指定したものが使用されます。sharding_strategy

  • sm_activation_offloading(Boolean) — SMP アクティベーションオフロード実装を有効にするかどうかを指定します。Falseオフロードがネイティブ実装を使用する場合。 PyTorch if True は SMP アクティベーションオフロード実装を使用します。また、 PyTorch スクリプトではアクティベーションオフロードラッパー (torch.distributed.algorithms._checkpoint.checkpoint_wrapper.offload_wrapper) を使用する必要があります。詳細については、「アクティベーションオフロード」を参照してください。デフォルト値は、Trueです。

  • activation_loading_horizon(Integer) — FSDP のアクティベーションオフロード範囲タイプを指定する整数。これは、入力を GPU メモリに同時に格納できる、チェックポイント付きまたはオフロードされたレイヤーの最大数です。詳細については、「アクティベーションオフロード」を参照してください。入力値は正の整数でなければなりません。デフォルト値は、2です。

  • fsdp_cache_flush_warnings(Boolean) — PyTorch メモリマネージャでキャッシュフラッシュが発生すると、計算パフォーマンスが低下する可能性があるため、それを検出して警告します。デフォルト値は、Trueです。

  • allow_empty_shards(Boolean) — テンソルが割り切れない場合にテンソルをシャーディングする際に空のシャードを許可するかどうか。これは、特定のシナリオにおけるチェックポイント処理中のクラッシュに対する実験的な修正です。これを無効にすると、元の動作に戻ります。 PyTorch デフォルト値は、Falseです。

  • tensor_parallel_degree(Integer) — テンソルの並列度を指定します。値はとの間でなければなりません。1 world_sizeデフォルト値は、1です。1 より大きい値を渡しても、テンソル並列処理は自動的に有効になりません。また、torch.sagemaker.transform API を使用してモデルをトレーニングスクリプトにラップする必要があります。詳細については、「テンソル並列性」を参照してください。

  • expert_parallel_degree(整数) — エキスパートの並列処理度を指定します。値は 1 ~ の範囲でなければなりません。world_sizeデフォルト値は、1です。1 より大きい値を渡しても、エキスパートによる並列処理は自動的に有効になりません。トレーニングスクリプトでは、必ず MoE モデルを torch.sagemaker.transform API でラップしてください。

  • random_seed(整数) — SMP テンソル並列処理またはエキスパート並列処理による分散モジュール内のランダム演算のシード番号。このシードはテンソル-パラレルランクまたはエキスパート-パラparallel ランクに追加され、各ランクの実際のシードが設定されます。テンソル-パラレルランクとエキスパート-パラparallel ランクごとに異なります。SMP v2は、テンソル-パラレルランクとエキスパート-パラparallel ランクで生成された乱数が、それぞれおよびのケースと一致することを確認します。 non-tensor-parallelism non-expert-parallelism

SMP v2 パッケージのリファレンス torch.sagemaker

このセクションは SMP v2 torch.sagemaker が提供するパッケージのリファレンスです。

torch.sagemaker.delayed_param.DelayedParamIniter

パラメータ初期化の遅延モデルに適用するための PyTorch API。

class torch.sagemaker.delayed_param.DelayedParamIniter( model: nn.Module, init_method_using_config : Callable = None, verbose: bool = False, )

パラメータ

  • model(nn.Module) — SMP v2 PyTorch の遅延パラメータ初期化機能をラップして適用するモデル。

  • init_method_using_config(呼び出し可能) — SMP v2 またはサポートされているテンソルparallel 実装を使用している場合は、このパラメーターをデフォルト値SMPテンソル並列処理に対応したHugging Face ストランスフォーマーモデル、つまりのままにしてください。Noneデフォルトでは、DelayedParamIniter API は指定されたモデルを正しく初期化する方法を見つけます。その他のモデルでは、カスタムパラメータ初期化関数を作成してスクリプトに追加する必要があります。次のコードスニペットは、SMP v2 init_method_using_config がに実装したデフォルト関数です。SMPテンソル並列処理に対応したHugging Face ストランスフォーマーモデル次のコードスニペットを参考にして、独自の初期化設定関数を作成し、スクリプトに追加して SMP API init_method_using_config のパラメータに渡してください。DelayedParamIniter

    from torch.sagemaker.utils.module_utils import empty_module_params, move_buffers_to_device # Define a custom init config function. def custom_init_method_using_config(module): d = torch.cuda.current_device() empty_module_params(module, device=d) if isinstance(module, (nn.Linear, Conv1D)): module.weight.data.normal_(mean=0.0, std=config.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=config.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.weight.data.fill_(1.0) module.bias.data.zero_() elif isinstance(module, LlamaRMSNorm): module.weight.data.fill_(1.0) move_buffers_to_device(module, device=d) delayed_initer = DelayedParamIniter(model, init_method_using_config=custom_init_method_using_config)

    torch.sagemaker.module_util前述のコードスニペット内の関数の詳細については、を参照してください。torch.sagemakerutil () 関数とプロパティ

  • verbose(Boolean) — 初期化および検証時に、より詳細なロギングを有効にするかどうか。デフォルト値は、Falseです。

方法

  • get_param_init_fn()— FSDP param_init_fn ラッパークラスの引数に渡すことができるパラメータ初期化関数を返します。 PyTorch

  • get_post_param_init_fn()— FSDP post_param_init_fn ラッパークラスの引数に渡すことができるパラメータ初期化関数を返します。 PyTorch これは、モデル内のウェイトを関連付けた場合に必要です。tie_weightsモデルにはメソッドを実装する必要があります。詳細については、同体重に関する注記を参照してくださいパラメータ初期化の遅延

  • count_num_params(module: nn.Module, *args: Tuple[nn.Parameter]) — パラメータ初期化関数によって初期化されているパラメータの数を追跡します。これにより、validate_params_and_buffers_inited以下のメソッドを実装しやすくなります。validate_params_and_buffers_initedこのメソッドはバックエンドで暗黙的にこのメソッドを呼び出すため、通常はこの関数を明示的に呼び出す必要はありません。

  • validate_params_and_buffers_inited(enabled: bool=True) — これは、初期化されたパラメーターの数がモデル内のパラメーターの総数と一致することを検証するのに役立つコンテキストマネージャーです。また、すべてのパラメーターとバッファーがメタデバイスではなく GPU デバイス上にあることも検証します。AssertionErrorsこれらの条件が満たされない場合に発生します。このコンテキストマネージャーはオプションであり、パラメーターの初期化にこのコンテキストマネージャーを使用する必要はありません。

torch.sagemaker.moe.moe_config.MoEConfig

SMP による混合エキスパート (MoE) 実装をセットアップするための設定クラスです。このクラスで MoE 設定値を指定し、API 呼び出しに渡すことができます。torch.sagemaker.transformこのクラスを MoE モデルのトレーニングに使用する方法の詳細については、を参照してくださいエキスパートによる並列処理

class torch.sagemaker.moe.moe_config.MoEConfig( smp_moe=True, random_seed=12345, moe_load_balancing="sinkhorn", global_token_shuffle=False, moe_all_to_all_dispatcher=True, moe_aux_loss_coeff=0.001, moe_z_loss_coeff=0.001 )
  • smp_moe(ブール値)-MoE の SMP 実装を使用するかどうか。デフォルト値は、Trueです。

  • random_seed(整数)-エキスパートparallel 分散モジュールにおけるランダム演算のシード番号。このシードはエキスパートparallel ランクに追加され、各ランクの実際のシードが設定されます。エキスパートのparallel ランクごとに異なります。デフォルト値は、12345です。

  • moe_load_balancing(文字列)-MoE ルーターの負荷分散タイプを指定します。有効なオプションはaux_loss、、sinkhornbalancednoneおよびです。デフォルト値は、sinkhornです。

  • global_token_shuffle(Boolean)-同じ EP グループ内の EP ランク間でトークンをシャッフルするかどうか。デフォルト値は、Falseです。

  • moe_all_to_all_dispatcher(ブール値)-MoE all-to-all での通信にディスパッチャーを使用するかどうか。デフォルト値は、Trueです。

  • moe_aux_loss_coeff(Float)-補助負荷分散損失の係数。デフォルト値は、0.001です。

  • moe_z_loss_coeff(フロート)-Z 損失の係数。デフォルト値は、0.001です。

torch.sagemaker.nn.attn.FlashSelfAttention

SMP v2 で使用するための API FlashAttention

class torch.sagemaker.nn.attn.FlashSelfAttention( attention_dropout_prob: float = 0.0, scale: Optional[float] = None, triton_flash_attention: bool = False, use_alibi: bool = False, )

パラメータ

  • attention_dropout_prob(float) — 注意が必要な場合に適用されるドロップアウト確率。デフォルト値は、0.0です。

  • scale(float) — 渡されると、このスケール係数がソフトマックスに適用されます。に設定した場合 None (これもデフォルト値)、スケール係数はです。1 / sqrt(attention_head_size)デフォルト値は、Noneです。

  • triton_flash_attention(bool) — 渡されると、フラッシュアテンションの Triton 実装が使用されます。これは線形バイアス (ALiBi) によるアテンションをサポートするために必要です (use_alibi以下のパラメーターを参照)。このバージョンのカーネルはドロップアウトをサポートしていません。デフォルト値は、Falseです。

  • use_alibi(bool) — 合格すると、指定されたマスクを使用して線形バイアスによる注意 (ALiBi) が有効になります。A を使用する場合はLiBi、以下のようにアテンションマスクを用意する必要があります。デフォルト値は、Falseです。

    def generate_alibi_attn_mask(attention_mask, batch_size, seq_length, num_attention_heads, alibi_bias_max=8): device, dtype = attention_mask.device, attention_mask.dtype alibi_attention_mask = torch.zeros( 1, num_attention_heads, 1, seq_length, dtype=dtype, device=device ) alibi_bias = torch.arange(1 - seq_length, 1, dtype=dtype, device=device).view( 1, 1, 1, seq_length ) m = torch.arange(1, num_attention_heads + 1, dtype=dtype, device=device) m.mul_(alibi_bias_max / num_attention_heads) alibi_bias = alibi_bias * (1.0 / (2 ** m.view(1, num_attention_heads, 1, 1))) alibi_attention_mask.add_(alibi_bias) alibi_attention_mask = alibi_attention_mask[..., :seq_length, :seq_length] if attention_mask is not None and attention_mask.bool().any(): alibi_attention_mask.masked_fill( attention_mask.bool().view(batch_size, 1, 1, seq_length), float("-inf") ) return alibi_attention_mask

方法

  • forward(self, qkv, attn_mask=None, causal=False, cast_dtype=None, layout="b h s d")— PyTorch 通常のモジュール関数。a module(x) が呼び出されると、SMP はこの関数を自動的に実行します。

    • qkvtorch.Tensor 次の形式です。(batch_size x seqlen x (3 x num_heads) x head_size)または(batch_size, (3 x num_heads) x seqlen x head_size)torch.Tensorsそれぞれの形が(batch_size x seqlen x num_heads x head_size)、またはの形のタプル。(batch_size x num_heads x seqlen x head_size)形状に基づいて適切なレイアウト引数を渡す必要があります。

    • attn_masktorch.Tensor (batch_size x 1 x 1 x seqlen) 以下の形式の。このアテンションマスクパラメータを有効にするには、triton_flash_attention=Trueとが必要ですuse_alibi=True。このメソッドを使用してアテンションマスクを生成する方法については、にあるコード例を参照してくださいFlashAttention。デフォルト値は、Noneです。

    • causalFalse 引数のデフォルト値であるに設定すると、マスクは適用されません。に設定するとTrueforwardメソッドは標準の下三角マスクを使用します。デフォルト値は、Falseです。

    • cast_dtype— 特定の値に設定するとdtypeqkvテンソルは前のものにキャストされます。dtype attnこれは、Hugging Face Transformer GPT-Neox モデルなど、q回転埋め込みの有無にかかわらず実装する場合に役立ちます。k fp32に設定すると、キャストは適用されません。Noneデフォルト値は、Noneです。

    • layout(文字列) — b h s d b s h d 使用可能な値はまたはです。適切な変換を適用できるように、qkv渡されたテンソルのレイアウトに設定する必要があります。attnデフォルト値は、b h s dです。

戻り値

形のある単一torch.Tensor(batch_size x num_heads x seq_len x head_size)

torch.sagemaker.nn.attn.FlashGroupedQueryAttention

SMP v2 FlashGroupedQueryAttention で使用するための API です。この API の使用方法の詳細については、を参照してください。 FlashAttention グループ化されたクエリの注意にはカーネルを使用してください。

class torch.sagemaker.nn.attn.FlashGroupedQueryAttention( attention_dropout_prob: float = 0.0, scale: Optional[float] = None, )

パラメータ

  • attention_dropout_prob(float) — 注意が必要な場合に適用できるドロップアウト確率。デフォルト値は、0.0です。

  • scale(float) — 渡されると、このスケール係数がソフトマックスに適用されます。に設定するとNone1 / sqrt(attention_head_size)スケール係数として使用されます。デフォルト値は、Noneです。

方法

  • forward(self, q, kv, causal=False, cast_dtype=None, layout="b s h d")— PyTorch 通常のモジュール関数。a module(x) が呼び出されると、SMP はこの関数を自動的に実行します。

    • qtorch.Tensor (batch_size x seqlen x num_heads x head_size) 以下の形式または(batch_size x num_heads x seqlen x head_size). 形状に基づいて適切なレイアウト引数を渡す必要があります。

    • kvtorch.Tensor 以下の形式の (batch_size x seqlen x (2 x num_heads) x head_size) or (batch_size, (2 x num_heads) x seqlen x head_size)、または 2 つの torch.Tensor s のタプル。(batch_size x seqlen x num_heads x head_size)それぞれがまたはの形でもかまいません。(batch_size x num_heads x seqlen x head_size)また、layout形状に応じて適切な引数を渡さなければなりません。

    • causalFalse 引数のデフォルト値であるに設定すると、マスクは適用されません。に設定するとTrueforwardメソッドは標準の下三角マスクを使用します。デフォルト値は、Falseです。

    • cast_dtype— 特定の dtype に設定すると、qkvテンソルをその前にその dtype にキャストします。attnこれは、Hugging Face ストランスフォーマー GPT-Neox などの実装に便利です。例えば、回転式埋め込み機能付きの実装です。q,k fp32に設定すると、キャストは適用されません。Noneデフォルト値は、Noneです。

    • レイアウト (文字列) — "b h s d" "b s h d" 使用可能な値はまたはです。適切な変換を適用できるように、qkv渡されたテンソルのレイアウトに設定する必要があります。attnデフォルト値は、"b h s d"です。

戻り値

torch.Tensor (batch_size x num_heads x seq_len x head_size)アテンション計算の出力を表す単精度浮動小数点を返します。

torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention

Llama FlashAttention モデルをサポートする API。この API は torch.sagemaker.nn.attn.FlashGroupedQueryAttention API を低レベルで使用します。この使用方法については、「」を参照してください FlashAttention グループ化されたクエリの注意にはカーネルを使用してください。

class torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention( config: LlamaConfig )

パラメータ

  • config— FlashAttention Llamaモデル用の設定。

方法

  • forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache)

    • hidden_states(torch.Tensor) — という形式のテンソルの隠れ状態 (batch_size x seq_len x num_heads x head_size)

    • attention_mask(torch.LongTensor) — という形式のパディングトークンインデックスに注意が向けられないようにするためのマスク。(batch_size x seqlen)デフォルト値は、Noneです。

    • position_ids(torch.LongTensor) — 存在しない場合はNone(batch_size x seqlen)位置埋め込みにおける各入力シーケンストークンの位置のインデックスを示す形式になります。デフォルト値は、Noneです。

    • past_key_value(キャッシュ) — 事前に計算された隠れ状態 (セルフアテンションブロックとクロスアテンションブロック内のキーと値)。デフォルト値は、Noneです。

    • output_attentions(bool) — すべてのアテンションレイヤーのアテンションテンソルを返すかどうかを示します。デフォルト値は、Falseです。

    • use_cache(bool) — past_key_values キー値の状態を返すかどうかを示します。デフォルト値は、Falseです。

戻り値

torch.Tensor (batch_size x num_heads x seq_len x head_size)アテンション計算の出力を表すシングルを返します。

torch.sagemaker.transform

SMP v2 には、Hugging Face ストランスフォーマーモデルを SMP モデル実装に変換し、SMP テンソル並列処理を有効にするためのこの torch.sagemaker.transform() API が用意されています。

torch.sagemaker.transform( model: nn.Module, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, config: Optional[Dict] = None, load_state_dict_from_rank0: bool = False )

SMP v2 は、Hugging Face Transformer モデルの設定を SMP SMPテンソル並列処理に対応したHugging Face ストランスフォーマーモデル トランスフォーマー設定に変換することで、のトランスフォーメーションポリシーを維持しています。

パラメータ

  • model(torch.nn.Module) — SMP SMPテンソル並列処理に対応したHugging Face ストランスフォーマーモデル ライブラリのテンソル並列処理機能を変換して適用するためのモデル。

  • device(torch.device) — 渡されると、このデバイス上に新しいモデルが作成されます。元のモジュールにメタデバイス上のパラメータがある場合 (「」を参照パラメータ初期化の遅延)、変換後のモジュールもメタデバイス上に作成され、ここで渡された引数は無視されます。デフォルト値は、Noneです。

  • dtype(torch.dtype) — 渡された場合、これをモデル作成用の dtype コンテキストマネージャーとして設定し、この dtype でモデルを作成します。fp32使用時にモデルを作成したいので、通常は不要でMixedPrecision、デフォルトの dtype fp32 になっています。 PyTorchデフォルト値は、Noneです。

  • config(dict) — SMP トランスフォーマーを設定するための辞書です。デフォルト値は、Noneです。

  • load_state_dict_from_rank0(Boolean) — デフォルトでは、このモジュールは新しい重みでモデルの新しいインスタンスを作成します。この引数がに設定されている場合True、 PyTorch SMPは元のモデルの状態辞書を0番目のランクから0番目のランクが属するテンソルparallel グループの変換モデルに読み込もうとします。これを設定するとTrue、ランク 0 はメタデバイスのパラメータを持つことができなくなります。この変換呼び出しの後、最初のテンソルparallel グループのみが0番目のランクから重みを設定します。最初のテンソルparallel sync_module_states True グループから他のすべてのプロセスにこれらの重みを取得するには、FSDP ラッパーで to を設定する必要があります。これを有効にすると、SMP ライブラリは元のモデルから状態辞書を読み込みます。SMPライブラリはstate_dict、変換前のモデルを取得し、変換されたモデルの構造と一致するように変換し、テンソルparallel ランクごとに分割し、この状態を0番目のランクから0番目のランクが属するテンソルparallel グループの他のランクに伝え、ロードします。デフォルト値は、Falseです。

を返します。

PyTorch FSDP でラップできる変換済みモデルを返します。をに設定するとload_state_dict_from_rank0True、ランク 0 を含むテンソルparallel グループの重みは、元の状態ディクショナリからランク 0 にロードされます。パラメータ初期化の遅延元のモデルで使用する場合、変換後のモデルのパラメーターとバッファー用の CPU 上の実際のテンソルがあるのはこれらのランクだけです。残りのランクでは、メモリを節約するために引き続きパラメータとバッファがメタデバイス上に残ります。

torch.sagemakerutil () 関数とプロパティ

トーチ.sagemaker ユーティリティ関数
  • torch.sagemaker.init(config: Optional[Union[str, Dict[str, Any]]] = None) -> None— SMP を使用してトレーニングジョブを初期化します。 PyTorch

  • torch.sagemaker.is_initialized() -> bool— トレーニングジョブが SMP で初期化されているかどうかを確認します。ジョブが SMP PyTorch で初期化されている間にネイティブにフォールバックすると、一部のプロパティは関連性がなくなりNone、以下のプロパティリストに示すようになります。

  • torch.sagemaker.utils.module_utils.empty_module_params(module: nn.Module, device: Optional[torch.device] = None, recurse: bool = False) -> nn.Moduledevice 指定したパラメータがあればそのパラメータに空のパラメータを作成し、指定すればすべてのネストされたモジュールに対して再帰的に指定できます。

  • torch.sagemaker.utils.module_utils.move_buffers_to_device(module: nn.Module, device: torch.device, recurse: bool = False) -> nn.Module— モジュールバッファを指定された場所に移動します。指定するとdevice、ネストされたすべてのモジュールに対して再帰的に指定できます。

プロパティ

torch.sagemaker.stateSMP with の初期化後に、複数の便利なプロパティを保持します。torch.sagemaker.init

  • torch.sagemaker.state.hybrid_shard_degree(int) — シャードデータの並列処理度。渡された SMP 設定のユーザー入力からのコピー。torch.sagemaker.init()詳細については、「 SageMaker モデル並列処理ライブラリ v2 を使い始める」を参照してください。

  • torch.sagemaker.state.rank(int) — デバイスのグローバルランク(の範囲内)。[0, world_size)

  • torch.sagemaker.state.rep_rank_process_group(torch.distributed.ProcessGroup) — 同じレプリケーションランクを持つすべてのデバイスを含むプロセスグループ。torch.sagemaker.state.tp_process_groupとの微妙ではあるが根本的な違いに注目してください。 PyTorchネイティブにフォールバックすると戻りますNone

  • torch.sagemaker.state.tensor_parallel_degree(int) — テンソルの並列処理度。SMP構成のユーザー入力からのコピーが渡されます。torch.sagemaker.init()詳細については、「 SageMaker モデル並列処理ライブラリ v2 を使い始める」を参照してください。

  • torch.sagemaker.state.tp_size(int) — へのエイリアス。torch.sagemaker.state.tensor_parallel_degree

  • torch.sagemaker.state.tp_rank(int) — テンソル並列度とランキングメカニズムによって決まる[0, tp_size)、範囲内のデバイスのテンソル並列度ランク。

  • torch.sagemaker.state.tp_process_group(torch.distributed.ProcessGroup) — 他の次元(シャーディングされたデータのparallel 処理と複製など)では同じランクであるが、テンソル並列ランクは異なるすべてのデバイスを含むテンソルparallel プロセスグループ。ネイティブにフォールバックすると戻ります。 PyTorch None

  • torch.sagemaker.state.world_size(int) — トレーニングに使用されたデバイスの総数。

SMP v1 から SMP v2 へのアップグレード

SMP v1 から SMP v2 に移行するには、スクリプトを変更して SMP v1 API を削除し、SMP v2 API を適用する必要があります。SMP v1 スクリプトから始めるのではなく、FSDP スクリプトから開始し、の指示に従うことをお勧めします。 PyTorch SageMaker モデル並列処理ライブラリ v2 を使い始める

SMP v1 モデルを SMP v2 に移行するには、SMP v1 ではモデル状態ディクショナリ全体を収集し、モデルステートディクショナリ上の変換関数を適用して Hugging Face Transformers モデルチェックポイント形式に変換する必要があります。これで、で説明したように、SMP v2 では SMP 使用中のチェックポイントの保存と読み込み Hugging Face Transformers モデルのチェックポイントを読み込んで、SMP v2 でもチェックポイント API を引き続き使用できます。 PyTorch PyTorch FSDP モデルで SMP を使用するには、必ず SMP v2 に移行し、FSDP やその他の最新機能を使用するようにトレーニングスクリプトを変更してください。 PyTorch

import smdistributed.modelparallel.torch as smp # Create model model = ... model = smp.DistributedModel(model) # Run training ... # Save v1 full checkpoint if smp.rdp_rank() == 0: model_dict = model.state_dict(gather_to_rank0=True) # save the full model # Get the corresponding translation function in smp v1 and translate if model_type == "gpt_neox": from smdistributed.modelparallel.torch.nn.huggingface.gptneox import translate_state_dict_to_hf_gptneox translated_state_dict = translate_state_dict_to_hf_gptneox(state_dict, max_seq_len=None) # Save the checkpoint checkpoint_path = "checkpoint.pt" if smp.rank() == 0: smp.save( {"model_state_dict": translated_state_dict}, checkpoint_path, partial=False, )

SMP v1 で使用できる翻訳機能については、を参照してください。Hugging Face Transformer モデルのサポート

SMP v2 でのモデルチェックポイントの保存と読み込みの手順については、を参照してください。SMP 使用中のチェックポイントの保存と読み込み