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

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

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

以下页面介绍了有关开始使用 Amazon SageMaker AI 进行分布式训练所需步骤的信息。如果您已熟悉如何进行分布式训练，请在以下选项中，选择与您偏好的策略或框架相符的选项以开始使用。如果您想全面了解分布式训练，请参阅分布式训练概念。

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

入门准备

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

可用区和网络背板

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

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

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

使用 SageMaker AI 分布式数据并行（SMDDP）库

SMDDP 库通过针对 Amazon 网络基础设施和 Amazon SageMaker AI ML 实例拓扑结构优化的 AllReduce 和 AllGather 集体通信操作的实现，改善了节点之间的通信。您可以使用 SMDDP 库作为基于 PyTorch 的分布式训练包的后端：PyTorch 分布式数据并行（DDP）、PyTorch 全分片数据并行性（FSDP）、DeepSpeed 和 Megatron-DeepSpeed。下面的代码示例说明了如何设置 PyTorch 估算器，以便在两个 ml.p4d.24xlarge 实例上启动分布式训练作业。

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    ...,
    instance_count=2,
    instance_type="ml.p4d.24xlarge",
    # Activate distributed training with SMDDP
    distribution={ "pytorchddp": { "enabled": True } }  # mpirun, activates SMDDP AllReduce OR AllGather
    # distribution={ "torch_distributed": { "enabled": True } }  # torchrun, activates SMDDP AllGather
    # distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }  # mpirun, activates SMDDP AllReduce OR AllGather
)

要了解如何准备训练脚本并在 SageMaker AI 上启动分布式数据并行训练作业，请参阅 利用 SageMaker AI 分布式数据并行库运行分布式训练。

使用 SageMaker AI 模型并行库（SMP）

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

要使用 SageMaker AI 的模型并行库，请配置 SageMaker AI 框架估算器的 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 AI 的模型并行库（另请参阅《SageMaker Python SDK 文档》中的分布式训练 API）。

使用开源分布式训练框架

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