SageMaker 模型并行库 v2 参考

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

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

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

  • 如果设置为介于 2 和之间的整数world_size，则在指定数量的 GPU 上进行分片。如果您未在 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.transform API 将模型封装在训练脚本中。要了解更多信息，请参阅张量并行性。

• expert_parallel_degree（整数）-指定专家并行度。该值必须介于 1 和之间world_size。默认值为 1。传递大于 1 的值不会自动启用专家并行性；请确保在训练脚本中使用 torch.sagemaker.transform API 封装 MoE 模型。

• random_seed（整数）— 分布式模块中按照 SMP 张量并行度或专家并行度进行随机运算的种子数。该种子将被添加到张量平行或专家并行等级中，以设置每个等级的实际种子。对于每个张量平行和专家并行等级，它都是独一无二的。SMP v2 确保在张量并行和专家并行等级中生成的随机数分别与和大小写相匹配。 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 张量并行度使用以下代码片段作为参考，创建自己的初始化配置函数，将其添加到脚本中，然后将其传递给 SMP AP DelayedParamIniter I 的init_method_using_config参数。

  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如果不满足这些条件，则会加注。此上下文管理器只是可选的，您无需使用此上下文管理器来初始化参数。

• 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)。使用 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 会自动运行此函数。

  • 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 Face Transformer GPT-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 Face Transformers GPT-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— 使用 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.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)-训练中使用的设备总数。