SageMaker 模型并行库 v2 参考

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

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

  • 如果设置为 1，当 tensor_parallel_degree 为 1 时，脚本中的 NO_SHARD 会返回到本地 PyTorch 实现和 API。否则，它就等同于任何给定张量并行组内的 NO_SHARD。

  • 如果设置为介于 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 在训练脚本中封装模型。要了解更多信息，请参阅张量并行性。

• 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 确保在张量并行和专家并行等级中生成的随机数分别与非张量并行和非专家并行的情况相匹配。

• model(nn.Module) - PyTorch 模型，用于封装和应用 SMP v2 的延迟参数初始化功能。

• init_method_using_config（可调用）- 如果您使用 SMP v2 或支持 Hugging Face 转换器模型兼容 SMP 张量并行 的张量并行实现，请将此参数保持为默认值，即 None。默认情况下，DelayedParamIniter API 会找出如何正确初始化给定模型的方法。对于其他模型，您需要创建一个自定义参数初始化函数并将其添加到脚本中。以下代码片段是 SMP v2 为 Hugging Face 转换器模型兼容 SMP 张量并行 实现的默认 init_method_using_config 函数。请参考以下代码片段创建自己的初始化配置函数，将其添加到脚本中，并将其传递给 SMP DelayedParamIniter API 的 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.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 (dict) - 必需。保存状态字典。

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

• storage_writer (StorageWriter) - 可选。PyTorch 中用于执行写入操作的 StorageWriter 实例。如果未指定，则使用默认配置 StorageWriter。

• planner (SavePlanner) - 可选。PyTorch 中 SavePlanner 的一个实例。如果未指定，则使用默认配置 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) - 必需。保存状态字典。

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

• storage_writer (StorageWriter) - 可选。PyTorch 中用于执行写入操作的 StorageWriter 实例。如果未指定，则使用默认配置 StorageWriter。

• planner (SavePlanner) - 可选。PyTorch 中 SavePlanner 的一个实例。如果未指定，则使用默认配置 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) - 可选。PyTorch 中 LoadPlanner 的一个实例。如果未指定，则使用默认配置 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、sinkhorn、balanced 和 none。默认值为 sinkhorn。

• global_token_shuffle（布尔值）：是否在同一 EP 组内对 EP 等级的令牌进行洗牌。默认值为 False。

• moe_all_to_all_dispatcher（布尔值）：是否在 MoE 中的通信中使用多对多调度程序。默认值为 True。

• moe_aux_loss_coeff (float)：辅助负载均衡损失系数。默认值为 0.001。

• moe_z_loss_coeff (float)：z 损失系数。默认值为 0.001。

• attention_dropout_prob (float)：应用于注意力的辍学概率。默认值为 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 模块的常规函数。当调用 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：以下表格的 (batch_size x 1 x 1 x seqlen) 的 torch.Tensor。要启用此注意掩码参数，需要 triton_flash_attention=True 和 use_alibi=True。要了解如何使用此方法生成注意掩码，请参阅 FlashAttention 的代码示例。默认值为 None。

  • causal：当设置为 False（参数的默认值）时，则不应用掩码。设置为 True 时，forward 方法使用标准下三角掩码。默认值为 False。

  • cast_dtype：当设置为特定的 dtype 时，它会在 attn 之前将 qkv 张量转换为此 dtype。这对于诸如 Hugging Face 转换器 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 (float)：应用于注意力的辍学概率。默认值为 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 模块的常规函数。当调用 module(x) 时，SMP 会自动运行此函数。

  • q：以下形式 (batch_size x seqlen x num_heads x head_size) 或 (batch_size x num_heads x seqlen x head_size) 的 torch.Tensor。必须根据形状传递适当的布局参数。

  • 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 时，它会在 attn 之前将 qkv 张量转换为此 dtype。这对于诸如 Hugging Face 转换器 GPT-NeoX 的实现非常有用，它在旋转嵌入后有 q,k 与 fp32。如果设置为 None，则不应用任何转换。默认值为 None。

  • 布局（字符串）：可用值为 "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)：来自 Hugging Face 转换器模型兼容 SMP 张量并行 的模型，用于转换和应用 SMP 库的张量并行功能。

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

• dtype(torch.dtype)：如果传递，则将其设置为创建模型的数据类型上下文管理器，并使用此数据类型创建模型。这通常是不必要的，因为当使用 MixedPrecision 时，我们希望使用 fp32 创建模型，而 fp32 是 PyTorch 的默认数据类型。默认值为 None。

• config(dict)：这是用于配置 SMP 转换器的字典。默认值为 None。

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

• cp_comm_type(str) - 确定上下文并行性实现，仅当 context_parallel_degree 大于 1 时适用。此参数的可用值为 p2p 和 all_gather。在注意计算过程中，p2p 利用点对点发送-接收调用进行键和值 (KV) 张量积累，以异步方式运行，允许通信与计算重叠。另一方面，all_gather 实现采用了 AllGather 通信集体操作来进行 KV 张量累积。默认值为 "p2p"。

torch.sagemaker util 函数

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

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