SageMaker 模型平行程式庫 v2 參考

• hybrid_shard_degree （整數） – 指定碎片平行處理程度。值必須是介於 0和 之間的整數world_size。預設值為 0。

  • 如果設定為 0，它會回到原生 PyTorch實作，並在 tensor_parallel_degree 為 1 時在指令碼API中。否則，它會hybrid_shard_degree根據 tensor_parallel_degree和 運算最大的可能world_size。回到原生 PyTorch FSDP使用案例時，如果 FULL_SHARD是您使用的策略，它會碎片分散到整個 叢集GPUs。如果 HYBRID_SHARD或 _HYBRID_SHARD_ZERO2是策略，則相當於 hybrid_shard_degree 8。啟用張量平行處理時，它會根據修訂的 碎片hybrid_shard_degree。

  • 如果設定為 1，它會回到原生 PyTorch實作，而當 tensor_parallel_degree為 1 時，則會在指令碼NO_SHARD中API針對 。否則，它相當於在任何指定的張量平行群組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.transformAPI來包裝模型。如需進一步了解，請參閱 內容平行處理。

• expert_parallel_degree （整數） – 指定專家平行程度。值必須介於 1 和 之間world_size。預設值為 1。請注意，傳遞大於 1 的值不會自動啟用內容平行處理；您也需要使用 torch.sagemaker.transformAPI來在訓練指令碼中包裝模型。如需進一步了解，請參閱 專家平行處理。

• random_seed （整數） – 透過SMP張量平行處理或專家平行處理，在分散式模組中隨機操作的種子編號。此種子會新增至張量平行或專家平行排名，以設定每個排名的實際種子。每個張量平行和專家平行排名都是唯一的。SMPv2 可確保在張量平行和專家平行排名之間產生的隨機數字分別符合 non-tensor-parallelism和 non-expert-parallelism案例。

• model (nn.Module) – 包裝和套用 SMP v2 延遲參數初始化功能的 PyTorch 模型。

• init_method_using_config （可呼叫） – 如果您使用 SMP v2 的張量平行實作或支援的 與SMP張量平行處理相容的擁抱臉部轉換器模型，請將此參數保持在預設值，即 None。根據預設， DelayedParamIniter API 會了解如何正確初始化指定的模型。對於任何其他模型，您需要建立自訂參數初始化函數，並將其新增至指令碼。下列程式碼片段是 SMP v2 針對 實作的預設init_method_using_config函數與SMP張量平行處理相容的擁抱臉部轉換器模型。使用下列程式碼片段做為建立您自己的初始化組態函數、將其新增至指令碼，並將其傳遞至 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() – 傳回參數初始化函數，您可以傳遞給包裝函式 PyTorch FSDP類別的post_param_init_fn引數。當您在模型中綁定權重時，這是必要的。模型必須實作方法 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 （裁剪） - 必要。要儲存的狀態索引。

• checkpoint_id (str) - 必要。儲存檢查點的儲存路徑。

• storage_writer (StorageWriter) - 選用。StorageWriter 中的執行個體， PyTorch 以執行寫入操作。如果未指定，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 （布爾） - 選用。指定是否等待所有排名完成錯誤處理的旗標。預設值為 True。

• force_check_all_plans （布爾） - 選用。一個旗標，可決定是否強制同步跨排名的計劃，即使在快取命中的情況下也是如此。預設值為 True。

• s3_region (str) - 選用。S3 儲存貯體所在的區域。如果未指定，則會從 推斷區域checkpoint_id。

• s3client_config (S3ClientConfig) - 選用。資料類別會公開 S3 用戶端的可設定參數。如果未提供，則會使用 S3ClientConfig 的預設組態。part_size 參數預設為 64MB。

• blocking （布爾） - 選用。如果 True，它將等待所有作用中的請求完成。否則，它只會完成已完成的非同步請求。預設值為 True。

• process_group (ProcessGroup) - 選用。要操作的程序群組。如果設定為 None，則會使用預設 （全域） 程序群組。

• 包含非同步呼叫索引的清單已成功完成。

