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

使用 Amazon SageMaker 进行分布式训练

SageMaker 提供分布式训练库，并支持适用于计算机视觉 (CV) 和自然语言处理 (NLP) 等深度学习任务的各种分布式训练选项。使用 SageMaker 的分布式训练库，您可以采用数据并行和模型并行方法，经济高效地运行高度可扩展的自定义深度学习训练作业。您也可以使用其他分布式训练框架和程序包，例如 PyTorch DistributedDataParallel (DDP)、torchrun、MPI (mpirun) 和参数服务器。本文中的说明和示例重点介绍如何使用 SageMaker Python SDK，针对深度学习设置分布式训练选项。

入门准备

SageMaker Training 支持在单个实例和多个实例上进行分布式训练，因此您可以大规模运行任何大小的训练。我们建议您在 SageMaker Python SDK 中使用框架估算器类，例如 PyTorch 和 TensorFlow，它们是训练作业启动程序，提供了各种分布式训练选项。创建估算器对象时，对象会设置分布式训练基础设施，在后端运行 CreateTrainingJob API，找到当前会话运行所在的区域，然后提取一个预先构建的 Amazon 深度学习容器，该容器预先打包了大量库，包括深度学习框架、分布式训练框架和 EFA 驱动程序。要将 FSx 文件系统挂载到训练实例，您需要将 VPC 子网和安全组 ID 传递给估算器。在 SageMaker 中运行分布式训练作业之前，请您先阅读以下有关基础设施设置的一般指导。

可用区和网络背板

使用多个实例（也称为节点）时，务必要了解连接实例的网络、它们如何读取训练数据以及如何在彼此之间共享信息。例如，当您运行分布式数据并行训练作业时，有许多因素，例如用于运行 AllReduce 操作的计算集群节点之间的通信，以及节点与 Amazon Simple Storage Service 或适用于 Lustre 的 Amazon FSx 中的数据存储之间的数据传输，对于实现计算资源的最佳使用和更快的训练速度起着至关重要的作用。为了减少通信开销，请确保在同一 Amazon Web Services 区域和可用区中配置实例、VPC 子网和数据存储。

具有更快网络速度和高吞吐量存储的 GPU 实例

从技术上讲，您可以使用任何实例进行分布式训练。在您需要运行多节点分布式训练作业来训练大型模型时，例如大型语言模型 (LLM) 和扩散模型，这些作业需要更快的节点间通信速度，我们建议使用 SageMaker 支持的启用 EFA 的 GPU 实例。特别是，为了在 SageMaker 中实现最高性能的分布式训练作业，我们建议采用配备 NVIDIA A100 GPU 的 P4d 和 P4de 实例。这些实例还配备了高吞吐量、低延迟的本地实例存储和更快的节点内网络。对于数据存储，我们推荐使用适用于 Lustre 的 Amazon FSx，它可为存储训练数据集和模型检查点提供高吞吐量。

开始在 Amazon SageMaker 中使用分布式训练

如果您已熟悉如何进行分布式训练，请在以下选项中，选择与您偏好的策略或框架相符的选项以开始使用。如果您想全面了解分布式训练，请参阅分布式训练基本概念。

SageMaker 分布式训练库针对 SageMaker 训练环境进行了优化，可以帮助您调整分布式训练作业以使其适于 SageMaker，并提高训练速度和吞吐量。该库提供了数据并行和模型并行训练策略。它们结合了软件和硬件技术来改进 GPU 间和节点间的通信，并可通过内置选项，只需对训练脚本进行最少的代码更改即可扩展 SageMaker 的训练功能。 

使用 SageMaker 数据并行性库 (SMDDP)

要使用 SageMaker 的数据并行性库，请配置 SageMaker 框架估算器的 distribution 参数。支持的框架估算器包括 PyTorch 和 TensorFlow。以下代码示例显示了如何在两个 ml.p4d.24xlarge 实例上，使用数据并行性库为分布式训练设置框架估算器。

from sagemaker.framework import Framework

estimator = Framework(
    ...,
    instance_count=2,
    instance_type="ml.p4d.24xlarge",
    distribution={"smdistributed" : {"dataparallel" : {"enabled" : True}}}
)

要了解如何准备训练脚本和启动分布式数据并行训练作业，请参阅 SageMaker 的数据并行性库（另请参阅《SageMaker Python SDK 文档》中的分布式训练 API）。

使用 SageMaker 模型并行性库 (SMP)

SageMaker 提供 SMP 库并支持各种分布式训练技术，例如分片数据并行性、管道处理、张量并行性、优化器状态分片等。要了解有关 SMP 库所提供功能的更多信息，请参阅 SageMaker 模型并行性库的核心功能。

要使用 SageMaker 的模型并行性库，请配置 SageMaker 框架估算器的 distribution 参数。支持的框架估算器包括 PyTorch 和 TensorFlow。以下代码示例显示了如何在两个 ml.p4d.24xlarge 实例上，使用数据并行性库为分布式训练构造框架估算器。

