SageMaker 模型并行库 v2 参考

• hybrid_shard_degree（整数）-指定分片并行度。该值必须是介于0和之间的整数world_size。默认值为 0。

  • 如果设置为0，则回退到本机 PyTorch实现，而脚本API中则tensor_parallel_degree为 1。否则，它将hybrid_shard_degree根据tensor_parallel_degree和world_size计算可能的最大值。当回退到原生 PyTorch FSDP用例时，如果你使用的策略FULL_SHARD是的话，它会在整个集群中进行分片。GPUs如果_HYBRID_SHARD_ZERO2是HYBRID_SHARD或曾经是策略，则相当于 hybrid_shard_degree 8。当启用张量并行度时，它会根据修订后的分片。hybrid_shard_degree

  • 如果设置为1，则回退到本机 PyTorch实现，而在脚本NO_SHARD中，当设置APItensor_parallel_degree为 1 时，它将回退到本机实现。否则，它等同于NO_SHARD在任何给定的张量 parallel 组中。

  • 如果设置为介于 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 的张量并行实现或支持Hugging Face Transformer 模型与张量并SMP行度兼容，请将此参数保持为默认值，即。None默认情况下，会DelayedParamIniterAPI找出如何正确初始化给定模型。对于任何其他模型，您需要创建一个自定义参数初始化函数并将其添加到脚本中。以下代码片段是 SMP v2 为实现的默认init_method_using_config函数。Hugging Face Transformer 模型与张量并SMP行度兼容使用以下代码片段作为参考，您可以创建自己的初始化配置函数，将其添加到脚本中，然后将其传递给的init_method_using_configSMPDelayedParamIniterAPI参数。

  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（布尔值）-是否在初始化和验证期间启用更详细的日志记录。默认值为 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)-可选。in 的实例 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（bool）-可选。一个标志，用于指定是否等待所有等级完成错误处理。默认值为 True。

• force_check_all_plans（bool）-可选。一个标志，用于确定是否强制跨等级同步计划，即使在缓存命中时也是如此。默认值为 True。

• s3_region(str)-可选。S3 存储桶所在的区域。如果未指定，则根据推断出该checkpoint_id区域。

• s3client_config(S3ClientConfig)-可选。为 S3 客户端公开可配置参数的数据类。如果未提供，则使用 S3 ClientConfig 的默认配置。默认情况下，该part_size参数设置为 64MB。

• blocking（bool）-可选。如果是True，它将等到所有活动请求完成。否则，它只会最终确定已经完成的异步请求。默认值为 True。

• process_group(ProcessGroup)-可选。要操作的流程组。如果设置为None，则使用默认（全局）进程组。

• 包含异步调用索引的列表已成功定稿。

• state_dict（字典）-必填项。要保存的州法典。

• checkpoint_id(str)-必填项。保存检查点的存储路径。

• storage_writer(StorageWriter)-可选。in 的实例 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 客户端公开可配置参数的数据类。如果未提供，则使用 S3 ClientConfig 的默认配置。默认情况下，该part_size参数设置为 64MB。

• state_dict（字典）-必填项。state_dict待加载的。

• checkpoint_id(str)-必填项。检查点的 ID。的含义checkpoint_id取决于存储空间。它可以是文件夹或文件的路径。如果存储是键值存储，则它也可以是密钥。

• storage_reader(StorageReader)-可选。in 的实例 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 客户端公开可配置参数的数据类。如果未提供，则使用 S3 ClientConfig 的默认配置。默认情况下，该part_size参数设置为 64MB。

• smp_moe（布尔值）-是否使用 MoE 的 SMP-实现。默认值为 True。

• random_seed（整数）-专家并行分布式模块中随机运算的种子数。此种子被添加到专家平行等级中，以设置每个等级的实际种子。对于每个专家并行等级，它都是独一无二的。默认值为 12345。

• moe_load_balancing（字符串）-指定 MoE 路由器的负载平衡类型。有效的选项是aux_losssinkhorn、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（float）— 如果通过，则此比例因子将应用于 softmax。如果设置为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 模块功能。调用 a 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)。必须根据形状传递适当的布局参数。

  • 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 Fac GPT e Transformer-NeoX 模型之类的实现非常有用，该模型q具有旋转嵌入k和后置嵌fp32入。如果设置为None，则不施放任何施法。默认值为 None。

  • layout（字符串）-可用值为b h s d或b s h d。应将其设置为传递的qkv张量的布局，以便可以应用适当的转换。attn默认值为 b h s d。

• attention_dropout_prob（浮动）-应用于注意力的辍学概率。默认值为 0.0。

• scale（float）— 如果通过，则此比例因子将应用于 softmax。如果设置为None，1 / sqrt(attention_head_size)则用作比例因子。默认值为 None。

• forward(self, q, kv, causal=False, cast_dtype=None, layout="b s h d")— 常规 PyTorch 模块功能。调用 a 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)，或者由两个 torch.Tensor s 组成的元组，每个元组的形状都可能为(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 Fac GPT e Transformers-Neox 之类的实现非常q,k有用，后者fp32带有旋转嵌入功能。如果设置为None，则不施放任何施法。默认值为 None。

  • 布局（字符串）-可用值为"b h s d"或"b s h d"。应将其设置为传递的qkv张量的布局，以便可以应用适当的转换。attn默认值为 "b h s d"。

• config— 美洲驼模型的 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) — 从模型转换Hugging Face Transformer 模型与张量并SMP行度兼容并应用库的张量并行度功能。SMP

• device(torch.device) — 如果通过，则在此设备上创建新模型。如果原始模块在元设备上有任何参数（参见参数初始化延迟），则转换后的模块也将在元设备上创建，忽略此处传递的参数。默认值为 None。

• dtype(torch.dtype) — 如果通过，则将其设置为用于创建模型的 dtype 上下文管理器，并使用此 dtype 创建模型。这通常是不必要的，因为我们想在使用fp32时创建模型MixedPrecision，并且fp32是中的 PyTorch默认 dtype。默认值为 None。

• config(dict) — 这是一本用于配置SMP变压器的字典。默认值为 None。

• load_state_dict_from_rank0（Boolean）-默认情况下，此模块使用新的权重创建模型的新实例。当此参数设置为时True，SMP会尝试将原始 PyTorch 模型的状态字典从第 0 个等级加载到第 0 个等级所属的张量 parallel 组的变换模型中。将其设置为时True，等级 0 在元设备上不能有任何参数。在此变换调用之后，只有第一个张量 parallel 组填充第 0 个等级的权重。你需要在FSDP包装器True中设置为sync_module_states，才能将这些权重从第一个张量 parallel 组传递到所有其他进程。激活此功能后，SMP库会从原始模型加载状态字典。该SMP库在变换之前获取模型state_dict的值，将其转换为与变换模型的结构相匹配，为每个张量平行等级对其进行分片，将这种状态从第 0 等级传达给第 0 等级所属的张量并行组中的其他等级，然后加载它。默认值为 False。

torch.sagemaker 实用程序函数

• 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在 with 的初始化后拥有多个有用的属性torch.sagemaker.init。SMP

• 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)-训练中使用的设备总数。