Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

分片数据并行性

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

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

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

SageMaker 通过实施 MiCS 来实现分片数据并行性，Amazon 博客文章在 Amazon 上以近线性扩展方式训练巨型模型介绍了此内容。在此实施中，您可以将分片度设置为可配置参数，该参数必须小于数据并行度。在每次向前和向后传递过程中，MiCS 都会通过 AllGather 操作临时重新组合所有 GPU 中的模型参数。在每层向前或向后传递后，MiCS 会再次对参数进行分片以节省 GPU 内存。在向后传递过程中，MiCS 减少梯度并同时通过 ReduceScatter 操作在多个 GPU 上对其分片。最后，MiCS 使用优化器状态的局部分片，将局部缩减梯度和分片梯度应用于其对应的局部参数分片。为了降低通信开销，SageMaker 模型并行性库在向前或向后传递中会预提取接下来使用的层，并将网络通信与计算重叠。

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

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

选择分片数据并行度

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

选择批次大小

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

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

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

调整 PyTorch 训练脚本

按照步骤 1：修改 PyTorch 训练脚本中的说明，使用 torch.nn.parallel 和 torch.distributed 模块的 smdistributed.modelparallel.torch 包装器来包装模型和优化器对象。

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

如果您的模型使用 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 估算器

在 步骤 2：使用 SageMaker Python SDK 启动训练作业 中配置 SageMaker PyTorch 估算器时，请添加分片数据并行性的参数。

要启用分片数据并行性，请将 sharded_data_parallel_degree 参数添加到 SageMaker PyTorch 估算器中。此参数指定对训练状态进行分片的 GPU 数量。sharded_data_parallel_degree 的值必须是介于 1 和数据并行度之间的整数，并且必须能够整除数据并行度。请注意，该库会自动检测 GPU 的数量，从而确定数据并行度。以下附加参数可用于配置分片数据并行度。

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

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 内存使用量。

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

当您想将大型语言模型 (LLM) 放入大规模集群中，同时使用序列长度更长的文本数据时，这会导致使用较小的批次大小，进而需要处理 GPU 内存使用量来对更长的文本序列训练 LLM，此时将分片数据并行性和张量并行性结合使用会非常有用。要了解更多信息，请参阅使用张量并行性的分片数据并行性。

有关案例研究、基准测试和更多配置示例，请参阅博客文章 Amazon SageMaker 模型并行库中的新性能改进。

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

SageMaker 数据并行性库提供针对 Amazon 基础设施进行了优化的集合体通信原语（SMDDP 集合体）。它使用 Elastic Fabric Adapter (EFA)，通过采用多对多类型的通信模式来实现优化，从而得到高吞吐量和对延迟较不敏感的集合体，将与通信相关的处理分载给 CPU，并腾出 GPU 周期用于计算。在大型集群上，与 NCCL 相比，SMDDP 集合体可以将分布式训练性能提高多达40%。有关案例研究和基准测试结果，请参阅博客 Amazon SageMaker 模型并行性库中的新性能改进。

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

使用 SMDDP 集合体进行训练

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

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

以下代码示例显示了如何使用分片数据并行性和 "ddp_dist_backend" 参数来设置 PyTorch 估算器，该参数默认情况下设置为 "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 集合体。

性能和内存调整

SMDDP 集合体使用额外的 GPU 内存。根据不同的模型训练使用场景，有两个环境变量可用于配置 GPU 内存使用情况。

缓冲区大小变量调整指南

环境变量的默认值应该适用于大多数使用场景。我们建议仅在训练遇到内存不足 (OOM) 错误时才调整这些变量。

以下列表讨论了一些调整技巧，可以减少 SMDDP 集合体的 GPU 内存占用，同时继续获得带来的性能提升。

有些集合体可能会回退为使用 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 位浮点数格式 (FP16) 或脑浮点数格式 (BF16)。

对于使用分片数据并行性的 FP16 训练

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

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

对于使用分片数据并行性的 BF16 训练

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

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

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

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

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

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

示例 1

假设我们要在包含 1536 个 GPU 的集群（192 个节点，每个节点 8 个 GPU）上训练模型，将分片数据并行度设置为 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 个 GPU 上。批次大小为 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 来应用。

例如，假设我们要设置张量并行度为 4 的 32 个 GPU（形成 4 个 GPU 的组），分片数据并行度为 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)。也就是说，四个 GPU 组成一个张量并行组。在这种情况下，张量并行组中第 0 个秩的 GPU 的缩减数据并行组将为 (0,4,8,12,16,20,24,28)。缩减数据并行组根据分片数据并行度 4 进行分片，从而得到两个复制组用于数据并行性。GPU (0,4,8,12) 构成一个分片组，该分片组共同存有第 0 个张量并行秩的所有参数的完整副本，GPU (16,20,24,28) 则构成另一个这样的分片组。其他张量并行秩也有类似的分片和复制组。

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

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

您还需要激活 prescaled_batch。这意味着，不是每个 GPU 读取自己的批次数据，而是每个张量并行组集体读取具有所选批次大小的一个组合批次。实际上，这并不是将数据集按照 GPU 数量（或数据并行大小 smp.dp_size()）等分成几个部分，而是将其等分成 GPU 数量再除以 tensor_parallel_degree（也称为缩减数据并行大小 smp.rdp_size()）个部分。有关预缩放批次的更多详细信息，请参阅《SageMaker Python SDK 文档》中的预缩放批次。另请参阅 SageMaker 示例 GitHub 存储库中适用于 GPT-2 的训练脚本 train_gpt_simple.py。

以下代码片段演示了基于示例 2 中前述场景创建 PyTorch 估算器对象的示例。

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 模型并行性库的分片数据并行性时，请考虑以下几点。