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

• hybrid_shard_degree (整数) — シャーディング並列処理の度数を指定します。値は 0～ の整数である必要がありますworld_size。デフォルト値は 0 です。

  • に設定すると0、ネイティブ実装にフォールバックし、 PyTorch tensor_parallel_degreeが 1 の場合はスクリプトAPIにフォールバックします。それ以外の場合は、 tensor_parallel_degreeと hybrid_shard_degreeに基づいて可能な限り最大の を計算しますworld_size。ネイティブ PyTorch FSDPユースケースにフォールバックする場合、 FULL_SHARDが使用する戦略であれば、 のクラスター全体にシャードされますGPUs。HYBRID_SHARD または _HYBRID_SHARD_ZERO2が戦略の場合、8 hybrid_shard_degreeに相当します。テンソル並列処理を有効にすると、改訂された に基づいてシャードされますhybrid_shard_degree。

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

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

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

• activation_loading_horizon (整数) — のアクティベーションオフロード期間タイプを指定する整数FSDP。これは、入力が同時にGPUメモリ内に存在できるチェックポイントレイヤーまたはオフロードレイヤーの最大数です。詳細については、「アクティベーションオフロード」を参照してください。入力値は正の整数である必要があります。デフォルト値は 2 です。

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

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

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

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

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

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

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

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

  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.sagemaker util 関数とプロパティ。

• verbose (ブール値) — 初期化と検証中により詳細なログ記録を有効にするかどうか。デフォルト値は False です。

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

• get_post_param_init_fn() – ラッパークラスの post_param_init_fn引 PyTorch FSDP数に渡すことができるパラメータ初期化関数を返します。これは、モデルで重みを結び付けている場合に必要です。モデルはメソッド を実装する必要があります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ない場合、 が発生します。このコンテキストマネージャーはオプションであり、このコンテキストマネージャーを使用してパラメータを初期化する必要はありません。

• state_dict (dict) - 必須。保存する状態 dict。

• checkpoint_id (str) - 必須。チェックポイントを保存するストレージパス。

• storage_writer (StorageWriter) - オプション。書き込みオペレーションを実行する PyTorch StorageWriterの インスタンス。これが指定されていない場合は、 のデフォルト設定StorageWriterが使用されます。

• planner (SavePlanner) - オプション。SavePlanner の インスタンス PyTorch。これが指定されていない場合は、 のデフォルト設定SavePlannerが使用されます。

• process_group (ProcessGroup) - オプション。作業するプロセスグループ。の場合None、デフォルトの (グローバル) プロセスグループが使用されます。

• coordinator_rank (int) - オプション。などの集合通信演算子を実行するときのコーディネーターのランクAllReduce。

• queue (AsyncRequestQueue) - オプション。使用する非同期スケジューラ。デフォルトでは、グローバルパラメータ が使用されますDEFAULT_ASYNC_REQUEST_QUEUE。

• sharded_strategy (PyTorchDistSaveShardedStrategy) - オプション。チェックポイントの保存に使用するシャード戦略。これが指定されていない場合、デフォルトで torch.sagemaker.distributed.checkpoint.state_dict_saver.PyTorchDistSaveShardedStrategy が使用されます。

• wait_error_handling (bool) - オプション。すべてのランクがエラー処理を完了するのを待つかどうかを指定するフラグ。デフォルト値は True です。

• force_check_all_plans (bool) - オプション。キャッシュヒットの場合でも、ランク間で計画を強制的に同期するかどうかを決定するフラグ。デフォルト値は True です。

• s3_region (str) - オプション。S3 バケットが配置されているリージョン。指定しない場合、リージョンは から推測されますcheckpoint_id。

• s3client_config (S3ClientConfig) - オプション。S3 クライアントの設定可能なパラメータを公開するデータクラス。指定しない場合、S3ClientConfig のデフォルト設定が使用されます。part_size パラメータはデフォルトで 64MB に設定されています。

