开始使用 Amazon A SageMaker I 进行分布式训练

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

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

入门准备

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

可用区和网络背板

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

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

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

使用 SageMaker AI 分布式数据并行度 (SMDDP) 库

SMDDP 库通过实施针对 AWS 网络基础设施AllReduce和 Amazon A SageMaker I ML 实例拓扑进行了优化的AllGather集体通信操作来改善节点之间的通信。您可以使用 SMDDP 库作为 PyTorch基于分布式训练包的后端：PyTorch 分布式数据并行 (DDP)、PyTorch 完全分片数据并行度 (FSDP) 和威震天-。DeepSpeedDeepSpeed下面的代码示例说明了如何设置 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参数。支持的框架估计器是和。PyTorchTensorFlow以下代码示例显示了如何在两个 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 的模型并行度库（另请参阅 Pyth SageMaker on SDK 文档 APIs中的分布式训练）。

使用开源分布式训练框架

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

• 要在 SageMaker AI 中将 PyTorch DistributedDataParallel (DDP) 与mpirun后端一起使用，请distribution={"pytorchddp": {"enabled": True}}添加到您的 PyTorch估算器中。有关更多信息，另请参阅 Pyth SageMaker on SDK 文档中的PyTorch 分布式训练和 SageMaker AI PyTorch 估算器的distribution论点。

  注意

  此选项适用于 PyTorch 1.12.0 及更高版本。

  from sagemaker.pytorch import PyTorch
  
  estimator = PyTorch(
      ...,
      instance_count=2,
      instance_type="ml.p4d.24xlarge",
      distribution={"pytorchddp": {"enabled": True}}  # runs mpirun in the backend
  )

• SageMaker AI 支持PyTorch torchrun启动器在基于 GPU 的亚马逊 EC2 实例（例如 P3 和 P4）以及由 Trainium 设备提供支持的 Trn1 上进行分布式训练。AWS

  要在 SageMaker AI 中将 PyTorch DistributedDataParallel (DDP) 与torchrun后端一起使用，请distribution={"torch_distributed": {"enabled": True}}添加到 PyTorch 估算器中。

  注意

  此选项适用于 PyTorch 1.13.0 及更高版本。

  以下代码片段显示了构建 A SageMaker I PyTorch 估计器以使用分布选项在两个ml.p4d.24xlarge实例上运行分布式训练的torch_distributed示例。

  from sagemaker.pytorch import PyTorch
  
  estimator = PyTorch(
      ...,
      instance_count=2,
      instance_type="ml.p4d.24xlarge",
      distribution={"torch_distributed": {"enabled": True}}   # runs torchrun in the backend
  )

  有关更多信息，请参阅 Pyth SageMaker on SDK 文档中的分布式 PyTorch 训练和 SageMaker AI PyTorch 估算器的distribution论点。

  在 Trn1 上进行分布式训练的注意事项

  一个 Trn1 实例由最多 16 个 Trainium 设备组成，每个 Trainium 设备由两台设备组成。NeuronCores有关 AWS Trainium 设备的规格，请参阅 Neuron 文档中的 Trainium 架构。AWS

  要在 Trainium 支持的实例上进行训练，您只需要在 AI 估算器类的instance_type参数中以字符串形式指定 Trn1 实例代码。ml.trn1.* SageMaker PyTorch 要查找可用的 Trn1 实例类型，请参阅《AWS Neuron 文档》中的 AWS Trn1 架构。

  注意

  SageMaker Amazon EC2 Trn1 实例上的训练目前仅适用于从 v1.11.0 开始的 Dee AWS p Learning Containers for Ne PyTorch uron 中的 PyTorch框架。要查找支持的 Ne PyTorch uron 版本的完整列表，请参阅 Dee AWS p Learning Conta iners 存储库中的 Neuron 容器 GitHub 。

  当你使用 Pyth SageMaker on SDK 在 Trn1 实例上启动训练作业时， SageMaker AI 会自动从 Deep Learning Containers 提供的AWS 神经元容器中提取并运行正确的容器。Neuron Containers 预先打包了训练环境设置和依赖项，便于您的训练作业适应 SageMaker 训练平台和 Amazon EC2 Trn1 实例。

  注意

  要使用 SageMaker AI 在 Trn1 实例上运行 PyTorch 训练作业，您应修改训练脚本以使用xla后端初始化进程组并使用 PyTorch /XLA。为了支持 XLA 的采用过程， AWS Neuron SDK 提供了 Ne PyTorch uron，它使用 XLA 将 PyTorch 操作转换为 Trainium 指令。要了解如何修改训练脚本，请参阅 Neuron 文档中的 PyTorch Neuron 训练开发者指南 (torch-neuronx)。AWS

  有关更多信息，请参阅 Python SDK 文档中的 Trn1 实例 PyTorch 上使用神经元进行分布式训练和 SageMaker AI E PyTorch stimator distribution 的论点。SageMaker

• 要在 SageMaker AI 中使用 MPI，请distribution={"mpi": {"enabled": True}}添加到您的估算器中。MPI 分发选项适用于以下框架： MXNet PyTorch、和。 TensorFlow

• 要在 SageMaker AI 中使用参数服务器，请distribution={"parameter_server": {"enabled": True}}添加到您的估算器中。参数服务器选项适用于以下框架： MXNet PyTorch、和 TensorFlow。

  提示

  有关使用每个框架的 MPI 和参数服务器选项的更多信息，请使用以下 SageMaker Python SDK 文档链接。

  • MXNet 分布式训练和 SageMaker AI MXNet 估算器的论点 distribution

  • PyTorch 分布式训练和 SageMaker AI PyTorch 估算器的论点 distribution

  • TensorFlow 分布式训练和SageMaker 人工智能 TensorFlow 估算器的论点。distribution