from sagemaker.framework import Framework

distribution={
    "smdistributed": {
        "modelparallel": {
            "enabled":True,
            "parameters": {
                ...   # enter parameter key-value pairs here
            }
        },
    },
    "mpi": {
        "enabled" : True,
        ...           # enter parameter key-value pairs here
    }
}

estimator = Framework(
    ...,
    instance_count=2,
    instance_type="ml.p4d.24xlarge",
    distribution=distribution
)

要了解如何调整训练脚本、在 estimator 类中配置分布参数以及启动分布式训练作业，请参阅 SageMaker 的模型并行性库（另请参阅《SageMaker Python SDK 文档》中的分布式训练 API）。

使用开源分布式训练框架

SageMaker 还支持以下选项，用于在后端操作 mpirun 和 torchrun。

分布式训练基本概念

SageMaker 的分布式训练库使用以下分布式训练术语和功能。

数据集和批处理

训练

实例和 GPU

分布式训练解决方案

高级概念

机器学习 (ML) 从业者在训练模型时通常会面临两个扩展挑战：扩展模型大小和扩展训练数据。虽然模型的大小和复杂程度会影响到准确性，但是可以放入单个 CPU 或 GPU 的模型大小有限制。此外，扩展模型大小可能会导致更多的计算量和更长的训练时间。

并非所有模型都能很好地处理训练数据的扩展，因为它们需要在内存中摄取所有训练数据用于训练。这些模型只能垂直扩展，不断扩展到更大的实例类型。在大多数情况下，扩展训练数据会导致更长的训练时间。

深度学习 (DL) 是一个特定的 ML 算法系列，由多层人工神经网络组成。最常见的训练方法是使用小批次随机梯度下降 (SGD)。在小批次 SGD 中，模型的训练方法是在减少误差的方向上，对其系数进行小的迭代变化。这些迭代在训练数据集的相同大小子样本上进行，这些子样本称为小批次。对于每个小批次，模型都对小批次的每个记录运行，衡量其误差和估算误差的梯度。然后，对小批次的所有记录测量平均梯度，并为每个模型系数提供更新方向。训练数据集的一次完整过程称为一个纪元。模型训练通常包括几十到数百个纪元。小批次 SGD 有几个好处：首先，它的迭代设计使训练时间理论上与数据集大小呈线性关系。其次，在给定的小批次中，模型会单独处理每条记录，除了最后的取梯度平均值外，不需要记录间的通信。因此，小批次的处理特别适合并行化和分布式。 

通过将小批次记录分布到不同的计算设备进行的并行化 SGD 训练被称为数据并行分布式训练，是最常用的 DL 分布模式。数据并行训练是一种相关的分发策略，可以扩展小批次大小，并更快地处理每个小批次。但是，数据并行训练带来了额外的复杂性，即必须使用来自所有工作线程的梯度，计算小批次梯度平均值，然后将其传输给所有工作线程，这是一个名为 allreduce 的步骤，随着训练集群的扩展，这可能带来更大的开销，如果实施方法不当或实施了不当的硬件减少，也可能严重地影响到训练时间。 

数据并行 SGD 仍然要求开发人员至少能够至少将模型和单个记录放入一个计算设备（例如单个 CPU 或 GPU）中。在训练非常大的模型时，例如自然语言处理 (NLP) 中的大型转换器，或者对高分辨率图像分段模型时，可能会出现这种做法不可行的情况。拆分工作负载的另一种替代方法是将模型在多个计算设备上分区，这种方法称为模型并行分布式训练。

策略

分布式训练通常按两种方法来拆分：数据并行和模型并行。数据并行是分布式训练的最常见方法：您有很多数据，将其分成批次，然后将数据块发送到多个 CPU 或 GPU（节点）以供神经网络或 ML 算法处理，然后合并结果。每个节点上的神经网络都是一样的。模型并行方法用于不能整体放到一个节点的内存中的大型模型；该方法拆分模型，不同的部分放在不同的节点上。在这种情况下，您需要将数据批次发送到各个节点，以便在模型的所有部分上处理数据。

术语网络和模型经常可以互换使用：大型模型实际上是一个具有许多层和参数的大型网络。使用大型网络进行训练会生成一个大型模型，将模型以及所有预训练的参数及其权重加载回网络，会将大型模型加载到内存中。当您拆分模型以在节点之间分布模型时，您也会拆分底层网络。网络由层组成，为了拆分网络，您可以将层放置在不同的计算设备上。

在设备之间拆分层时，一个常见的内在缺陷是 GPU 的利用率严重不足。无论是向前还是向后传递，训练在本质上都是顺序的，在给定时间，只有一个 GPU 可以有效地进行计算，而其他 GPU 则等待发送激活信号。为了解决这个问题，现代模型并行库使用管道执行计划来提高设备利用率。但是，只有 Amazon SageMaker 的分布式模型并行库包含自动模型拆分。该库有两个核心功能：自动模型拆分和管道执行计划，通过制定自动化决策以实现高效设备利用，简化了实施模型并行的过程。