• blocking (bool) - オプション。の場合True、アクティブなリクエストがすべて完了するまで待機します。それ以外の場合は、すでに完了した非同期リクエストのみを確定します。デフォルト値は True です。

• process_group (ProcessGroup) - オプション。操作するプロセスグループ。に設定するとNone、デフォルトの (グローバル) プロセスグループが使用されます。

• 非同期呼び出しのインデックスを含むリストは正常に確定されます。

• state_dict (dict) - 必須。保存する状態 dict。

• checkpoint_id (str) - 必須。チェックポイントを保存するストレージパス。

• storage_writer (StorageWriter) - オプション。書き込みオペレーションを実行する PyTorch StorageWriterの インスタンス。これを指定しない場合、 のデフォルト設定StorageWriterが使用されます。

• planner (SavePlanner) - オプション。SavePlanner の インスタンス PyTorch。これを指定しない場合、 のデフォルト設定SavePlannerが使用されます。

• process_group (ProcessGroup) - オプション。作業するプロセスグループ。の場合None、デフォルトの (グローバル) プロセスグループが使用されます。

• coordinator_rank (int) - オプション。などの集合通信演算子を実行するときのコーディネーターのランクAllReduce。

• wait_error_handling (bool) - オプション。すべてのランクがエラー処理を完了するのを待つかどうかを指定するフラグ。デフォルト値は True です。

• force_check_all_plans (bool) - オプション。キャッシュヒットの場合でも、ランク間で計画を強制的に同期するかどうかを決定するフラグ。デフォルト値は True です。

• s3_region (str) - オプション。S3 バケットが配置されているリージョン。指定しない場合、リージョンは から推測されますcheckpoint_id。

• s3client_config (S3ClientConfig) - オプション。S3 クライアントの設定可能なパラメータを公開するデータクラス。指定しない場合、S3ClientConfig のデフォルト設定が使用されます。part_size パラメータはデフォルトで 64MB に設定されています。

• state_dict (dict) - 必須。ロードstate_dictする 。

• checkpoint_id (str) - 必須。チェックポイントの ID。の意味はストレージcheckpoint_idによって異なります。フォルダまたはファイルへのパスを指定できます。ストレージがキーバリューストアである場合は、キーにすることもできます。

• storage_reader (StorageReader) - オプション。読み取りオペレーションを実行する PyTorch StorageReaderの インスタンス。指定しない場合、分散チェックポイントは に基づいてリーダーを自動的に推測しますcheckpoint_id。checkpoint_id も の場合None、例外エラーが発生します。

• planner (StorageReader) - オプション。LoadPlanner の インスタンス PyTorch。指定しない場合、 のデフォルト設定LoadPlannerが使用されます。

• check_keys_matched (bool) - オプション。有効にすると、すべてのランクのstate_dictキーが を使用して一致しているかどうかをチェックしますAllGather。

• s3_region (str) - オプション。S3 バケットが配置されているリージョン。指定しない場合、リージョンは から推測されますcheckpoint_id。

• s3client_config (S3ClientConfig) - オプション。S3 クライアントの設定可能なパラメータを公開するデータクラス。指定しない場合、S3ClientConfig のデフォルト設定が使用されます。part_size パラメータはデフォルトで 64MB に設定されています。

• smp_moe (ブール値) - MoE の SMP-implementation を使用するかどうか。デフォルト値は True です。

• random_seed (整数) - エキスパート並列分散モジュールのランダムオペレーションのシード番号。このシードはエキスパート並列ランクに追加され、各ランクの実際のシードを設定します。これは、エキスパートの並列ランクごとに一意です。デフォルト値は 12345 です。

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

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

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

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

• moe_z_loss_coeff (浮動小数点) - z 損失の係数。デフォルト値は 0.001 です。

• attention_dropout_prob (浮動小数点) — 注意に適用するドロップアウト確率。デフォルト値は 0.0 です。

• scale (浮動小数点) – 合格すると、このスケール係数がソフトマックスに適用されます。（デフォルト値でもある) に設定すると、スケール係数Noneは になります1 / sqrt(attention_head_size)。デフォルト値は None です。

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

