本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。
# SageMaker 模型并行库 v2 参考
以下是 SageMaker 模型并行库 v2 (SMP v2) 的参考资料。
**Topics**
+ [SMP v2 核心功能配置参数](#distributed-model-parallel-v2-reference-init-config)
+ [SMP v2 `torch.sagemaker` 软件包的参考](#model-parallel-v2-torch-sagemaker-reference)
+ [从 SMP v1 升级到 SMP v2](#model-parallel-v2-upgrade-from-v1)
## SMP v2 核心功能配置参数
以下是激活和配置 [SageMaker 模型并行度库 v2 的核心功能](model-parallel-core-features-v2.md) 的完整参数列表。它们必须以 JSON 格式编写，然后在 Pyth SageMaker on SDK 中传递给 PyTorch 估算器或另存为 JSON 文件。 SageMaker HyperPod
```
{
"hybrid_shard_degree": Integer,
"sm_activation_offloading": Boolean,
"activation_loading_horizon": Integer,
"fsdp_cache_flush_warnings": Boolean,
"allow_empty_shards": Boolean,
"tensor_parallel_degree": Integer,
"context_parallel_degree": Integer,
"expert_parallel_degree": Integer,
"random_seed": Integer
}
```
+ `hybrid_shard_degree`（整数）- 指定分片并行度。此值必须为介于 `0` 和 `world_size` 之间的整数。默认值为 `0`。
+ 如果设置为`0`，则当为 1 时`tensor_parallel_degree`，它将回退到脚本中的原生 PyTorch实现和 API。否则，它会根据 `tensor_parallel_degree` 和 `world_size` 计算出可能的最大 `hybrid_shard_degree`。当回退到原生 PyTorch FSDP 用例时，如果你使用的策略`FULL_SHARD`是的话，它会在整个集群中进行分片。 GPUs如果 `HYBRID_SHARD` 或 `_HYBRID_SHARD_ZERO2` 是策略，则相当于 8 的 `hybrid_shard_degree`。当启用张量并行后，它会根据修订后的 `hybrid_shard_degree` 进行分片。
+ 如果设置为`1`，则当为 1 时`tensor_parallel_degree`，它将回退到脚本`NO_SHARD`中的原生 PyTorch实现和 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`)。要了解更多信息，请参阅[激活分载](model-parallel-core-features-v2-pytorch-activation-offloading.md)。默认值为 `True`。
+ `activation_loading_horizon`（整数）- 指定 FSDP 激活卸载范围类型的整数。这是其输入可同时存在 GPU 内存中的检查点或卸载层的最大数量。要了解更多信息，请参阅[激活分载](model-parallel-core-features-v2-pytorch-activation-offloading.md)。输入值必须为正整数。默认值为 `2`。
+ `fsdp_cache_flush_warnings`（布尔值）-检测 PyTorch 内存管理器中是否发生缓存刷新并发出警告，因为它们会降低计算性能。默认值为 `True`。
+ `allow_empty_shards`（布尔值）- 如果张量不可分割，是否允许在张量分片时使用空分片。这是针对某些情况下检查点操作过程中崩溃的实验性修复。禁用此功能会回到最初的 PyTorch 行为。默认值为 `False`。
+ `tensor_parallel_degree`（整数）- 指定张量并行度。此值必须在 `1` 到 `world_size` 之间。默认值为 `1`。请注意，传递大于 1 的值并不能自动启用上下文并行；您还需要使用 [`torch.sagemaker.transform`](#model-parallel-v2-torch-sagemaker-reference-transform) API 在训练脚本中封装模型。要了解更多信息，请参阅[张量并行性](model-parallel-core-features-v2-tensor-parallelism.md)。
+ `context_parallel_degree`（整数）- 指定上下文并行度。此值必须在 `1` 和 `world_size` 之间，且必须是 `<= hybrid_shard_degree`。默认值为 `1`。请注意，传递大于 1 的值并不能自动启用上下文并行；您还需要使用 [`torch.sagemaker.transform`](#model-parallel-v2-torch-sagemaker-reference-transform) API 在训练脚本中封装模型。要了解更多信息，请参阅[上下文并行性](model-parallel-core-features-v2-context-parallelism.md)。
+ `expert_parallel_degree`（整数）- 指定专家并行度。此值必须在 1 到 `world_size` 之间。默认值为 `1`。请注意，传递大于 1 的值并不能自动启用上下文并行；您还需要使用 [`torch.sagemaker.transform`](#model-parallel-v2-torch-sagemaker-reference-transform) API 在训练脚本中封装模型。要了解更多信息，请参阅[专家并行性](model-parallel-core-features-v2-expert-parallelism.md)。
+ `random_seed`（整数）- 通过 SMP 张量并行或专家并行性在分布式模块中进行随机操作的种子数。此种子被添加到张量并行或专家并行等级中，以设置每个等级的实际种子。每个张量并行等级和专家并行等级都是唯一的。SMP v2 确保在张量并行和专家并行等级中生成的随机数分别与和大小写相匹配。 non-tensor-parallelism non-expert-parallelism
## SMP v2 `torch.sagemaker` 软件包的参考
本节是 SMP v2 提供的 `torch.sagemaker` 软件包的参考。
**Topics**
+ [`torch.sagemaker.delayed_param.DelayedParamIniter`](#model-parallel-v2-torch-sagemaker-reference-delayed-param-init)
+ [`torch.sagemaker.distributed.checkpoint.state_dict_saver.async_save`](#model-parallel-v2-torch-sagemaker-reference-checkpoint-async-save)
+ [`torch.sagemaker.distributed.checkpoint.state_dict_saver.maybe_finalize_async_calls`](#model-parallel-v2-torch-sagemaker-reference-checkpoint-state-dict-saver)
+ [`torch.sagemaker.distributed.checkpoint.state_dict_saver.save`](#model-parallel-v2-torch-sagemaker-reference-checkpoint-save)
+ [`torch.sagemaker.distributed.checkpoint.state_dict_loader.load`](#model-parallel-v2-torch-sagemaker-reference-checkpoint-load)
+ [`torch.sagemaker.moe.moe_config.MoEConfig`](#model-parallel-v2-torch-sagemaker-reference-moe)
+ [`torch.sagemaker.nn.attn.FlashSelfAttention`](#model-parallel-v2-torch-sagemaker-reference-flashselfattention)
+ [`torch.sagemaker.nn.attn.FlashGroupedQueryAttention`](#model-parallel-v2-torch-sagemaker-reference-flashGroupedQueryAttn)
+ [`torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention`](#model-parallel-v2-torch-sagemaker-reference-llamaFlashAttn)
+ [`torch.sagemaker.transform`](#model-parallel-v2-torch-sagemaker-reference-transform)
+ [`torch.sagemaker` util 函数和属性](#model-parallel-v2-torch-sagemaker-reference-utils)
### `torch.sagemaker.delayed_param.DelayedParamIniter`
用于应用于 PyTorch模型[延迟参数初始化](model-parallel-core-features-v2-delayed-param-init.md)的 API。
```
class torch.sagemaker.delayed_param.DelayedParamIniter(
model: nn.Module,
init_method_using_config : Callable = None,
verbose: bool = False,
)
```
**参数**
+ `model`(`nn.Module`) — 用于封装和应用 SMP v2 延迟参数初始化功能的 PyTorch 模型。
+ `init_method_using_config`（可调用）- 如果您使用 SMP v2 或支持 [Hugging Face 转换器模型兼容 SMP 张量并行](model-parallel-core-features-v2-tensor-parallelism.md#model-parallel-core-features-v2-tensor-parallelism-supported-models) 的张量并行实现，请将此参数保持为默认值，即 `None`。默认情况下，`DelayedParamIniter` API 会找出如何正确初始化给定模型的方法。对于其他模型，您需要创建一个自定义参数初始化函数并将其添加到脚本中。以下代码片段是 SMP v2 为 [Hugging Face 转换器模型兼容 SMP 张量并行](model-parallel-core-features-v2-tensor-parallelism.md#model-parallel-core-features-v2-tensor-parallelism-supported-models) 实现的默认 `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 函数和属性](#model-parallel-v2-torch-sagemaker-reference-utils)。
+ `verbose`（布尔值）- 是否在初始化和验证过程中启用更详细的日志记录。默认值为 `False`。
**Methods**
+ `get_param_init_fn()`— 返回参数初始化函数，您可以将其传递给 PyTorch FSDP 包装器类的`param_init_fn`参数。
+ `get_post_param_init_fn()`— 返回参数初始化函数，您可以将其传递给 PyTorch FSDP 包装器类的`post_param_init_fn`参数。在模型中绑定权重时需要使用此功能。模型必须实现方法 `tie_weights`。有关更多信息，请参阅 [延迟参数初始化](model-parallel-core-features-v2-delayed-param-init.md) 中的**绑定权重说明**。
+ `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`。此上下文管理器只是可选的，您无需使用此上下文管理器来初始化参数。
### `torch.sagemaker.distributed.checkpoint.state_dict_saver.async_save`
异步保存的入口 API。使用此方法可将 `state_dict` 异步保存到指定的 `checkpoint_id`。
```
def async_save(
state_dict: STATE_DICT_TYPE,
*,
checkpoint_id: Union[str, os.PathLike, None] = None,
storage_writer: Optional[StorageWriter] = None,
planner: Optional[SavePlanner] = None,
process_group: Optional[dist.ProcessGroup] = None,
coordinator_rank: int = 0,
queue : AsyncCallsQueue = None,
sharded_strategy: Union[SaveShardedStrategy, Tuple[str, int], None] = None,
wait_error_handling: bool = True,
force_check_all_plans: bool = True,
s3_region: Optional[str] = None,
s3client_config: Optional[S3ClientConfig] = None
) -> None:
```
**参数**
+ `state_dict` (dict) - 必需。保存状态字典。
+ `checkpoint_id` (str) - 必需。保存检查点的存储路径。
+ `storage_writer`(StorageWriter)-可选。in 的实例 [https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageWriter](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageWriter)， PyTorch 用于执行写入操作。如果未指定，则使用默认配置 [https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageWriter](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageWriter)。
+ `planner`(SavePlanner)-可选。中的一个实[https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner)例 PyTorch。如果未指定，则使用默认配置 [https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.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](https://github.com/awslabs/s3-connector-for-pytorch/blob/main/s3torchconnector/src/s3torchconnector/_s3client/s3client_config.py#L7) 的默认配置。`part_size` 参数默认设置为 64MB。
### `torch.sagemaker.distributed.checkpoint.state_dict_saver.maybe_finalize_async_calls`
此功能允许一个训练过程监控多个异步请求的完成情况。
```
def maybe_finalize_async_calls(
blocking=True,
process_group=None
) -> List[int]:
```
**参数**
+ `blocking` (bool) - 可选。如果是 `True`，则会等待所有活动请求完成。否则，它只会最终处理已经完成的异步请求。默认值为 `True`。
+ `process_group`(ProcessGroup)-可选。操作的进程组。如果设置为 `None`，则使用默认（全局）进程组。
**返回值**
+ 包含成功完成异步调用的索引的列表。
### `torch.sagemaker.distributed.checkpoint.state_dict_saver.save`
使用此方法可将 `state_dict` 同步保存到指定的 `checkpoint_id`。
```
def save(
state_dict: STATE_DICT_TYPE,
*,
checkpoint_id: Union[str, os.PathLike, None] = None,
storage_writer: Optional[StorageWriter] = None,
planner: Optional[SavePlanner] = None,
process_group: Optional[dist.ProcessGroup] = None,
coordinator_rank: int = 0,
wait_error_handling: bool = True,
force_check_all_plans: bool = True,
s3_region: Optional[str] = None,
s3client_config: Optional[S3ClientConfig] = None
) -> None:
```
**参数**
+ `state_dict` (dict) - 必需。保存状态字典。
+ `checkpoint_id` (str) - 必需。保存检查点的存储路径。
+ `storage_writer`(StorageWriter)-可选。in 的实例 [https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageWriter](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageWriter)， PyTorch 用于执行写入操作。如果未指定，则使用默认配置 [https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageWriter](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageWriter)。
+ `planner`(SavePlanner)-可选。中的一个实[https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner)例 PyTorch。如果未指定，则使用默认配置 [https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.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](https://github.com/awslabs/s3-connector-for-pytorch/blob/main/s3torchconnector/src/s3torchconnector/_s3client/s3client_config.py#L7) 的默认配置。`part_size` 参数默认设置为 64MB。
### `torch.sagemaker.distributed.checkpoint.state_dict_loader.load`
加载分布式模型的状态字典 (`state_dict`)。
```
def load(
state_dict: Dict[str, Any],
*,
checkpoint_id: Union[str, os.PathLike, None] = None,
storage_reader: Optional[StorageReader] = None,
planner: Optional[LoadPlanner] = None,
process_group: Optional[dist.ProcessGroup] = None,
check_keys_matched: bool = True,
coordinator_rank: int = 0,
s3_region: Optional[str] = None,
s3client_config: Optional[S3ClientConfig] = None
) -> None:
```
**参数**
+ `state_dict` (dict) - 必需。要加载的 `state_dict`。
+ `checkpoint_id` (str) - 必需。检查点的 ID。`checkpoint_id` 的含义取决于存储空间。它可以是文件夹或文件的路径。如果存储是键值存储，则它也可以是一个键。
+ `storage_reader`(StorageReader)-可选。in 的实例 [https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageReader](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.StorageReader)， PyTorch 用于执行读取操作。如果未指定，则分布式检查点将根据 `checkpoint_id` 自动推理读取器。如果 `checkpoint_id` 也是 `None`，则会出现异常错误。
+ `planner`(StorageReader)-可选。中的一个实[https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.LoadPlanner](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.LoadPlanner)例 PyTorch。如果未指定，则使用默认配置 [https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.LoadPlanner](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.LoadPlanner)。
+ `check_keys_matched` (bool) - 可选。如果启用，则检查所有级别的 `state_dict` 键是否使用 `AllGather` 匹配。
+ `s3_region` (str) - 可选。S3 存储桶所在的区域。如果未指定，则从 `checkpoint_id` 中推理区域。
+ `s3client_config`(S3ClientConfig)-可选。为 S3 客户端提供可配置参数的数据类。如果未提供，则使用 [S3 ClientConfig](https://github.com/awslabs/s3-connector-for-pytorch/blob/main/s3torchconnector/src/s3torchconnector/_s3client/s3client_config.py#L7) 的默认配置。`part_size` 参数默认设置为 64MB。
### `torch.sagemaker.moe.moe_config.MoEConfig`
用于设置 Mixture-of-Experts (MoE) 的 SMP 实现的配置类。您可以通过此类指定 MoE 配置值，并将其传递给 [https://docs.aws.amazon.com/sagemaker/latest/dg/distributed-model-parallel-v2-reference.html#model-parallel-v2-torch-sagemaker-reference-transform](https://docs.aws.amazon.com/sagemaker/latest/dg/distributed-model-parallel-v2-reference.html#model-parallel-v2-torch-sagemaker-reference-transform) API 调用。要了解此类用于训练 MoE 模型的更多信息，请参阅 [专家并行性](model-parallel-core-features-v2-expert-parallelism.md)。
```
class torch.sagemaker.moe.moe_config.MoEConfig(
smp_moe=True,
random_seed=12345,
moe_load_balancing="sinkhorn",
global_token_shuffle=False,
moe_all_to_all_dispatcher=True,
moe_aux_loss_coeff=0.001,
moe_z_loss_coeff=0.001
)
```
**参数**
+ `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 中使用 all-to-all调度器进行通信。默认值为 `True`。
+ `moe_aux_loss_coeff` (float)：辅助负载均衡损失系数。默认值为 `0.001`。
+ `moe_z_loss_coeff` (float)：z 损失系数。默认值为 `0.001`。
### `torch.sagemaker.nn.attn.FlashSelfAttention`
在 SMP v2 中使用 [FlashAttention](model-parallel-core-features-v2-flashattention.md) 的 API。
```
class torch.sagemaker.nn.attn.FlashSelfAttention(
attention_dropout_prob: float = 0.0,
scale: Optional[float] = None,
triton_flash_attention: bool = False,
use_alibi: bool = False,
)
```
**参数**
+ `attention_dropout_prob` (float)：应用于注意力的辍学概率。默认值为 `0.0`。
+ `scale` (float)：如果通过，则此比例因子将用于 softmax。如果设置为 `None`（也是默认值），则比例因子为 `1 / sqrt(attention_head_size)`。默认值为 `None`。
+ `triton_flash_attention`(bool)：如果通过，则使用 Triton 实现的闪存注意。这是支持线性偏差注意力 (B ALi i) 所必需的（参见以下`use_alibi`参数）。此版本的内核不支持停用。默认值为 `False`。
+ `use_alibi`(bool) — 如果通过，则使用提供的掩码启用线性偏差注意力 (ALiBi)。使用 ALi Bi 时，需要准备好注意面具，如下所示。默认值为 `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
```
**Methods**
+ `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](model-parallel-core-features-v2-flashattention.md) 的代码示例。默认值为 `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`。
**返回值**
形状为 `(batch_size x num_heads x seq_len x head_size)` 的单个 `torch.Tensor`。
### `torch.sagemaker.nn.attn.FlashGroupedQueryAttention`
在 SMP v2 中使用 `FlashGroupedQueryAttention` 的 API。要了解有关此 API 用法的更多信息，请参阅 [使用 FlashAttention 内核进行分组查询注意](model-parallel-core-features-v2-flashattention.md#model-parallel-core-features-v2-flashattention-grouped-query)。
```
class torch.sagemaker.nn.attn.FlashGroupedQueryAttention(
attention_dropout_prob: float = 0.0,
scale: Optional[float] = None,
)
```
**参数**
+ `attention_dropout_prob` (float)：应用于注意力的辍学概率。默认值为 `0.0`。
+ `scale` (float)：如果通过，则此比例因子将用于 softmax。如果设置为 `None`，则使用 `1 / sqrt(attention_head_size)` 作为缩放因子。默认值为 `None`。
**Methods**
+ `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"`。
**返回值**
返回代表注意计算输出的单个 `torch.Tensor (batch_size x num_heads x seq_len x head_size)`。
### `torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention`
一个支持 Llama 模型 FlashAttention 的 API。此 API 在低级别使用 [`torch.sagemaker.nn.attn.FlashGroupedQueryAttention`](#model-parallel-v2-torch-sagemaker-reference-flashGroupedQueryAttn) API。要了解如何使用，请参阅 [使用 FlashAttention 内核进行分组查询注意](model-parallel-core-features-v2-flashattention.md#model-parallel-core-features-v2-flashattention-grouped-query)。
```
class torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention(
config: LlamaConfig
)
```
**参数**
+ `config`— 美洲驼模型的 FlashAttention 配置。
**Methods**
+ `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`。
**返回值**
返回代表注意计算输出的单个 `torch.Tensor (batch_size x num_heads x seq_len x head_size)`。
### `torch.sagemaker.transform`
SMP v2 提供了 `torch.sagemaker.transform()` API，用于将 Hugging Face 转换器模型转换为 SMP 模型实现，并启用 SMP 张量并行。
```
torch.sagemaker.transform(
model: nn.Module,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
config: Optional[Dict] = None,
load_state_dict_from_rank0: bool = False,
cp_comm_type: str = "p2p"
)
```
SMP v2 通过将 Hugging Face 转换器模型的配置转换为 SMP 转换器配置，来维护 [Hugging Face 转换器模型兼容 SMP 张量并行](model-parallel-core-features-v2-tensor-parallelism.md#model-parallel-core-features-v2-tensor-parallelism-supported-models) 的转换策略。
**参数**
+ `model`(`torch.nn.Module`)：来自 [Hugging Face 转换器模型兼容 SMP 张量并行](model-parallel-core-features-v2-tensor-parallelism.md#model-parallel-core-features-v2-tensor-parallelism-supported-models) 的模型，用于转换和应用 SMP 库的张量并行功能。
+ `device`(`torch.device`)：如果传递，将在此设备上创建新模型。如果原始模块在元设备上有任何参数（见 [延迟参数初始化](model-parallel-core-features-v2-delayed-param-init.md)），则转换后的模块也将在元设备上创建，忽略此处传递的参数。默认值为 `None`。
+ `dtype`(`torch.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 级不能在元设备上设置任何参数。在此转换调用之后，只有第一个并行张量组会从第 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`实现利用 peer-to-peer发送-接收调用 key-and-value (KV) 张量累积，异步运行并允许通信与计算重叠。另一方面，`all_gather` 实现采用了 `AllGather` 通信集体操作来进行 KV 张量累积。默认值为 `"p2p"`。
**返回值**
返回一个可以用 PyTorch FSDP 封装的转换后的模型。当 `load_state_dict_from_rank0` 设置为 `True` 时，涉及第 0 级的张量并行组的权重将从第 0 级的原始状态字典中加载。在原始模型[延迟参数初始化](model-parallel-core-features-v2-delayed-param-init.md)上使用时，只有这些等级才会 CPUs 为变换后的模型的参数和缓冲区开启实际张量。其余等级的参数和缓冲区继续保留在元设备上，以节省内存。
### `torch.sagemaker` util 函数和属性
**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`，如果指定，可以递归到所有嵌套模块。
**Properties**
在使用 `torch.sagemaker.init` 初始化 SMP 后，`torch.sagemaker.state` 拥有多种有用的属性。
+ `torch.sagemaker.state.hybrid_shard_degree`(int) - 分片数据并行度，SMP 配置中用户输入的副本，传递给 `torch.sagemaker.init()`。要了解更多信息，请参阅[使用 SageMaker 模型并行度库 v2](model-parallel-use-api-v2.md)。
+ `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](model-parallel-use-api-v2.md)。
+ `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) - 训练中使用的设备总数。
## 从 SMP v1 升级到 SMP v2
要从 SMP v1 迁移到 SMP v2，必须更改脚本以删除 SMP v1 并应用 SMP v APIs 2。 APIs我们建议您不要从 SMP v1 脚本开始，而是从 PyTorch FSDP 脚本开始，然后按照中的说明进行操作。[使用 SageMaker 模型并行度库 v2](model-parallel-use-api-v2.md)
要将 SMP v1 *模型*引入 SMP v2，必须在 SMP v1 中收集完整的模型状态字典，并在模型状态字典上应用转换函数，将其转换为 Hugging Face 转换器模型检查点格式。然后，在 SMP v2 中，如中所述[使用 SMP 的检查点](model-parallel-core-features-v2-checkpoints.md)，你可以加载 Hugging Face Transformers 模型检查点，然后在 SMP v2 中继续使用 PyTorch 该 APIs 检查点。要将 SMP 与 PyTorch FSDP 模型一起使用，请务必移至 SMP v2 并更改训练脚本以使用 PyTorch FSDP 和其他最新功能。
```
import smdistributed.modelparallel.torch as smp
# Create model
model = ...
model = smp.DistributedModel(model)
# Run training
...
# Save v1 full checkpoint
if smp.rdp_rank() == 0:
model_dict = model.state_dict(gather_to_rank0=True) # save the full model
# Get the corresponding translation function in smp v1 and translate
if model_type == "gpt_neox":
from smdistributed.modelparallel.torch.nn.huggingface.gptneox import translate_state_dict_to_hf_gptneox
translated_state_dict = translate_state_dict_to_hf_gptneox(state_dict, max_seq_len=None)
# Save the checkpoint
checkpoint_path = "checkpoint.pt"
if smp.rank() == 0:
smp.save(
{"model_state_dict": translated_state_dict},
checkpoint_path,
partial=False,
)
```
要查找 SMP v1 中可用的转换功能，请参阅 [支持 Hugging Face 转换器模型](model-parallel-extended-features-pytorch-hugging-face.md)。
有关在 SMP v2 中保存和加载模型检查点的说明，请参阅 [使用 SMP 的检查点](model-parallel-core-features-v2-checkpoints.md)。