使用数据并行和模型并行进行训练

如果您在使用大型数据集进行训练，请首先使用数据并行方法。如果您在训练期间内存不足，则可能需要切换到模型并行方法，或尝试混合使用模型并行和数据并行。您还可以尝试以下方法，通过数据并行来提高性能：

尝试梯度压缩（FP16、INT8）：

尝试减小输入大小：

检查是否使用了批次标准化，因为这可能会影响收敛性。在使用分布式训练时，您的批次在 GPU 之间拆分，而小得多的批次大小可能会带来更高的错误率，进而干扰模型的收敛。例如，如果您在批次大小为 64 的单个 GPU 上对网络进行原型设计，然后扩展到使用 4 个 p3dn.24xlarge，那么现在有 32 个 GPU，而每个 GPU 的批次大小从 64 降至 2。这可能会破坏您在单个节点上可以实现的收敛。

在以下情况下从模型并行训练开始：

要了解有关 SageMaker 分布式库的更多信息，请参阅以下内容：

优化分布式训练

根据您的使用场景和数据自定义超参数，以获得最佳的扩展效率。在下面的讨论中，我们重点介绍了一些最具影响力的训练变量，并提供了对先进实施方案的参考，以便您了解更多的可用选项。此外，建议您参考您偏好的框架的分布式训练文档。

批次大小

SageMaker 分布式工具包通常允许您在较大的批次上进行训练。例如，如果某个模型能够放入单个设备，但只能使用小批次大小进行训练，则模型并行训练或数据并行训练可以让您实验较大的批次大小。

请注意，批次大小会控制每次迭代时模型更新中的噪声量，进而会直接影响模型准确性。增加批次大小可以减少梯度估算中的噪声量，这在从非常小的批次大小开始增加时会有所帮助，但随着批量大小增加到较大值，可能会导致模型准确性降低。 

现在已经开发了许多技术，以便在批次增加时保持良好的模型收敛性。

小批次大小

在 SGD 中，小批次大小量化了梯度估算中存在的噪声量。较小的小批次会产生噪声非常多的小批次梯度，这并不能代表数据集的真正梯度。较大的小批次会得到接近数据集上真实梯度的微批次梯度，并且有可能噪声不够，因而可能会保持锁定在无关的最小点中。

要了解有关这些技术的更多信息，请参阅以下文章：

场景

以下部分介绍了您可能需要纵向扩展训练的场景，以及如何使用 Amazon 资源完成扩展。

从单个 GPU 扩展到多个 GPU

根据机器学习中使用的数据量或模型的大小，可能会出现训练模型的时间超过您愿意等待的时间的情况。有时，由于模型或训练数据太大，训练可能完全无法正常工作。这种情况的解决方案之一是增加用于训练的 GPU 数量。在具有多个 GPU 的实例上，例如具有八个 GPU 的 p3.16xlarge，数据和处理会在八个 GPU 之间拆分。当您使用分布式训练库时，这可以为您训练模型所用的时间带来近乎线性的加速。相比在只有一个 GPU 的 p3.2xlarge 上，此时所需的时间仅比 1/8 略多一点。

从单个实例扩展到多个实例

如果您想进一步扩展训练，则可以使用更多的实例。但是，您应先考虑选择更大的实例类型，然后再考虑添加更多实例。查看上表以了解每种 p3 实例类型中有多少个 GPU。

当您已经从 p3.2xlarge 上的 1 个 GPU 提升到 p3.8xlarge 上的 4 个 GPU，但是还需要更多的处理能力时，在增加实例数量之前，如果您选择 p3.16xlarge，则可能会发现可以获得更好的性能并且产生的成本更低。根据您使用的库，保持在单个实例上进行训练时，相比使用多个实例的场景性能会更好，成本也更低。

当您准备好扩展实例数量时，您可以使用 SageMaker Python SDK estimator 功能，通过设置 instance_count 来扩展。例如，您可以创建 instance_type = p3.16xlarge 和 instance_count = 2。您现在不是使用单个 p3.16xlarge 上的 8 个 GPU，而是使用 2 个相同实例上的 16 个 GPU。下图显示了扩展和吞吐量的变化，从单个实例上的 8 个 GPU 开始，增加到 64 个实例，总共 256 个 GPU。

自定义训练脚本

虽然 SageMaker 可以根据您选择的框架，轻松部署和扩展实例及 GPU 的数量，但管理数据和结果可能非常具有挑战性，这就是经常会使用外部支持库的原因。这种最基本的分布式训练形式需要修改训练脚本以管理数据分布。

SageMaker 还支持 Horovod 以及各个主要深度学习框架原生的分布式训练实施。如果您选择使用这些框架中的示例，请关注 SageMaker 的容器指南有关深度学习容器的内容，以及各种演示实施的示例笔记本。