• use_alibi (bool) – 渡されると、提供されたマスクを使用して線形バイアス (ALiBi) による注意が有効になります。を使用する場合はALiBi、次のように準備されたアテンションマスクが必要です。デフォルト値は 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 モジュール関数。module(x) が呼び出されると、 はこの関数を自動的にSMP実行します。

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

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

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

  • cast_dtype – 特定の に設定するとdtype、テンソルは qkv dtypeより前のテンソルにキャストされますattn。これは、Hugging Face Transformer GPT-NeoX モデルなどの実装に便利です。このモデルには、引き継ぎ埋め込みqkとfp32引き継ぎ埋め込みがあります。に設定するとNone、キャストは適用されません。デフォルト値は None です。

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

• attention_dropout_prob (浮動小数点) — 注意に適用するドロップアウト確率。デフォルト値は 0.0 です。

• scale (浮動小数点) – 合格すると、このスケール係数がソフトマックスに適用されます。に設定するとNone、 1 / sqrt(attention_head_size)がスケール係数として使用されます。デフォルト値は None です。

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

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

  • kv – torch.Tensor次の形式(batch_size x seqlen x (2 x num_heads) x head_size)または の (batch_size, (2 x num_heads) x seqlen x head_size)、または 2 つの のタプルtorch.Tensor。それぞれが (batch_size x seqlen x num_heads x head_size)または の形になります(batch_size x num_heads x seqlen x head_size)。また、シェイプに基づいて適切なlayout引数を渡す必要があります。

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

  • cast_dtype – 特定の dtype に設定すると、 qkv の前にテンソルをその dtype にキャストしますattn。これは、Hugging Face Transformers GPT-NeoX などの実装に便利です。これは、 fp32 q,kにポストアセンブルメント埋め込みを備えています。に設定するとNone、キャストは適用されません。デフォルト値は None です。

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

• config – Llama モデル FlashAttention の設定。

• 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 です。

• model (torch.nn.Module) – SMPライブラリのテンソル並列処理機能を変換して適用SMP テンソル並列処理と互換性のあるHugging Face Transformer モデルするための からのモデル。

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

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

• config (dict) – これはトランスSMPフォーマーを設定するためのディクショナリです。デフォルト値は None です。

• load_state_dict_from_rank0 (ブール値) – デフォルトでは、このモジュールは新しい重みを持つモデルの新しいインスタンスを作成します。この引数を に設定するとTrue、 SMP は元の PyTorch モデルの状態ディクショナリを 0 ランクから、0 ランクが属するテンソル並列グループの変換されたモデルにロードしようとします。これを に設定するとTrue、ランク 0 はメタデバイスにパラメータを持つことができません。この変換呼び出しの後、最初のテンソル並列グループのみが 0 ランクの重みを入力します。これらの重みを最初のテンソル並列グループから他のすべてのプロセスに取得するには、FSDPラッパーTrueで を sync_module_statesに設定する必要があります。有効にすると、SMPライブラリは元のモデルから状態ディクショナリをロードします。SMP ライブラリは変換前にstate_dictモデルの を受け取り、変換されたモデルの構造と一致するように変換し、テンソル並列ランクごとにシャードし、この状態を 0 番目のランクから 0 番目のランクが属するテンソル並列グループの他のランクに通信してロードします。デフォルト値は False です。

torch.sagemaker util 関数

• torch.sagemaker.init(config: Optional[Union[str, Dict[str, Any]]] = None) -> None – で PyTorch トレーニングジョブを初期化しますSMP。

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

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

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

プロパティ

torch.sagemaker.state は、 SMPで を初期化した後、複数の有用なプロパティを保持します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) – 他のディメンションで同じランクを持つすべてのデバイス (シャーディングデータ並列処理やレプリケーションなど) を含むテンソル並列プロセスグループ。ネイティブ にフォールバックすると PyTorch、 が返されますNone。

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