本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

# SageMaker 模型并行度库的核心功能
核心功能

Amazon SageMaker AI 的模型并行库提供分发策略和节省内存的技术，例如分片数据并行性、张量并行性、按层划分模型以进行管道调度以及检查点。模型并行性策略和技术有助于将大型模型分布在多个设备上，同时优化训练速度和内存使用。该库还提供 Python 帮助程序函数、上下文管理器和封装器函数，用于调整训练脚本以实现模型的自动或手动分区。

在训练作业中实现模型并行性时，将保持 “使用[模型并行度运行分布式 SageMaker 训练作业” 部分中显示的相同的两步工作](https://docs.amazonaws.cn/sagemaker/latest/dg/model-parallel-use-api.html)流程。对于调整训练脚本，您无需或者仅需在训练脚本中添加几行额外的代码。要启动调整后的训练脚本的训练作业，您需要设置分布配置参数以激活节省内存的功能或者传递用于并行度的值。

要开始使用示例，请参阅以下 Jupyter 笔记本，其中演示了如何使用 SageMaker 模型并行度库。
+ [PyTorch 示例笔记本](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel)
+ [TensorFlow 示例笔记本](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/tensorflow/model_parallel/mnist)

要深入了解库的核心功能，请参阅以下主题。

**注意**  
 SageMaker 分布式训练库可通过 Hugging Face 的 Amazon PyTorch深度学习容器获得， TensorFlow 也可以在训练平台 SageMaker 内使用。要使用分布式训练库的功能，我们建议您使用 SageMaker Python SDK。如果你 SageMaker APIs 通过适用于 Python 的 SDK (Boto3) 或，也可以使用 JSON 请求语法手动配置。 Amazon Command Line Interface在整篇文档中，说明和示例都侧重于如何将分布式训练库与 SageMaker Python SDK 配合使用。

**重要**  
 SageMaker 模型并行度库支持的所有核心功能，并支持流水线并行性。 PyTorch TensorFlow

**Topics**
+ [

# 分片数据并行性
](model-parallel-extended-features-pytorch-sharded-data-parallelism.md)
+ [

# 模型管道传输
](model-parallel-core-features-pipieline-parallelism.md)
+ [

# 张量并行性
](model-parallel-extended-features-pytorch-tensor-parallelism.md)
+ [

# 优化器状态分片
](model-parallel-extended-features-pytorch-optimizer-state-sharding.md)
+ [

# 激活检查点
](model-parallel-extended-features-pytorch-activation-checkpointing.md)
+ [

# 激活分载
](model-parallel-extended-features-pytorch-activation-offloading.md)
+ [

# FP16 使用模型并行度进行训练
](model-parallel-extended-features-pytorch-fp16.md)
+ [

# Support FlashAttention
](model-parallel-attention-head-size-for-flash-attention.md)

# 分片数据并行性


分@@ *片数据并行性*是一种节省内存的分布式训练技术，它在数据并行组中拆分模型的状态（模型参数、梯度和优化器状态）。 GPUs 

**注意**  
分片数据并行度可在 SageMaker 模型并行度库 v1.11.0 及更高版本 PyTorch 中使用。

将训练作业扩展到大型 GPU 集群时，您可以通过将模型的训练状态分为多个集群来减少模型的每 GPU 内存占用。 GPUs这会带来两个好处：您可以容纳在标准数据并行性下会耗尽内存的更大模型；或者您可以使用腾出的 GPU 内存来增加批次大小。

标准的数据并行度技术在数据并行组 GPUs中复制训练状态，并根据操作执行梯度聚合。`AllReduce`分片数据并行性修改了标准的数据并行分布式训练过程，以考虑优化器状态的分片性质。用于对模型和优化器状态进行分片的一组秩称为*分片组*。*分片数据并行技术对模型的可训练参数以及分片组中相应的梯度和优化器状态进行分片。 GPUs *

SageMaker 人工智能通过实现 MIC 实现分片数据并行性， Amazon 博客文章巨型模型训练的[近线性扩展](https://www.amazon.science/blog/near-linear-scaling-of-gigantic-model-training-on-aws)对此进行了讨论。 Amazon在此实施中，您可以将分片度设置为可配置参数，该参数必须小于数据并行度。在每次向前和向后传递过程中，MIC 都会在 GPUs 整个操作过程中临时重新组合模型参数。`AllGather`在每层向前或向后传递后，MiCS 会再次对参数进行分片以节省 GPU 内存。在向后传递过程中，MIC 会降低梯度，同时在整个操作中 GPUs 将其分片。`ReduceScatter`最后，MiCS 使用优化器状态的局部分片，将局部缩减梯度和分片梯度应用于其对应的局部参数分片。为了降低通信开销， SageMaker 模型并行度库在向前或向后通道中预取即将到来的层，并将网络通信与计算重叠。

模型的训练状态在分片组之间复制。这意味着，在将梯度应用于参数之前，除了在分片组内进行的 `ReduceScatter` 操作之外，还必须跨分片组进行 `AllReduce` 操作。

实际上，分片数据并行性在通信开销和 GPU 内存效率之间进行了权衡。使用分片数据并行性会增加通信成本，但是在每个 GPU 上占用的内存量（不包括激活导致的内存使用量）会除以分片数据并行度，因此可以将更大的模型放入 GPU 集群。

**选择分片数据并行度**

当您为分片数据并行度选择一个值时，该值必须能够整除数据并行度。例如，对于 8 路数据并行性作业，请选择 2、4 或 8 作为分片数据并行度。在选择分片数据并行度时，我们建议您从一个较小的数字开始，然后逐渐增加它，直到模型与所需的批次大小均适合内存。

**选择批次大小**

设置分片数据并行性后，请确保您找到了可以在 GPU 集群上成功运行的最佳训练配置。要训练大型语言模型 (LLM)，请从批次大小 1 开始，然后逐渐增加批次大小，直到达到接收 out-of-memory (OOM) 错误的程度。如果您即使使用最小的批次也会遇到 OOM 错误，请应用更高的分片数据并行度，或者将分片数据并行性和张量并行性结合使用。

**Topics**
+ [

## 如何将分片数据并行性应用于训练作业
](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use)
+ [

## 参考配置
](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use-config-sample)
+ [

## 使用 SMDDP 集合体的分片数据并行性
](#model-parallel-extended-features-pytorch-sharded-data-parallelism-smddp-collectives)
+ [

## 混合精度训练与分片数据并行性
](#model-parallel-extended-features-pytorch-sharded-data-parallelism-16bits-training)
+ [

## 使用张量并行性的分片数据并行性
](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism)
+ [

## 使用分片数据并行性的提示和注意事项
](#model-parallel-extended-features-pytorch-sharded-data-parallelism-considerations)

## 如何将分片数据并行性应用于训练作业


要开始使用分片数据并行性，请对训练脚本进行必要的修改，然后使用参数设置 SageMaker PyTorch 估计器。 sharded-data-parallelism-specific另外还可以考虑以参考值和示例笔记本为起点。

### 调整您的 PyTorch 训练脚本


按照[步骤 1：修改 PyTorch 训练脚本](model-parallel-customize-training-script-pt.md)中的说明使用和模块的`smdistributed.modelparallel.torch`包装器封装模型`torch.nn.parallel`和`torch.distributed`优化器对象。

**（可选）用于注册外部模型参数的额外修改**

如果您的模型使用 `torch.nn.Module` 构建并使用了未在模型类中定义的参数，您应手动将这些参数注册到模块，以便 SMP 收集完整参数。要将参数注册到模块，请使用 `smp.register_parameter(module, parameter)`。

```
class Module(torch.nn.Module):
    def __init__(self, *args):
        super().__init__(self, *args)
        self.layer1 = Layer1()
        self.layer2 = Layer2()
        smp.register_parameter(self, self.layer1.weight)

    def forward(self, input):
        x = self.layer1(input)
        # self.layer1.weight is required by self.layer2.forward
        y = self.layer2(x, self.layer1.weight)
        return y
```

### 设置 SageMaker PyTorch 估算器


在中配置 SageMaker PyTorch 估算器时[第 2 步：使用 SageMaker Python 软件开发工具包启动训练 Job](model-parallel-sm-sdk.md)，请添加分片数据并行度参数。

要开启分片数据并行性，请将`sharded_data_parallel_degree`参数添加到 Estimator 中。 SageMaker PyTorch此参数指定分片训练状态的数量。 GPUs `sharded_data_parallel_degree` 的值必须是介于 1 和数据并行度之间的整数，并且必须能够整除数据并行度。请注意，该库会自动检测数字 GPUs ，因此数据并行度。以下附加参数可用于配置分片数据并行度。
+ `"sdp_reduce_bucket_size"`*（int，默认值：5e8）*— 以默认 dtype 的元素数量指定 [PyTorch DDP 渐变桶](https://pytorch.org/docs/stable/notes/ddp.html#internal-design)的大小。
+ `"sdp_param_persistence_threshold"`*（整数，默认值：1e6）* – 以每个 GPU 上可以持续存在的元素数量来指定参数张量的大小。分片数据并行性将每个参数张量分成 GPUs 一个数据并行组。如果参数张量中的元素数量小于此阈值，则不会拆分参数张量；这有助于减少通信开销，因为参数张量是跨数据并行复制的。 GPUs
+ `"sdp_max_live_parameters"`*（整数，默认值：1e9）* – 指定在向前和向后传递期间，可以同时处于重新组合训练状态的最大参数数量。当活动参数的数量达到给定阈值时，提取参数的 `AllGather` 操作将暂停。请注意，增加此参数会增加内存占用。
+ `"sdp_hierarchical_allgather"`*（布尔值，默认值：True）* – 如果设置为 `True`，则 `AllGather` 操作按层次运行：首先在每个节点内运行，然后跨节点运行。对于多节点分布式训练作业，分层 `AllGather` 操作会自动激活。
+ `"sdp_gradient_clipping"`*（浮点数，默认值：1.0）* – 指定一个阈值，用于在通过模型参数向后传播梯度之前，先裁剪梯度的 L2 范数。当分片数据并行性被激活时，梯度裁剪也被激活。默认阈值为 `1.0`。如果您遇到梯度爆炸问题，请调整此参数。

以下代码演示了如何配置分片数据并行性的示例。

```
import sagemaker
from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled": True,
    "parameters": {
        # "pipeline_parallel_degree": 1,    # Optional, default is 1
        # "tensor_parallel_degree": 1,      # Optional, default is 1
        "ddp": True,
        # parameters for sharded data parallelism
        "sharded_data_parallel_degree": 2,              # Add this to activate sharded data parallelism
        "sdp_reduce_bucket_size": int(5e8),             # Optional
        "sdp_param_persistence_threshold": int(1e6),    # Optional
        "sdp_max_live_parameters": int(1e9),            # Optional
        "sdp_hierarchical_allgather": True,             # Optional
        "sdp_gradient_clipping": 1.0                    # Optional
    }
}

mpi_options = {
    "enabled" : True,                      # Required
    "processes_per_host" : 8               # Required
}

smp_estimator = PyTorch(
    entry_point="your_training_script.py", # Specify your train script
    role=sagemaker.get_execution_role(),
    instance_count=1,
    instance_type='ml.p3.16xlarge',
    framework_version='1.13.1',
    py_version='py3',
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": mpi_options
    },
    base_job_name="sharded-data-parallel-job"
)

smp_estimator.fit('s3://my_bucket/my_training_data/')
```

## 参考配置


 SageMaker 分布式训练团队提供了以下参考配置，您可以将其用作起点。您可以根据上述配置进行推断，以实验并估算模型配置的 GPU 内存使用量。

使用 SMDDP 集合体的分片数据并行性


| 模型/参数数量 | 实例数 | 实例类型 | 序列长度 | 全局批次大小 | 小批次大小 | 分片数据并行度 | 
| --- | --- | --- | --- | --- | --- | --- | 
| GPT-NEOX-20B | 2 | ml.p4d.24xlarge | 2048 | 64 | 4 | 16 | 
| GPT-NEOX-20B | 8 | ml.p4d.24xlarge | 2048 | 768 | 12 | 32 | 

例如，如果您增加了 200 亿参数模型的序列长度或将模型的大小增加到 650 亿个参数，则您需要先尝试减小批次大小。如果在最小的批次大小（批次大小为 1）模型仍然不适合，请尝试提高模型的并行度。

使用张量并行性和 NCCL 集合体的分片数据并行性


| 模型/参数数量 | 实例数 | 实例类型 | 序列长度 | 全局批次大小 | 小批次大小 | 分片数据并行度 | 张量并行度 | 激活分载 | 
| --- | --- | --- | --- | --- | --- | --- | --- | --- | 
| GPT-NEOX-65B | 64 | ml.p4d.24xlarge | 2048 | 512 | 8 | 16 | 8 | Y | 
| GPT-NEOX-65B | 64 | ml.p4d.24xlarge | 4096 | 512 | 2 | 64 | 2 | Y | 

当你想将大型语言模型 (LLM) 拟合到大规模集群中，同时使用序列长度更长的文本数据，从而使用较小的批处理大小，从而处理 GPU 内存使用量来针对更长的文本序列进行训练时，将分片数据并行性和张量并行性结合起来非常有用。 LLMs 要了解更多信息，请参阅[使用张量并行性的分片数据并行性](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism)。

有关案例研究、基准测试和更多配置示例，请参阅博客文章 [Amazon A SageMaker I 模型并行库中的新性能改进](https://www.amazonaws.cn/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)。

## 使用 SMDDP 集合体的分片数据并行性


 SageMaker 数据并行度库提供针对基础架构进行了优化的集体通信基元（SMDDP 集合）。 Amazon 它通过使用 Ela [stic Fabric Adapter (EFA) Adapter (EFA)](https://www.amazonaws.cn/hpc/efa/) 来实现优化，从而产生高吞吐量且对延迟不敏感的集合，将与通信相关的处理分流给 CPU，并腾出 GPU 周期用于计算。 all-to-all-type在大型集群上，与 NCCL 相比，SMDDP 集合体可以将分布式训练性能提高多达40%。有关案例研究和基准测试结果，请参阅博客 [Amazon A SageMaker I 模型并行度库中的新性能改进](https://www.amazonaws.cn/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)。

**注意**  
 SageMaker 模型并行度库 v1.13.0 及更高版本以及数据并行库 v1.6.0 及更高版本中提供了与 SMDDP Collectives 的分片数据并行性。 SageMaker 有关在 SMDDP 集合体中使用分片数据并行性的信息，另请参阅[Supported configurations](#sharded-data-parallelism-smddp-collectives-supported-config)。

在分片数据并行性（这是在大规模分布式训练中常用的技术）中，`AllGather` 集合体用于针对向前和向后传递计算重构分片层参数，与 GPU 计算并行处理。对于大型模型，高效执行 `AllGather` 操作对于避免 GPU 瓶颈问题以及避免降低训练速度至关重要。激活分片数据并行性后，SMDDP 集合体会放入这些对性能至关重要的 `AllGather` 集合体中，从而提高训练吞吐量。

**使用 SMDDP 集合体进行训练**

当您的训练作业已激活分片数据并行性并满足[Supported configurations](#sharded-data-parallelism-smddp-collectives-supported-config)时，SMDDP 集合体会自动激活。在内部，SMDDP Collectives 会优化`AllGather`集体，使其在 Amazon 基础设施上保持高性能，而所有其他集体的性能则回归到 NCCL。此外，在不支持的配置下，所有集合体（包括 `AllGather`）都会自动使用 NCCL 后端。

从 SageMaker 模型并行度库版本 1.13.0 起，该`"ddp_dist_backend"`参数已添加到选项中。`modelparallel`此配置参数的默认值为 `"auto"`，这将尽可能使用 SMDDP 集合体，否则回退到 NCCL。要强制库始终使用 NCCL，请将 `"nccl"` 指定为 `"ddp_dist_backend"` 配置参数。

以下代码示例显示了如何使用分片数据并行度设置 PyTorch 估算器，`"ddp_dist_backend"`参数`"auto"`默认设置为，因此可以选择添加。

```
import sagemaker
from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled":True,
    "parameters": {                        
        "partitions": 1,
        "ddp": True,
        "sharded_data_parallel_degree": 64
        "bf16": True,
        "ddp_dist_backend": "auto"  # Specify "nccl" to force to use NCCL.
    }
}

mpi_options = {
    "enabled" : True,                      # Required
    "processes_per_host" : 8               # Required
}

smd_mp_estimator = PyTorch(
    entry_point="your_training_script.py", # Specify your train script
    source_dir="location_to_your_script",
    role=sagemaker.get_execution_role(),
    instance_count=8,
    instance_type='ml.p4d.24xlarge',
    framework_version='1.13.1',
    py_version='py3',
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": mpi_options
    },
    base_job_name="sharded-data-parallel-demo",
)

smd_mp_estimator.fit('s3://my_bucket/my_training_data/')
```

**支持的配置**

满足以下所有配置要求后，训练作业中将激活 `AllGather` 操作和 SMDDP 集合体。
+ 分片数据并行度大于 1
+ `Instance_count` 大于 1 
+ `Instance_type` 等于 `ml.p4d.24xlarge` 
+ SageMaker 适用于 PyTorch v1.12.1 或更高版本的训练容器
+  SageMaker 数据并行度库 v1.6.0 或更高版本
+  SageMaker 模型并行度库 v1.13.0 或更高版本

**性能和内存调整**

SMDDP 集合体使用额外的 GPU 内存。根据不同的模型训练使用场景，有两个环境变量可用于配置 GPU 内存使用情况。
+ `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` – 在 SMDDP `AllGather` 操作期间，`AllGather` 输入缓冲区被复制到临时缓冲区中，用于节点间通信。`SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 变量控制此临时缓冲区的大小（以字节为单位）。如果临时缓冲区的大小小于 `AllGather` 输入缓冲区的大小，则 `AllGather` 集合体将回退到使用 NCCL。
  + 默认值：16 \$1 1024 \$1 1024 (16 MB)
  + 可接受值：8192 的任意倍数
+  `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` – `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` 变量用于调整临时缓冲区的大小（以字节为单位），以保存从节点间通信中收集的数据。如果临时缓冲区的大小小于 `1/8 * sharded_data_parallel_degree * AllGather input size`，则 `AllGather` 集合体将回退到使用 NCCL。
  + 默认值：128 \$1 1024 \$1 1024 (128 MB)
  + 可接受值：8192 的任意倍数

**缓冲区大小变量调整指南**

环境变量的默认值应该适用于大多数使用场景。我们建议只有在训练遇到 out-of-memory (OOM) 错误时才调整这些变量。

以下列表讨论了一些调整技巧，可以减少 SMDDP 集合体的 GPU 内存占用，同时继续获得带来的性能提升。
+ 调整 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES`
  + 对于较小的模型，`AllGather` 输入缓冲区的大小会更小。因此，对于参数较少的模型，所需的 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 大小可能会更小。
  + `AllGather`输入缓冲区的大小会随着`sharded_data_parallel_degree`增加而减小，因为模型会被分片得更多 GPUs。因此，对于 `sharded_data_parallel_degree` 具有较大值的训练作业，所需的 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 大小可能会更小。
+ 调整 `SMDDP_AG_SORT_BUFFER_SIZE_BYTES`
  + 对于参数较少的模型，从节点间通信中收集的数据量较少。因此，对于参数数量较少的此类模型，所需的 `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` 大小可能会更小。

有些集合体可能会回退为使用 NCCL；因此，您可能无法获得优化 SMDDP 集合体所带来的性能提升。如果有额外的 GPU 内存可供使用，您可以考虑增加 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 和 `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` 的值以从性能提升中受益。

以下代码显示了如何通过将环境变量附加到 PyTorch 估计器的分布参数`mpi_options`中来配置环境变量。

```
import sagemaker
from sagemaker.pytorch import PyTorch

smp_options = {
    .... # All modelparallel configuration options go here
}

mpi_options = {
    "enabled" : True,                      # Required
    "processes_per_host" : 8               # Required
}

# Use the following two lines to tune values of the environment variables for buffer
mpioptions += " -x SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES=8192" 
mpioptions += " -x SMDDP_AG_SORT_BUFFER_SIZE_BYTES=8192"

smd_mp_estimator = PyTorch(
    entry_point="your_training_script.py", # Specify your train script
    source_dir="location_to_your_script",
    role=sagemaker.get_execution_role(),
    instance_count=8,
    instance_type='ml.p4d.24xlarge',
    framework_version='1.13.1',
    py_version='py3',
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": mpi_options
    },
    base_job_name="sharded-data-parallel-demo-with-tuning",
)

smd_mp_estimator.fit('s3://my_bucket/my_training_data/')
```

## 混合精度训练与分片数据并行性


要通过半精度浮点数和分片数据并行度进一步节省 GPU 内存，您可以通过向分布式训练配置中添加一个附加参数来激活 16 位[浮点格式 (FP16BF16) 或 Brain 浮点格式](https://en.wikichip.org/wiki/brain_floating-point_format) ()。

**注意**  
 SageMaker 模型并行度库 v1.11.0 及更高版本中提供了具有分片数据并行性的混合精度训练。

**用于使用分片数据并行度进行 FP16 训练**

要使用分片数据并行度运行 FP16 训练，请`"fp16": True"`添加到配置字典中。`smp_options`在训练脚本中，您可以通过 `smp.DistributedOptimizer` 模块在静态和动态损失缩放选项之间进行选择。有关更多信息，请参阅 [FP16 使用模型并行度进行训练](model-parallel-extended-features-pytorch-fp16.md)。

```
smp_options = {
    "enabled": True,
    "parameters": {
        "ddp": True,
        "sharded_data_parallel_degree": 2,
        "fp16": True
    }
}
```

**用于使用分片数据并行度进行 BF16 训练**

 SageMaker AI 的分片数据并行功能支持数据类型训练。 BF16 BF16 数据类型使用 8 位来表示浮点数的指数，而 FP16 数据类型使用 5 位。保留指数的 8 位可以保持 32 位单精度浮点 () FP32 数的指数的相同表示形式。这使得 FP32 和之间的转换变得 BF16更简单，并且不太容易导致 FP16 训练中经常出现的溢出和下溢问题，尤其是在训练较大的模型时。虽然这两种数据类型总共使用 16 位，但 BF16 格式中指数表示范围的增加是以降低精度为代价的。对于训练大型模型，这种精度缩减通常会视为面向范围和训练稳定性作出的可接受取舍。

**注意**  
当前，只有激活分片数据并行性后， BF16 训练才有效。

要使用分片数据并行度运行 BF16 训练，请`"bf16": True`添加到配置字典中。`smp_options`

```
smp_options = {
    "enabled": True,
    "parameters": {
        "ddp": True,
        "sharded_data_parallel_degree": 2,
        "bf16": True
    }
}
```

## 使用张量并行性的分片数据并行性


如果您使用分片数据并行性并且还需要减小全局批次大小，请考虑将[张量并行性](https://docs.amazonaws.cn/sagemaker/latest/dg/model-parallel-extended-features-pytorch-tensor-parallelism.html)与分片数据并行性结合使用。在非常大的计算集群（通常包含 128 个节点或更多）上训练采用分片数据并行性的大型模型时，即使每个 GPU 的批次大小很小，也会得到非常大的全局批次大小。这可能会导致收敛问题或计算性能低下问题。当单个批次已经很大并且无法进一步缩小时，有时候仅使用分片数据并行性无法减少每个 GPU 上的批次大小。在这种情况下，将分片数据并行性与张量并行性结合使用，有助于减少全局批次大小。

最佳分片数据并行度和张量并行度的选择，取决于模型的规模、实例类型以及模型能够合理收敛的全局批次大小。我们建议您从低张量 parallel 度开始，以使全局批量大小适合计算集群，从而解决 CUDA out-of-memory 错误并获得最佳性能。请参阅以下两个示例案例，了解张量并行性和分片数据并行性的组合如何帮助您通过对模型并行性进行分组 GPUs 来调整全局批次大小，从而减少模型副本的数量和更小的全局批次大小。

**注意**  
此功能可从 SageMaker 模型并行度库 v1.15 中获得，并支持 v1.13.1。 PyTorch 

**注意**  
此功能可通过库的张量并行度功能，面向支持的模型提供。要查找支持的模型列表，请参阅[对 Hugging Face 转换器模型的支持](https://docs.amazonaws.cn/sagemaker/latest/dg/model-parallel-extended-features-pytorch-hugging-face.html)。另请注意，在修改训练脚本时，您需要将 `tensor_parallelism=True` 传递给 `smp.model_creation` 参数。要了解更多信息，请参阅 *SageMaker AI 示例 GitHub 存储库[https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L793](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L793)*中的训练脚本。

### 示例 1


假设我们要 GPUs 在 1536 个集群 GPUs （192 个节点，每个节点 8 个）上训练模型，将分片数据并行度设置为 32 (`sharded_data_parallel_degree=32`)，将每个 GPU 的批处理大小设置为 1，其中每个批次的序列长度为 4096 个令牌。在本例中有 1536 个模型副本，全局批次大小变为 1536，每个全局批次包含大约 600 万个令牌。

```
(1536 GPUs) * (1 batch per GPU) = (1536 global batches)
(1536 batches) * (4096 tokens per batch) = (6,291,456 tokens)
```

向其添加张量并行性，使其可以减少全局批次大小。在一种配置示例中，可以将张量并行度设置为 8，将每个 GPU 的批次大小设置为 4。这形成了 192 个张量并行组或 192 个模型副本，其中每个模型副本分布在 8 个副本上。 GPUs批次大小为 4 是每次迭代中每个张量并行组的训练数据量；也就是说，每个模型副本每次迭代消耗 4 个批次。在这种情况下，全局批次大小变为 768，每个全局批次包含大约 300 万个令牌。这样，与之前仅使用分片数据并行性的情况相比，全局批次大小减少了一半。

```
(1536 GPUs) / (8 tensor parallel degree) = (192 tensor parallelism groups)
(192 tensor parallelism groups) * (4 batches per tensor parallelism group) = (768 global batches)
(768 batches) * (4096 tokens per batch) = (3,145,728 tokens)
```

### 示例 2


当分片数据并行性和张量并行性均已激活时，库会首先应用张量并行性，并在该维度上对模型分片。对于每个张量并行秩，数据并行性根据 `sharded_data_parallel_degree` 来应用。

例如，假设我们要设置 32 GPUs ，张量 parallel 度为 4（形成一组 4 GPUs），分片数据并行度为 4，最终复制度为 2。该分配基于如下所示张量并行度创建八个 GPU 组：`(0,1,2,3)`、`(4,5,6,7)`、`(8,9,10,11)`、`(12,13,14,15)`、`(16,17,18,19)`、`(20,21,22,23)`、`(24,25,26,27)`、`(28,29,30,31)`。也就是说，四 GPUs 形成一个张量 parallel 群。在这种情况下，张量并行组第 0 个等级 GPUs 的简化数据并行组将是。`(0,4,8,12,16,20,24,28)`简化的数据并行组根据分片数据并行度 4 进行分片，从而产生两个用于数据并行性的复制组。 GPUs`(0,4,8,12)`形成一个分片组，该分片组共同保存第 0 张量 parallel 等级的所有参数的完整副本， GPUs`(16,20,24,28)`然后形成另一个这样的分片组。其他张量并行秩也有类似的分片和复制组。

![\[图 1：张量并行组。\]](http://docs.amazonaws.cn/sagemaker/latest/dg/images/distributed/model-parallel/sdp_tp_group_tp.jpg)


图 1：（节点、分片数据并行度、张量并行度）=（4, 4, 4）的张量并行组，其中每个矩形代表一个指数从 0 到 31 的 GPU。从 TPG 到 T 0 PG 的 GPUs 形式张量并行度分组。7复制组为（\$1TPG0, TPG4\$1、\$1TPG1, TPG5\$1、\$1TPG2, TPG6\$1 和 \$1TPG3, TPG7\$1）；每个复制组对的颜色相同，但填充方式不同。

![\[图 2：分片数据并行组。\]](http://docs.amazonaws.cn/sagemaker/latest/dg/images/distributed/model-parallel/sdp_tp_group_sdp.jpg)


图 2：（节点、分片数据并行度、张量并行度）=（4, 4, 4）的分片数据并行组，其中每个矩形代表一个指数从 0 到 31 的 GPU。该 GPUs 表单将从 SDPG 到 SDPG 的数据并行分组划分0。7复制组为（\$1SDPG 0, SDPG 4\$1、\$1SDPG 1, SDPG 5\$1、\$1SDPG 2, SDPG 6\$1 和 \$1SDPG 3, SDPG 7\$1）；每个复制组对的颜色相同，但填充方式不同。

### 如何激活使用张量并行性的分片数据并行性


要将分片数据并行性与张量并行性结合使用，您需要在创建估计器类的对象时`tensor_parallel_degree`在配置中`distribution`同时设置`sharded_data_parallel_degree`和。 SageMaker PyTorch

您还需要激活 `prescaled_batch`。这意味着，不是每个 GPU 读取自己的批次数据，而是每个张量并行组集体读取具有所选批次大小的一个组合批次。实际上，它不是将数据集分成等于数量 GPUs （或数据并行大小`smp.dp_size()`）的部分，而是将其分成等于 GPUs 除以数的部分`tensor_parallel_degree`（也称为缩小后的数据并行大小，`smp.rdp_size()`）。有关预缩放批处理的更多详细信息，请参阅 *SageMaker Python* SDK 文档中的 Pres [caled Bat](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#prescaled-batch) ch。另请参阅 *SageMaker AI 示例 GitHub 存储库中的 GPT-2 训练脚本[https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L164](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L164)示例*。

以下代码片段显示了基于中上述场景创建 PyTorch 估算器对象的示例。[示例 2](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism-ex2)

```
mpi_options = "-verbose --mca orte_base_help_aggregate 0 "
smp_parameters = {
    "ddp": True,
    "fp16": True,
    "prescaled_batch": True,
    "sharded_data_parallel_degree": 4,
    "tensor_parallel_degree": 4
}

pytorch_estimator = PyTorch(
    entry_point="your_training_script.py",
    role=role,
    instance_type="ml.p4d.24xlarge",
    volume_size=200,
    instance_count=4,
    sagemaker_session=sagemaker_session,
    py_version="py3",
    framework_version="1.13.1",
    distribution={
        "smdistributed": {
            "modelparallel": {
                "enabled": True, 
                "parameters": smp_parameters,
            }
        },
        "mpi": {
            "enabled": True,
            "processes_per_host": 8,
            "custom_mpi_options": mpi_options,
        },
    },
    source_dir="source_directory_of_your_code",
    output_path=s3_output_location
)
```

## 使用分片数据并行性的提示和注意事项


在使用 SageMaker 模型并行度库的分片数据并行度时，请考虑以下几点。
+ 分片数据并行性与训练兼容。 FP16 要进行 FP16训练，请参阅一[FP16 使用模型并行度进行训练](model-parallel-extended-features-pytorch-fp16.md)节。
+ 分片数据并行性与张量并行性兼容。将分片数据并行性与张量并行性结合使用时，您可能需要考虑以下几点。
  + 将分片数据并行性与张量并行性结合使用时，嵌入层也会自动分布在张量并行组中。换而言之，`distribute_embedding` 参数会自动设置为 `True`。有关张量并行性的更多信息，请参阅[张量并行性](model-parallel-extended-features-pytorch-tensor-parallelism.md)。
  + 请注意，具有张量并行性的分片数据并行性目前使用 NCCL 集合体作为分布式训练策略的后端。

  要了解更多信息，请参阅[使用张量并行性的分片数据并行性](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism)一节。
+ 分片数据并行性目前不兼容[管道并行性](model-parallel-intro.md#model-parallel-intro-pp)或[优化器状态分片](model-parallel-extended-features-pytorch-optimizer-state-sharding.md)。要激活分片数据并行性，请关闭优化器状态分片，并将管道并行度设置为 1。
+ [激活检查点](model-parallel-extended-features-pytorch-activation-checkpointing.md)和[激活分载](model-parallel-extended-features-pytorch-activation-offloading.md)功能与分片数据并行性兼容。
+ 要将分片数据并行性与梯度累积一起使用，请在使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.DistributedModel](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.DistributedModel) 模块包装模型时，将 `backward_passes_per_step` 参数设置为累积步骤数。这可确保跨模型复制组（分片组）的梯度 `AllReduce` 操作发生在梯度累积的边界上。
+ 您可以使用库的检查点检查使用分片数据并行度训练的模型，以及。 APIs `smp.save_checkpoint` `smp.resume_from_checkpoint`有关更多信息，请参阅 [对分布式 PyTorch 模型执行检查点操作（适用于 SageMaker 模型并行度库 v1.10.0 及更高版本）](distributed-model-parallel-checkpointing-and-finetuning.md#model-parallel-extended-features-pytorch-checkpoint)。
+ 在分片数据并行性下，[https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.delay_param_initialization](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.delay_param_initialization) 配置参数的行为会发生变化。同时启用这两个功能时，在分片方式下，将在模型创建时立即初始化参数，而不是延迟参数初始化，这样每个秩都会初始化并存储自己的参数分片。
+ 当分片数据并行性被激活时，该库会在 `optimizer.step()` 调用运行时在内部执行梯度裁剪。您无需使用实用工具 APIs 进行渐变剪裁，例如[https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html](https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html)。要调整梯度裁剪的阈值，可以在构造 SageMaker PyTorch 估计器时通过分布参数配置的参数对其进行设置，如部分所示。`sdp_gradient_clipping` [如何将分片数据并行性应用于训练作业](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use)

# 模型管道传输


模型并行度库 SageMaker的核心功能之一是*管道并行性*，它决定了模型训练期间进行计算的顺序和跨设备处理数据的顺序。Pipelining 是一种通过在不同的数据样本上同时进行计算，从而在模型并行性中实现真正的并行化，并克服顺序 GPUs 计算导致的性能损失的技术。使用管道并行性时，训练作业以管道方式按照微批次执行，以最大限度地提高 GPU 使用率。

**注意**  
管道并行度（也称为模型分区）可用于和。 PyTorch TensorFlow有关支持的框架版本，请参阅[支持的框架和 Amazon Web Services 区域](distributed-model-parallel-support.md)。

## 管道执行计划


Pipelining 基于将迷你批次拆分为微批，这些微批量被输入到训练管道中， one-by-one并遵循库运行时定义的执行时间表。*微批次*是给定训练小批次的更小的子集。管道计划决定了每个时段由哪个设备执行哪个微批次。

例如，根据流水线计划和模型分区，GPU `i` 可能会在微批量上执行（向前或向后）计算，`b`而 GPU 则在微批处理上`i+1`执行计算`b+1`，从而使两者同时 GPUs 处于活动状态。在单次向前或向后传递期间，根据分区决策，单个微批次的执行流程可能会多次访问同一个设备。例如，位于模型开头的操作，与位于模型末尾的操作可能会放在同一设备上，而两者之间的操作位于不同的设备上，这意味着此设备被访问了两次。

该库提供两种不同的流水线计划，*简单*调度和*交错排*程，可以使用 Pyth SageMaker on SDK 中的`pipeline`参数进行配置。在大多数情况下，交错流水线可以通过 GPUs 更有效地利用流水线来实现更好的性能。

### 交错管道


在交错管道中，尽可能优先考虑向后执行微批次。这样可以更快地释放内存用于激活，从而更有效地使用内存。它还允许扩大微粒的数量，从而缩短微粒的闲置时间。 GPUs在稳态下，每个设备在向前传递和向后传递之间交替运行。这意味着一个微批次的向后传递可能会在另一个微批次的向前传递完成之前运行。

![\[超过 2 的交错管道的执行计划示例。 GPUs\]](http://docs.amazonaws.cn/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-execution.png)


上图说明了超过 2 的交错管道的执行计划示例。 GPUs在图中，F0 表示微批次 0 的向前传递，B1 表示微批次 1 的向后传递。**更新**表示优化器对参数的更新。 GPU0 始终尽可能优先考虑向后传递（例如，在 F2 之前执行 B0），这样可以清除用于更早激活的内存。

### 简单管道


相比之下，简单管道则为每个微批次运行向前传递，然后再开始向后传递。这意味着它仅对自身中的向前传递和向后传递阶段进行管道传输。下图说明了其工作原理的示例，超过 2 GPUs.

![\[在管道上的示例，每个微批次都先运行前向传递，然后再开始后向传递。\]](http://docs.amazonaws.cn/sagemaker/latest/dg/images/distributed/model-parallel/simple-pipeline-execution.png)


### 特定框架中的管道执行


使用以下章节来了解特定于框架的管道调度决策 SageMaker的模型并行度库对和的影响。 TensorFlow PyTorch

#### 使用管道执行 TensorFlow


下图是使用自动模型拆分由模型并行度库划分的 TensorFlow 图形的示例。拆分图形时，生成的每个子图形都会被复制 B 次（变量除外），其中 B 是微批次的数量。在此图中，每个子图形被复制 2 次 (B=2)。在子图形的每个输入中插入一个 `SMPInput` 运算，在每个输出中插入一个 `SMPOutput` 运算。这些操作与库后端通信，以在彼此之间传输张量。

![\[使用自动模型拆分由模型并行度库划分的 TensorFlow 图形示例。\]](http://docs.amazonaws.cn/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-tf.png)


下图是 2 个子图形的示例，以 B=2 进行拆分，并添加了梯度运算。`SMPInput` 运算的梯度是 `SMPOutput` 运算，反之亦然。这使梯度能够在反向传播期间向后流动。

![\[2 个子图形的示例，以 B=2 进行拆分，并添加了梯度运算。\]](http://docs.amazonaws.cn/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-tf.gif)


此 GIF 演示了一个示例交错管道执行计划，其中 B=2 个微批次和 2 个子图形。每个设备按顺序执行其中一个子图形副本，以提高 GPU 利用率。随着 B 增大，空闲时段的比例趋近于零。每当需要对特定的子图形副本执行（向前或向后）计算时，管道层都会向相应的蓝色 `SMPInput` 运算发出信号以开始执行。

计算出单个小批次中所有微批次的梯度后，该库就会合并各个库的梯度，然后将其应用于参数。

#### 使用管道执行 PyTorch


从概念上讲，流水线也遵循类似的想法。 PyTorch但是，由于 PyTorch 不涉及静态图，因此模型并行度库的 PyTorch 功能使用了更动态的流水线模式。

例如 TensorFlow，每个批次都被分成多个微粒，这些微批次在每个设备上一次执行一个。但是，执行计划由在每个设备上启动的执行服务器来处理。只要当前设备需要放置在另一台设备上的子模块的输出，就会向远程设备的执行服务器发送执行请求，同时将输入张量发送到该子模块。然后，服务器使用给定的输入执行此模块，并将响应返回给当前设备。

由于当前设备在远程子模块执行期间处于空闲状态，当前微批处理的本地执行暂停，库运行时系统会将执行切换到当前设备可以主动处理的另一个微批次。微批次的优先级由所选的管道计划决定。对于交错管道计划，只要有可能，就会优先考虑处于计算后退阶段的微批次。

# 张量并行性


*张量并行性*是模型并行性的一种，它在设备之间拆分特定的模型权重、梯度和优化器状态。管道并行性保持单个权重不变但对权重*集*进行拆分，张量并行性则与之相反，会拆分单个权重。这通常涉及对模型的特定运算、模块或层进行分布式计算。

在单个参数占用大部分 GPU 内存的情况下（例如词汇表很大的大型嵌入表，或者具有大量类的大型 softmax 层），则需要张量并行性。在这种情况下，将这种大张量或运算视为原子单元的效率会很低，并且会阻碍内存负载的平衡。

在仅仅依靠管道并不足以满足要求的超大型模型中，张量并行性也很有用。例如，在 GPT-3 规模的模型中需要对数十个实例进行分区，仅使用微批次管道传输效率低下，因为管道深度会变得过深，开销会变得过大。

**注意**  
张量并行度可在 SageMaker 模型并行度库 v1.6.0 及更高版本 PyTorch 中使用。

**Topics**
+ [

# 张量并行性的工作原理
](model-parallel-extended-features-pytorch-tensor-parallelism-how-it-works.md)
+ [

# 使用张量并行度运行 SageMaker 分布式模型并行训练 Job
](model-parallel-extended-features-pytorch-tensor-parallelism-examples.md)
+ [

# 支持 Hugging Face 转换器模型
](model-parallel-extended-features-pytorch-hugging-face.md)
+ [

# 将管道并行性与张量并行性结合使用时的秩评定机制
](model-parallel-extended-features-pytorch-ranking-mechanism.md)

# 张量并行性的工作原理
工作方式

张量并行性在级别 `nn.Modules` 上实现；它在张量并行秩之间对模型中的特定模块分区。这是对管道并行性中使用的*模块集*的现有分区的补充。

通过张量并行性对模块进行分区时，其向前和向后传播是分布式的。库处理设备间的必要通信，以实施这些模块的分布式执行。这些模块在多个数据并行秩之间进行分区。与传统的工作负载分布相反，在使用库的张量并行性时，每个数据并行秩并**没有**完整的模型副本。相反，在所有未分布而保持完整的模块之外，每个数据并行秩可能只有分布式模块的一个分区。

**示例：**考虑跨数据并行秩的张量并行性，其中数据并行度为 4，张量并行度为 2。假设在对模块集进行分区后，您有一个存有以下模块树的数据并行组。

```
A
├── B
|   ├── E
|   ├── F
├── C
└── D
    ├── G
    └── H
```

假设模块 B、G 和 H 支持张量并行性。此模型的张量并行分区的一种可能结果可能是：

```
dp_rank 0 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 1 (tensor parallel rank 1): A, B:1, C, D, G:1, H
dp_rank 2 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 3 (tensor parallel rank 1): A, B:1, C, D, G:1, H
```

每行代表存储在该 `dp_rank` 中的一组模块，标记 `X:y` 表示模块 `X` 的第 `y` 个部分。注意以下几点：

1. 分区跨数据并行秩的子集进行，我们调用 `TP_GROUP`，而不是整个 `DP_GROUP`，因此精确的模型分区会跨 `dp_rank` 0 和 `dp_rank` 2 复制，类似地跨 `dp_rank` 1 和 `dp_rank` 3 复制。

1. 模块 `E` 和 `F` 不再是模型的一部分，因为它们的父模块 `B` 已分区，并且通常属于 `E` 和 `F` 的任何执行在（已分区）`B` 模块内进行。

1. 尽管 `H` 支持张量并行性，但在此示例中，它没有进行分区，这强调了是否对模块进行分区取决于用户输入。模块支持张量并行性并不一定意味着它需要分区。

## 库如何调整张量并行度以适应模块 PyTorch `nn.Linear`


对数据并行秩执行张量并行性时，*对于所分区的模块*，参数、梯度和优化器状态的子集在张量并行设备上分区。对于其余模块，张量并行设备以常规的数据并行方式操作。为了执行分区模块，设备首先在同一个张量并行度组中，跨对等设备收集*所有数据样本*的必要部分。然后，设备在所有这些数据样本上运行模块的本地部分，然后进行另一轮同步，该同步既合并了每个数据样本的输出部分，又将组合的数据样本返回到数据样本最初来源的样本。 GPUs 下图显示了在已分区 `nn.Linear` 模块上执行此过程的示例。

![\[两张图显示了两个张量并行概念。\]](http://docs.amazonaws.cn/sagemaker/latest/dg/images/distributed/model-parallel/tensor-parallel-concept.png)


第一张图显示了一个带有 `nn.Linear` 模块的小模型，带有在两个张量并行性秩上的数据并行性。`nn.Linear` 模块复制到两个并行秩中。

第二张图显示了在拆分 `nn.Linear` 模块时应用于较大模型的张量并行性。每个 `tp_rank` 存有一半的线性模块，以及剩余的全部操作。当线性模块运行时，每个 `tp_rank` 都会收集所有数据样本中相关的一半，并传递到自己的一半 `nn.Linear` 模块中。结果必须是减少分散的（使用求和作为归约运算），这样每个秩都有自己数据样本的最终线性输出。模型的其余部分以典型的数据并行方式运行。

# 使用张量并行度运行 SageMaker 分布式模型并行训练 Job
使用张量并行性运行训练作业

在本部分中，您将学习：
+ 如何配置 SageMaker PyTorch 估计器和 SageMaker 模型并行度选项以使用张量并行度。
+ 如何使用扩展的 `smdistributed.modelparallel` 模块来调整训练脚本以用于张量并行性。

要了解有关这些`smdistributed.modelparallel`模块的更多信息，请参阅 *SageMaker Python SDK 文档 APIs*中的[并行SageMaker 模型](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)。

**Topics**
+ [

## 仅使用张量并行性
](#model-parallel-extended-features-pytorch-tensor-parallelism-alone)
+ [

## 张量并行性与管道并行性相结合
](#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism)

## 仅使用张量并行性


以下分布式训练选项示例单独激活张量并行性，不使用管道并行性。配置`mpi_options`和`smp_options`字典以为 SageMaker `PyTorch`估计器指定分布式训练选项。

**注意**  
扩展的内存节省功能可通过 Deep Learning Containers for 获得 PyTorch，该容器实现了 SageMaker 模型并行度库 v1.6.0 或更高版本。

**配置 SageMaker PyTorch 估算器**

```
mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8,               # 8 processes
    "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}
               
smp_options = {
    "enabled":True,
    "parameters": {
        "pipeline_parallel_degree": 1,    # alias for "partitions"
        "placement_strategy": "cluster",
        "tensor_parallel_degree": 4,      # tp over 4 devices
        "ddp": True
    }
}
              
smp_estimator = PyTorch(
    entry_point='your_training_script.py', # Specify
    role=role,
    instance_type='ml.p3.16xlarge',
    sagemaker_session=sagemaker_session,
    framework_version='1.13.1',
    py_version='py36',
    instance_count=1,
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": mpi_options
    },
    base_job_name="SMD-MP-demo",
)

smp_estimator.fit('s3://my_bucket/my_training_data/')
```

**提示**  
要查找的完整参数列表`distribution`，请参阅 Pyth SageMaker on SDK 文档中的[模型并行度配置参数](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html)。

**调整您的 PyTorch 训练脚本**

以下示例训练脚本展示了如何根据训练脚本调整 SageMaker 模型并行度库。在此示例中，假设脚本命名为 `your_training_script.py`。

```
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchnet.dataset import SplitDataset
from torchvision import datasets

import smdistributed.modelparallel.torch as smp

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        return F.log_softmax(x, 1)

def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by
        # the current process, based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target, reduction="mean")
        loss.backward()
        optimizer.step()

# smdistributed: Initialize the backend
smp.init()

# smdistributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")

# smdistributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
if smp.local_rank() == 0:
    dataset = datasets.MNIST("../data", train=True, download=False)
smp.barrier()

# smdistributed: Shard the dataset based on data parallel ranks
if smp.dp_size() > 1:
    partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}
    dataset = SplitDataset(dataset, partitions=partitions_dict)
    dataset.select(f"{smp.dp_rank()}")

train_loader = torch.utils.data.DataLoader(dataset, batch_size=64)

# smdistributed: Enable tensor parallelism for all supported modules in the model
# i.e., nn.Linear in this case. Alternatively, we can use
# smp.set_tensor_parallelism(model.fc1, True)
# to enable it only for model.fc1
with smp.tensor_parallelism():
    model = Net()

# smdistributed: Use the DistributedModel wrapper to distribute the
# modules for which tensor parallelism is enabled
model = smp.DistributedModel(model)

optimizer = optim.AdaDelta(model.parameters(), lr=4.0)
optimizer = smp.DistributedOptimizer(optimizer)

train(model, device, train_loader, optimizer)
```

## 张量并行性与管道并行性相结合


以下是一个分布式训练选项的示例，该选项支持张量并行性与流水线并行性相结合。 设置`mpi_options`和`smp_options`参数，以便在配置估计器时使用张量并行度指定模型并行度选项。 SageMaker `PyTorch`

**注意**  
扩展的内存节省功能可通过 Deep Learning Containers for 获得 PyTorch，该容器实现了 SageMaker 模型并行度库 v1.6.0 或更高版本。

**配置 SageMaker PyTorch 估算器**

```
mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8,               # 8 processes
    "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}
               
smp_options = {
    "enabled":True,
    "parameters": {
    "microbatches": 4,
        "pipeline_parallel_degree": 2,    # alias for "partitions"
        "placement_strategy": "cluster",
        "tensor_parallel_degree": 2,      # tp over 2 devices
        "ddp": True
    }
}
              
smp_estimator = PyTorch(
    entry_point='your_training_script.py', # Specify
    role=role,
    instance_type='ml.p3.16xlarge',
    sagemaker_session=sagemaker_session,
    framework_version='1.13.1',
    py_version='py36',
    instance_count=1,
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": mpi_options
    },
    base_job_name="SMD-MP-demo",
)

smp_estimator.fit('s3://my_bucket/my_training_data/')  
```

<a name="model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism-script"></a>**调整您的 PyTorch 训练脚本**

以下示例训练脚本展示了如何根据训练脚本调整 SageMaker 模型并行度库。请注意，训练脚本现在包括 `smp.step` 修饰器：

```
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchnet.dataset import SplitDataset
from torchvision import datasets

import smdistributed.modelparallel.torch as smp

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        return F.log_softmax(x, 1)


# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(model, data, target):
    output = model(data)
    loss = F.nll_loss(output, target, reduction="mean")
    model.backward(loss)
    return output, loss

def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by
        # the current process, based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # Return value, loss_mb is a StepOutput object
        _, loss_mb = train_step(model, data, target)

        # smdistributed: Average the loss across microbatches.
        loss = loss_mb.reduce_mean()

        optimizer.step()

# smdistributed: Initialize the backend
smp.init()

# smdistributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")

# smdistributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
if smp.local_rank() == 0:
    dataset = datasets.MNIST("../data", train=True, download=False)
smp.barrier()

# smdistributed: Shard the dataset based on data parallel ranks
if smp.dp_size() > 1:
    partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}
    dataset = SplitDataset(dataset, partitions=partitions_dict)
    dataset.select(f"{smp.dp_rank()}")

# smdistributed: Set drop_last=True to ensure that batch size is always divisible
# by the number of microbatches
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

model = Net()

# smdistributed: enable tensor parallelism only for model.fc1
smp.set_tensor_parallelism(model.fc1, True)

# smdistributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)

optimizer = optim.AdaDelta(model.parameters(), lr=4.0)
optimizer = smp.DistributedOptimizer(optimizer)

train(model, device, train_loader, optimizer)
```

# 支持 Hugging Face 转换器模型


 SageMaker 模型并行度库的张量并行度为以下 Hugging Face Transformer 模型提供 out-of-the-box支持：
+ GPT-2、BERT 和 RoBERTa （在 SageMaker 模型并行度库 1.7.0 及更高版本中可用）
+ GPT-J（在 SageMaker 模型并行度库 1.8.0 及更高版本中可用）
+ GPT-neo（在 SageMaker 模型并行度库 v1.10.0 及更高版本中可用）

**注意**  
对于任何其他转换器模型，您需要使用 [smdistributed.modelparallel.torch.tp\$1register\$1with\$1module()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch_tensor_parallel.html#smdistributed.modelparallel.torch.tp_register_with_module) API 来应用张量并行性。

**注意**  
要使用张量并行度来训练 Hugging Face Transformer 模型，请务必使用 SageMaker具有模型并行度库 v1.7.0 及更高版本的 Hugging Face Deep Lear PyTorch ning Containers。有关更多信息，请参阅[SageMaker 模型并行度库发行说明](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_release_notes/smd_model_parallel_change_log.html)。

## 现成支持的模型


对于开箱即用库支持的 Hugging Face 变压器模型，您无需手动实现挂钩即可将 Transformer `smdistributed` 转换为 APIs 变压器层。[你可以使用上下文管理器 smdistributed.modelparallel.torch.tensor\$1parallelism () 激活张量并行性，然后用 [smdistributed.modelparallel.torch 包装模型。](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch_tensor_parallel.html#smdistributed.modelparallel.torch.tensor_parallelism) DistributedModel()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.DistributedModel)。您无需使用 `smp.tp_register` API 手动注册钩子用于张量并行性。

Hugging Face Transformers 与 `smdistributed.modelparallel` 之间的 `state_dict` 转换函数可以如下所示访问。
+  `smdistributed.modelparallel.torch.nn.huggingface.gpt2.translate_state_dict_to_hf_gpt2(state_dict, max_seq_len=None)`
+  `smdistributed.modelparallel.torch.nn.huggingface.gpt2.translate_hf_state_dict_to_smdistributed_gpt2(state_dict)` 
+  `smdistributed.modelparallel.torch.nn.huggingface.bert.translate_state_dict_to_hf_bert(state_dict, max_seq_len=None)` 
+  `smdistributed.modelparallel.torch.nn.huggingface.bert.translate_hf_state_dict_to_smdistributed_bert(state_dict)` 
+  `smdistributed.modelparallel.torch.nn.huggingface.roberta.translate_state_dict_to_hf_roberta(state_dict, max_seq_len=None)` 
+  `smdistributed.modelparallel.torch.nn.huggingface.roberta.translate_hf_state_dict_to_smdistributed_roberta(state_dict)` 
+ `smdistributed.modelparallel.torch.nn.huggingface.gptj.translate_state_dict_to_hf_gptj(state_dict, max_seq_len=None)`（在 SageMaker 模型并行度库 v1.8.0 及更高版本中可用）
+ `smdistributed.modelparallel.torch.nn.huggingface.gptj.translate_hf_gptj_state_dict_to_smdistributed_gptj`（在 SageMaker 模型并行度库 v1.8.0 及更高版本中可用）
+ `smdistributed.modelparallel.torch.nn.huggingface.gptneo.translate_state_dict_to_hf_gptneo(state_dict, max_seq_len=None)`（在 SageMaker 模型并行度库 v1.10.0 及更高版本中可用）
+ `smdistributed.modelparallel.torch.nn.huggingface.gptneo.translate_hf_state_dict_to_smdistributed_gptneo(state_dict)`（在 SageMaker 模型并行度库 v1.10.0 及更高版本中可用）

**GPT-2 转换函数的使用示例**

首先包装模型，如以下代码所示。

```
from transformers import AutoModelForCausalLM

with smp.tensor_parallelism():
    model = AutoModelForCausalLM.from_config(hf_gpt2_config)

model = smp.DistributedModel(model)
```

对于 `DistributedModel` 对象中的 `state_dict`，您可以使用 `translate_state_dict_to_hf_gpt2` 函数将权重加载到原始 Hugging Face GPT-2 模型中，如以下代码所示。

```
from smdistributed.modelparallel.torch.nn.huggingface.gpt2 \
                                      import translate_state_dict_to_hf_gpt2
max_seq_len = 1024

# [... code block for training ...]

if smp.rdp_rank() == 0:
    state_dict = dist_model.state_dict()
    hf_state_dict = translate_state_dict_to_hf_gpt2(state_dict, max_seq_len)

    # can now call model.load_state_dict(hf_state_dict) to the original HF model
```

**Ro BERTa 翻译函数的用法示例**

同样，给定支持的 HuggingFace 模型`state_dict`，您可以使用`translate_hf_state_dict_to_smdistributed`函数将其转换为可读的格式`smp.DistributedModel`。这在迁移学习使用场景中非常有用，此时预训练的模型加载到 `smp.DistributedModel` 中用于模型并行微调：

```
from smdistributed.modelparallel.torch.nn.huggingface.roberta \
                                      import translate_state_dict_to_smdistributed

model = AutoModelForMaskedLM.from_config(roberta_config)
model = smp.DistributedModel(model)

pretrained_model = AutoModelForMaskedLM.from_pretrained("roberta-large")
translated_state_dict =
        translate_state_dict_to_smdistributed(pretrained_model.state_dict())

# load the translated pretrained weights into the smp.DistributedModel
model.load_state_dict(translated_state_dict)

# start fine-tuning...
```

# 将管道并行性与张量并行性结合使用时的秩评定机制
秩评定机制

此部分解释了模型并行性的秩评定机制如何与张量并行性结合使用。这是从 [SageMaker 模型并行度库的核心功能](model-parallel-core-features.md) 的[秩评定基础知识](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#ranking-basics)中扩展而来的。在张量并行度中，该库引入了三种类型的排名和处理组 APIs：用于`smp.tp_rank()`张量并行等级，用于`smp.pp_rank()`流水线并行等级，以及用于`smp.rdp_rank()`简化数据的并行等级。对应的通信进程组是张量并行组 (`TP_GROUP`)、管道并行组 (`PP_GROUP`) 和缩减数据并行组 (`RDP_GROUP`)。这些组的定义如下：
+ *张量并行组* (`TP_GROUP`) 是数据并行组中一个可均匀分割的子集，在其中完成模块的张量并行分布。当管道并行度为 1 时，`TP_GROUP` 与*模型并行组* (`MP_GROUP`) 相同。
+ *管道并行组* (`PP_GROUP`) 是完成管道并行性的进程组。当张量并行度为 1 时，`PP_GROUP` 与 `MP_GROUP` 相同。
+ *缩减数据并行组* (`RDP_GROUP`) 是一组进程，同时容纳相同的管道并行性分区和相同的张量并行分区，并在它们自身中执行数据并行性。之所以将其称为缩减数据并行组，是因为它是整个数据并行性组 `DP_GROUP` 的子集。对于分布在 `TP_GROUP` 中的模型参数，梯度 `allreduce` 运算仅对缩减数据并行组执行，而对于未分布的参数，梯度 `allreduce` 在整个 `DP_GROUP` 上进行。
+ 模型并行组 (`MP_GROUP`) 是指一组共同存储整个模型的进程。它由当前进程的 `TP_GROUP` 中，所有秩的 `PP_GROUP` 并集组成。当张量并行度为 1 时，`MP_GROUP` 等于 `PP_GROUP`。它也与先前 `smdistributed` 版本中的 `MP_GROUP` 现有定义一致。请注意，当前 `TP_GROUP` 是当前 `DP_GROUP` 和当前 `MP_GROUP` 的子集。

要详细了解 SageMaker 模型并行度库 APIs 中的通信过程，请参阅 Pyth *SageMaker on* SDK 文档中的[通用 API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_common_api.html#) 和[PyTorch特定 APIs](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html)的。

![\[张量并行性的排名机制、参数分布和关联 AllReduce 运算。\]](http://docs.amazonaws.cn/sagemaker/latest/dg/images/distributed/model-parallel/tensor-parallel-ranking-mechanism.png)


例如，考虑具有8的单个节点的处理组 GPUs，其中张量并行度为2，流水线并行度为2，数据并行度为4。上图的顶部居中部分显示了一个包含 4 层的模型的示例。图的左下角和右下部分说明了 GPUs 使用流水线并行性和张量并行度分布在 4 层的 4 层模型，其中中间的两层使用张量并行性。下方的两个图是简单的副本，用于说明不同的组边界线。在 GPUs 0-3 和 4-7 之间复制分区模型以实现数据并行性。左下图显示了`MP_GROUP`、`PP_GROUP` 和 `TP_GROUP` 的定义。右下角的图显示`WORLD`了`RDP_GROUP``DP_GROUP`、和在同一组上 GPUs。为了实现数据并行性，具有相同颜色的层和层切片的梯度通过 `allreduce` 分在一起。例如，第一层（浅蓝色）获取 `DP_GROUP` 上的 `allreduce` 运算，而第二层中的深橙色切片只能获取其进程的 `RDP_GROUP` 中的 `allreduce` 运算。加粗深红色箭头表示张量及其完整 `TP_GROUP` 的批次。

```
GPU0: pp_rank 0, tp_rank 0, rdp_rank 0, dp_rank 0, mp_rank 0
GPU1: pp_rank 1, tp_rank 0, rdp_rank 0, dp_rank 0, mp_rank 1
GPU2: pp_rank 0, tp_rank 1, rdp_rank 0, dp_rank 1, mp_rank 2
GPU3: pp_rank 1, tp_rank 1, rdp_rank 0, dp_rank 1, mp_rank 3
GPU4: pp_rank 0, tp_rank 0, rdp_rank 1, dp_rank 2, mp_rank 0
GPU5: pp_rank 1, tp_rank 0, rdp_rank 1, dp_rank 2, mp_rank 1
GPU6: pp_rank 0, tp_rank 1, rdp_rank 1, dp_rank 3, mp_rank 2
GPU7: pp_rank 1, tp_rank 1, rdp_rank 1, dp_rank 3, mp_rank 3
```

在此示例中，管道并行性跨 GPU 对 (0,1)、(2,3)、(4,5) 和 (6,7) 进行。此外，数据并行度 (`allreduce`) 在 GPUs 0、2、4、6 之间进行，并在 GPUs 1、3、5、7 上独立进行。张量并行性发生在 `DP_GROUP` 的子集上，跨 GPU 对 (0,2)、(1,3)、(4,6) 和 (5,7)。

  对于这种管道和张量并行性的混合，`data_parallel_degree` 的算术仍然是 `data_parallel_degree = number_of_GPUs / pipeline_parallel_degree`。该库根据关系 `reduced_data_parallel_degree * tensor_parallel_degree = data_parallel_degree`，进一步计算缩减数据并行度。  

# 优化器状态分片


*优化器状态分片*是一种非常有用的内存节省技术，它可以跨数据并行设备组对优化器状态（描述优化器状态的一组权重）进行分片。每当使用有状态优化器（例如 Adam）或优化器（存储参数 FP16 和参数 FP32 副本）时，都可以使用 FP16 优化器状态分片。

**注意**  
优化器状态分片可在 SageMaker 模型并行度库 v1.6.0 及更高版本 PyTorch 中使用。

## 如何使用优化器状态分片


您可以通过在 `modelparallel` 配置设置 `"shard_optimizer_state": True` 来启用*优化器状态分片*。

启用此功能后，库会根据数据并行度对模型参数集进行分区。与第 `i` 个分区对应的梯度，仅在第 `i` 个数据并行秩处缩减。在对 `smp.step` 修饰器函数的第一次调用结束时，被 `smp.DistributedOptimizer` 包装的优化器重新定义了其参数，使其仅限于与当前数据并行秩的分区相对应的参数。重新定义的参数称为*虚拟参数*，它们与原始参数共享底层存储。在第一次调用 `optimizer.step` 期间，优化器状态是根据这些重新定义的参数创建的，这些状态由于原始分区而被分片。优化器更新后，该 AllGather操作（作为`optimizer.step`调用的一部分）在数据 parallel 等级中运行，以实现一致的参数状态。

**提示**  
当数据并行度大于 1 且模型的参数超过 10 亿时，优化器状态分片可能很有用。  
数据并行性由 `(processes_per_host * instance_count / pipeline_parallel_degree)` 计算得出，`smp.dp_size()` 函数在后台处理大小调整。

**配置 SageMaker PyTorch 估算器**

```
mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8,               # 8 processes
    "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}

smp_options = {
    "enabled":True,
    "parameters": {
        "microbatches": 4,
        "pipeline_parallel_degree": 2,    # alias for "partitions"
        "placement_strategy": "cluster",
        "tensor_parallel_degree": 2,      # tp over 2 devices
        "ddp": True,
        "shard_optimizer_state": True
    }
}
```

**调整您的 PyTorch 训练脚本**

请参阅 *Tensor 并行度与流水线并行度相结合部分*中的[调整 PyTorch 训练脚本](model-parallel-extended-features-pytorch-tensor-parallelism-examples.md#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism-script)。无需对脚本进行其他修改。

# 激活检查点


*激活检查点*（或*梯度检查点*）技术通过清除某些层的激活并在向后传递期间重新计算它们，来减少内存使用量。实际上，这是用额外的计算时间来换取内存使用量的减少。如果对模块执行了检查点操作，则在向前传递结束时，该模块的输入和输出将保留在内存中。在向前传递期间，任何本应是模块内部计算一部分的中间张量都会被释放。在有检查点的模块的向后传递过程中，会重新计算这些张量。此时，有检查点的模块之外的层已经完成其向后传递，因此检查点操作的峰值内存使用量可能会更低。

**注意**  
此功能可在 SageMaker 模型并行度库 v1.6.0 及更高版本 PyTorch 中使用。

## 如何使用激活检查点


使用 `smdistributed.modelparallel`，您可以按模块使用激活检查点。对于除 `torch.nn.Sequential` 之外的所有 `torch.nn` 模块，只有当从管道并行性的角度来看，模块树位于一个分区内时，才能对模块树执行检查点操作。对于 `torch.nn.Sequential` 模块，顺序模块内的每个模块树必须完全位于一个分区内，激活检查点才能起作用。使用手动分区时，请注意这些限制。

使用[自动模型分区](https://docs.amazonaws.cn/sagemaker/latest/dg/model-parallel-core-features.html#model-parallel-automated-model-splitting)时，您可在训练作业日志中找到以 `Partition assignments:` 开头的分区分配日志。如果一个模块在多个秩（例如，一个后代属于一个秩，另一个后代处于不同的秩）上分区，则库会忽略对模块执行检查点的尝试，并发出一条警告消息，说明该模块没有检查点。

**注意**  
 SageMaker 模型并行度库支持重叠和非重叠操作以及检查点`allreduce`操作。

**注意**  
PyTorch的本机检查点 API 与不兼容。`smdistributed.modelparallel`

**示例 1：**以下示例代码演示了当脚本中有模型定义时，如何使用激活检查点操作。

```
import torch.nn as nn
import torch.nn.functional as F

from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        # This call of fc1 will be checkpointed
        x = checkpoint(self.fc1, x)
        x = self.fc2(x)
        return F.log_softmax(x, 1)
```

**示例 2：**以下示例代码演示了当脚本中有顺序模型时，如何使用激活检查点操作。

```
import torch.nn as nn
from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint_sequential

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.seq = nn.Sequential(
            nn.Conv2d(1,20,5),
            nn.ReLU(),
            nn.Conv2d(20,64,5),
            nn.ReLU()
        )

    def forward(self, x):
        # This call of self.seq will be checkpointed
        x = checkpoint_sequential(self.seq, x)
        return F.log_softmax(x, 1)
```

**示例 3：**以下示例代码显示了在从库（例如和 Hugging Face Transformers PyTorch ）导入预建模型时如何使用激活检查点。无论您是否对顺序模型执行检查点操作，请完成以下过程：

1. 使用 `smp.DistributedModel()` 包装模型。

1. 为顺序层定义一个对象。

1. 使用 `smp.set_activation_checkpointig()` 包装顺序层对象。

```
import smdistributed.modelparallel.torch as smp
from transformers import AutoModelForCausalLM

smp.init()
model = AutoModelForCausalLM(*args, **kwargs)
model = smp.DistributedModel(model)

# Call set_activation_checkpointing API
transformer_layers = model.module.module.module.transformer.seq_layers
smp.set_activation_checkpointing(
    transformer_layers, pack_args_as_tuple=True, strategy='each')
```

# 激活分载


当激活检查点和管道并行性均已启用并且微批次数量大于 1 时，*激活分载*是可以进一步减少内存使用量的附加功能。对于当前未在 CPU 中运行的微批次，激活分载会异步移动与其对应的有检查点的激活。就在 GPU 需要激活以便微批次向后传递之前，此功能会将分载的激活从 CPU 预取回来。

**注意**  
此功能可在 SageMaker 模型并行度库 v1.6.0 及更高版本 PyTorch 中使用。

## 如何使用激活分载


当**微批次的数量大于 1 并启用了激活检查点时**，请使用激活分载以减少内存使用量（请参阅[激活检查点](model-parallel-extended-features-pytorch-activation-checkpointing.md)）。如果不使用激活检查点，则激活分享不会生效。在仅与一个微批次一起使用时，这不会节省内存。

要使用激活分载，请在 `modelparallel` 配置中设置 `"offload_activations": True`。

激活分载将 `nn.Sequential` 模块中的有检查点的激活异步移至 CPU。通过 PCIe 链路传输的数据与 GPU 计算重叠。在计算了特定检查点层的向前传递后，会立即进行分载。在特定微批次的向后传递需要激活之前的片刻，激活将加载回 GPU。CPU-GPU 传输同样与计算重叠。

要调整激活加载回 GPU 的时间提前量，您可以使用配置参数 `"activation_loading_horizon"`（默认设置为 4，必须为大于 0 的 `int`）。较大的激活加载范围会导致激活更早地加载回 GPU。如果范围太大，激活分载所带来的内存节省影响可能会减弱。如果范围太小，激活可能无法及时加载回，从而减少重叠量，导致性能下降。

**提示**  
激活分载对于参数超过一千亿的大型模型非常有用。

**配置 SageMaker PyTorch 估算器**

```
mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8,               # 8 processes
    "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}

smp_options = {
    "enabled":True,
    "parameters": {
        "microbatches": 4,
        "pipeline_parallel_degree": 2,    # alias for "partitions"
        "placement_strategy": "cluster",
        "tensor_parallel_degree": 2,      # tp over 2 devices
        "ddp": True,
        "offload_activations": True,
        "activation_loading_horizon": 4   # optional. default is 4.
    }
}
```

# FP16 使用模型并行度进行训练


要进行 FP16 训练，请对训练脚本和估算器进行以下修改。

**注意**  
此功能可在 SageMaker 模型并行度库 v1.10.0 及更高版本 PyTorch 中使用。

**调整您的 PyTorch 训练脚本**

1. 使用 [smdistributed.modelparallel.torch.model\$1creation()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.model_creation) 上下文管理器包装您的模型。

   ```
   # fp16_training_script.py
   
   import torch
   import smdistributed.modelparallel.torch as smp
   
   with smp.model_creation(
       dtype=torch.float16 if args.fp16 else torch.get_default_dtype()
   ):
       model = ...
   ```
**提示**  
如果您使用的是张量并行性，请将 `tensor_parallelism=smp.tp_size() > 1` 添加到 `smp.model_creation` 上下文管理器中。添加此行还有助于自动检测是否已激活张量并行性。  

   ```
   with smp.model_creation(
       ... ,
       tensor_parallelism=smp.tp_size() > 1
   ):
       model = ...
   ```

1. 当您使用 `smdistributed.modelparallel.torch.DistributedOptimizer` 包装优化器时，请设置 `static_loss_scaling` 或 `dynamic_loss_scaling` 参数。默认情况下，`static_loss_scaling` 设置为 `1.0`，`dynamic_loss_scaling` 设置为 `False`。如果您设置 `dynamic_loss_scale=True`，则可以通过 `dynamic_loss_args` 参数将动态损失缩放选项作为字典输入。在大多数情况下，我们建议您使用带有默认选项的动态损失缩放。[有关优化器包装器函数的更多信息、选项和示例，请参阅 smdistributed.modelparallel.torch。 DistributedOptimizer](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed-modelparallel-torch-distributedoptimizer)API。

   以下代码是使用动态损失缩放来包装`Adadelta`优化器对象以进行 FP16 训练的示例。

   ```
   optimizer = torch.optim.Adadelta(...)
   optimizer = smp.DistributedOptimizer(
       optimizer,
       static_loss_scale=None,
       dynamic_loss_scale=True,
       dynamic_loss_args={
           "scale_window": 1000,
           "min_scale": 1,
           "delayed_shift": 2
       }
   )
   ```

**配置 SageMaker PyTorch 估算器**

在创建 SageMaker PyTorch 估计器对象时，将 FP16 参数 (`"fp16"`) 添加到分布配置中以实现模型并行性。有关模型并行性配置参数的完整列表，请参阅 [`smdistributed` 的参数](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#parameters-for-smdistributed)。

```
from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled": True,
    "parameters":  {
        "microbatches":  4,
        "pipeline_parallel_degree":  2,
        "tensor_parallel_degree":  2,
        ...,

        "fp16": True
    }
}

fp16_estimator = PyTorch(
    entry_point="fp16_training_script.py", # Specify your train script
    ...,

    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": {...}
    }
)

fp16_estimator.fit(...)
```

 FP16 训练开始时，模型和优化器`FP16_Optimizer`分别由`FP16_Module`和包装，它们是 [Apex](https://nvidia.github.io/apex/fp16_utils.html#apex-fp16-utils) 实用程序的修改`smdistributed`版本。 `FP16_Module`将模型转换为 FP16 dtype 并处理向前传入。 FP16

**提示**  
您可以在 `optimizer.step` 之前通过调用 `clip_master_grads` 来应用梯度剪裁。  

```
optimizer.clip_master_grads(max_norm)     # max_norm(float or int): max norm of the gradients
```

**提示**  
使用`torch.optim.lr_scheduler`和 FP16 训练时，需要传递`optimizer.optimizer`给 LR 调度器而不是优化器。请参阅以下示例代码。  

```
from torch.optim.lr_scheduler import StepLR

scheduler = StepLR(
    optimizer.optimizer if smp.state.cfg.fp16 else optimizer,
    step_size=1,
    gamma=args.gamma
)
```

# Support FlashAttention


Suppor FlashAttention t for 是该库的一项功能，仅适用于*分布式变压器模型，分布式变压器*模型是为模型并行训练而封装的 Trans [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed-modelparallel-torch-distributedmodel](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed-modelparallel-torch-distributedmodel)former 模型。此功能还与 [张量并行性](model-parallel-extended-features-pytorch-tensor-parallelism.md) 兼容。

该[FlashAttention](https://github.com/HazyResearch/flash-attention)库仅在设置`attention_head_size`为 8 的倍数且小于 128 的值时才支持模型。因此，在训练分布式变压器并确保其 FlashAttention 正常工作时，应调整参数以使注意力头大小符合要求。有关更多信息，另请参阅*FlashAttention GitHub存储库*中的[安装和功能](https://github.com/HazyResearch/flash-attention#installation-and-features)。

例如，假设您使用 `hidden_width=864` 和 `num_heads=48` 配置转换器模型。的头部大小计算公式 FlashAttention 为`attention_head_size = hidden_width / num_heads = 864 / 48 = 18`。要启用 FlashAttention，您需要将`num_heads`参数调整为 `54``attention_head_size = hidden_width / num_heads = 864 / 54 = 16`，即 8 的倍数。