• state_dict （裁剪） - 必要。要儲存的狀態索引。

• checkpoint_id (str) - 必要。儲存檢查點的儲存路徑。

• storage_writer (StorageWriter) - 選用。StorageWriter 中的執行個體， PyTorch 用於執行寫入操作。如果未指定，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 （裁剪） - 必要。state_dict 要載入的 。

• checkpoint_id (str) - 必要。檢查點的 ID。的意義checkpoint_id取決於儲存體。它可以是資料夾或 檔案的路徑。如果儲存體是金鑰值存放區，它也可以是金鑰。

• storage_reader (StorageReader) - 選用。StorageReader 中的執行個體， PyTorch 以執行讀取操作。如果未指定，分散式檢查點會根據 自動推斷讀取器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實作。預設值為 True。

• random_seed （整數） - 專家平行分散式模組中隨機操作的種子號碼。此種子會新增至專家平行排名，以設定每個排名的實際種子。每個專家平行排名都是唯一的。預設值為 12345。

• moe_load_balancing （字串） - 指定 MoE 路由器的負載平衡類型。有效選項為 aux_loss、balanced、 sinkhorn和 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-loss 的係數。預設值為 0.001。

• attention_dropout_prob （浮點數） – 套用至注意力的退出機率。預設值為 0.0。

• scale （浮點數） – 如果傳遞，此縮放係數會套用至 softmax。如果設定為 None（這也是預設值），則縮放係數為 1 / sqrt(attention_head_size)。預設值為 None。

• triton_flash_attention （布爾） – 如果通過，則會使用 Triton 快速注意實作。這是支援注意力與線性偏誤 (ALiBi) 的必要條件 （請參閱下列use_alibi參數）。此版本的核心不支援退出。預設值為 False。

• use_alibi （布林值） – 如果傳遞，則會使用提供的遮罩啟用具有線性偏誤的注意力 (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)或 (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)。必須根據形狀傳遞適當的配置 arg。

  • 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 模型等實作很有用，這些模型在旋轉嵌入fp32後具有 k q和 。如果設定為 None，則不會套用任何轉換。預設值為 None。

  • layout （字串） – 可用的值為 b h s d或 b s h d。這應該設定為傳遞qkv張量的配置，以便將適當的轉換套用至 attn。預設值為 b h s d。

• attention_dropout_prob （浮點數） – 套用至注意力的退出機率。預設值為 0.0。

• scale （浮點數） – 如果傳遞，此縮放係數會套用至 softmax。如果設定為 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)。必須根據形狀傳遞適當的配置 arg。

  • 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)，或兩個 的組合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 之前將張量轉換為該 dtypeattn。這對於 Hugging Face Transformer 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 （布爾） – 指示是否傳回所有注意力層的注意力張量。預設值為 False。

  • use_cache (bool) – 指示是否要傳回past_key_values金鑰值狀態。預設值為 False。

• model (torch.nn.Module) – 從 與SMP張量平行處理相容的擁抱臉部轉換器模型轉換和套用程式SMP庫張量平行處理的模型。

• device (torch.device) – 如果傳遞，則會在此裝置上建立新的模型。如果原始模組在中繼裝置上有任何參數 （請參閱延遲參數初始化)，則轉換後的模組也會在中繼裝置上建立，忽略傳遞到這裡的引數。預設值為 None。

• dtype (torch.dtype) – 如果傳遞，請將此設定為建立模型的 dtype 內容管理員，並使用此 dtype 建立模型。這通常是不必要的，因為我們想要在使用 fp32時使用 建立模型MixedPrecision， fp32是預設的 dtype 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。

• cp_comm_type (str) – 判斷內容平行處理實作，且僅適用於 context_parallel_degree 大於 1 時。此參數的可用值為 p2p和 all_gather。p2p 實作會在注意力運算期間利用 (KV) 張量累積的 peer-to-peer key-and-value收發呼叫，以非同步方式執行，並允許通訊與運算重疊。另一方面，all_gather實作採用 KV 張量累積的AllGather通訊集體操作。預設值為 "p2p"。

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) – 訓練中使用的裝置總數。