Amazon SageMaker HyperPod：大规模训练基础模型的解决方案

关键要点

Amazon SageMaker HyperPod 现已全面上线，可实现高达40更快的基础模型训练。SageMaker HyperPod 提供坚如磐石的训练环境，减少了大规模训练集群的管理复杂性。引入自动恢复和自我修复能力，从而增强训练作业的韧性。支持与 AWS 中的其他服务集成，以提供优化的性能和便利。

构建基础模型FMs需要建立、管理和优化大型集群，以在海量数据上训练数十亿至数百亿参数的模型。创建一个能够应对故障和环境变化的稳健环境，是一个运营挑战，这需要实施集群扩展、主动健康监控、作业检查点以及在发生故障或问题时自动恢复训练的能力。

我们很高兴地宣布，Amazon SageMaker HyperPod现已正式推出，它能够以最多40更快的速度，利用数千个加速器训练基础模型。SageMaker HyperPod 提供一个高弹性的训练环境，消除了运行大规模训练集群所需的繁重管理工作。通过 SageMaker HyperPod，机器学习ML从业者能够在数周甚至数月内无干扰地训练基础模型，而无需担心硬件故障的问题。

像 Stability AI 这样的客户正在使用 SageMaker HyperPod 来训练他们的基础模型，包括 Stable Diffusion。

“作为领先的开源生成性 AI 公司，我们的目标是最大限度地提高现代 AI 的可及性。我们正在构建数十亿参数的基础模型，这需要优化训练性能的基础设施。借助 SageMaker HyperPod 的管理基础设施和优化库，我们可以将训练时间和成本减少超过50。这使得我们的模型训练更具韧性和性能，从而更快地构建尖端模型。”

Emad Mostaque，Stability AI 创始人兼 CEO。

为了使基础模型开发的全周期在硬件故障面前具有韧性，SageMaker HyperPod 帮助您创建集群、监控集群健康、即时修复和更换故障节点、保存频繁的检查点，并在不丢失进展的情况下自动恢复训练。此外，SageMaker HyperPod 预配置了 Amazon SageMaker 分布式训练库，包括 SageMaker 数据并行库SMDDP 和 SageMaker 模型并行库SMP，通过将训练数据和模型拆分成更小的块并在集群节点之间并行处理，从而提升模型训练性能，同时充分利用集群的计算和网络基础设施。SageMaker HyperPod 集成了 Slurm 工作负载管理器，以便进行集群和训练作业协调。

Slurm 工作负载管理器概述

Slurm以前称为 Simple Linux Utility for Resource Management是一个用于在分布式计算集群上运行作业的作业调度器。它还提供了一个框架，用于使用 NVIDIA 集体通信库NCCL 或 消息传递接口MPI 标准运行并行作业。Slurm 是一个广泛使用的开源集群资源管理系统，主要应用于高性能计算HPC以及生成性 AI 和基础模型训练工作负载。SageMaker HyperPod 提供了一种简单的方法，可以在几分钟内启动和运行 Slurm 集群。

下图展示了用户如何与 SageMaker HyperPod 互动，以及各种集群组件如何相互联系，并与其他 AWS 服务如 Amazon FSx for Lustre 和 Amazon Simple Storage Service进行交互。

Slurm 作业是通过命令行提交的。运行 Slurm 作业的命令是 srun 和 sbatch。srun 命令以交互和阻塞模式运行训练作业，而 sbatch 则以批处理和非阻塞模式运行。srun 通常用于运行即时的作业，而 sbatch 则可用于后续的作业运行。

有关更多 Slurm 命令和配置的信息，请参见 Slurm 工作负载管理器文档。

自动恢复和自我修复能力

SageMaker HyperPod 的一个新特性是作业的自动恢复能力。以前，当工作节点在训练或微调作业运行中出现故障时，用户需要检查作业状态、从最新检查点重新启动作业，并在整个运行过程中持续监控该作业。由于训练作业或微调作业可能需要连续运行数天、数周甚至数月，这会导致用户需花费额外的时间监控和维护作业，尤其是在节点崩溃的情况下，造成高昂的管理开销，以及昂贵的加速计算实例闲置时间成本。

SageMaker HyperPod 通过使用自动健康检查、节点替换和作业恢复，解决了作业的韧性问题。SageMaker HyperPod 中的 Slurm 作业通过使用 SageMaker 自定义 Slurm 插件基于 SPANK 框架进行监控。当训练作业失败时，SageMaker HyperPod 会通过一系列健康检查检查集群健康。如果发现集群中存在故障节点，SageMaker HyperPod 将自动将该节点从集群中移除，并用健康节点替换它，然后重新启动训练作业。当在训练作业中使用检查点时，任何中断或失败的作业都可以从最新的检查点恢复。

解决方案概述

要部署您的 SageMaker HyperPod，您首先需要通过配置 Amazon 虚拟私有云Amazon VPC网络和安全组，部署在 VPC 中支持的服务如 FSx for Lustre，并将 Slurm 生命周期脚本发布到 S3 存储桶。然后，您可以部署和配置 SageMaker HyperPod，并连接到头节点以开始训练作业。

先决条件

在创建 SageMaker HyperPod 之前，您需要首先配置 VPC，创建 FSx for Lustre 文件系统，并建立一个具有所需集群生命周期脚本的 S3 存储桶。您还需要最新版本的 AWS 命令行界面AWS CLI和为 AWS 会话管理器 安装的 CLI 插件，这是一项 AWS 系统管理器 的功能。

SageMaker HyperPod 与 VPC 完全集成。有关创建新 VPC 的信息，请参阅 创建默认 VPC 或 创建 VPC。为了确保资源间最高性能的无缝连接，您应该在同一区域和可用区中创建所有资源，并确保关联的安全组规则允许集群资源之间的连接。

接下来，您可以 创建 FSx for Lustre 文件系统。这将作为我们模型训练期间使用的高性能文件系统。确保 FSx for Lustre 和集群安全组允许集群资源与 FSx for Lustre 文件系统之间的入站和出站通信。

要设置集群生命周期脚本，需要创建一个 S3 存储桶，随后复制并自定义默认的生命周期脚本。在此示例中，我们将所有生命周期脚本存储在一个存储桶前缀 lifecyclescripts 下。

首先，您需要从 GitHub 仓库 下载示例生命周期脚本。您应根据所需的集群行为进行自定义。

然后，创建一个 S3 存储桶以存储自定义的生命周期脚本。

bashaws s3 mb s3//ltyourbucketnamegt

接下来，使用 aws s3 sync 将默认生命周期脚本从本地目录复制到您选定的存储桶和前缀：

bashaws s3 sync s3//ltyourbucketnamegt/lifecyclescripts

最后，为了简化与集群头节点的连接，您应该 安装或更新 AWS CLI，并安装 AWS 会话管理器 CLI 插件，以支持交互式终端连接，用于管理集群并运行训练作业。

您可以选择使用可用的按需资源或通过 SageMaker 请求容量预留来创建 SageMaker HyperPod 集群。要创建容量预留，您可以在服务配额仪表板中提交配额增加请求，以预留特定的计算实例类型和容量分配。

设置训练集群

要创建 SageMaker HyperPod 集群，请完成以下步骤：

在 SageMaker 控制台中，选择导航栏中的 HyperPod 集群 下的 集群管理。

选择 创建集群。

提供集群名称，并 optionally 为集群资源应用任何标签，然后选择 下一步。

选择 创建实例组，并指定实例组名称、所需实例类型、所需实例数量，以及您之前复制的集群生命周期脚本存储的 S3 存储桶和前缀路径。

建议为用于管理集群和提交作业的控制节点和用于运行训练作业的工作节点配置不同的实例组。您还可以选择为登录节点配置额外的实例组。

首先创建控制器实例组，其中将包括集群头节点。对于此实例组的 AWS 身份与访问管理IAM角色，选择 创建新角色 并指定任何想要让实例组中的集群实例访问的 S3 存储桶。

所生成的角色默认获得对指定存储桶的只读访问权限。

选择 创建角色。

在每个实例创建时要运行的脚本的创建提示中输入脚本名称。在此示例中，oncreate脚本称为 oncreatesh。

选择 保存。

选择 创建实例组 以创建您的工作实例组。提供所需的所有信息，包括实例类型和所需数量。

本示例中使用了四个 mltrn132xl 加速实例来执行我们的训练任务。您可以使用与之前相同的 IAM 角色，或者为工作实例自定义角色。类似地，您可以为此工作实例组使用不同的 oncreate 生命周期脚本。

选择 下一步 以继续。选择集群实例的所需 VPC、子网和安全组。

我们将集群实例托管在单个可用区和子网中，以确保低延迟。

请注意，如果您将频繁访问 S3 数据，建议创建与私有子网路由表关联的 VPC 终端节点，以降低潜在的数据传输成本。

选择 下一步。

审查集群详细信息摘要，然后选择 提交。

或者，要使用 AWS CLI 创建 SageMaker HyperPod，首先自定义用于创建集群的 JSON 参数：

json// createclusterslurmdefaultvpcjson{ ClusterName sagemakerdemocluster InstanceGroups [ { InstanceGroupName mycontrollergroup InstanceType mlm5xlarge InstanceCount 1 lifecycleConfig { SourceS3Uri s3//ltyours3bucketgt/ltlifecyclescriptdirectorygt/ OnCreate oncreatesh } ExecutionRole arnawsiam111122223333role/myroleforcluster ThreadsPerCore 1 } { InstanceGroupName workergroup1 InstanceType mltrn132xlarge InstanceCount 4 lifecycleConfig { SourceS3Uri s3//ltyours3bucketgt/ltlifecyclescriptdirectorygt/ OnCreate oncreatesh } ExecutionRole arnawsiam111122223333role/myroleforcluster ThreadsPerCore 1 } ]}

然后，使用以下命令使用提供的输入创建集群：

bashaws sagemaker createcluster createclusterslurmdefaultvpcjson

使用 Llama 2 运行第一个训练作业

请注意，Llama 2 模型的使用受到 Meta 许可证的管理。要下载模型权重和分词器，请访问 Meta 官网 并接受许可证，然后在 Meta 的 Hugging Face 网站 上申请访问权限。

集群运行后，使用集群 ID、实例组名称和实例 ID 使用会话管理器登录。使用以下命令查看您的集群详细信息：

bashaws sagemaker describecluster clustername ltclusternamegt

注意响应中的集群 ARN 中包含的集群 ID。

json“ClusterArn” “arnawssagemakeruswest2111122223333cluster/ltclusteridgt”

使用以下命令检索所需的实例组名称和实例 ID，以便登录集群。

bashaws sagemaker listclusternodes clustername ltclusternamegt

请记录响应中的 InstanceGroupName 和 InstanceId，因为它们将用于与会话管理器连接到实例。

现在，您可以使用会话管理器登录到头节点或其中一个登录节点，并运行训练作业：

bashaws ssm startsession target sagemakerclusterltclusteridgtltinstancegroupnamegtltinstanceidgt

接下来，我们要准备环境并下载 Llama 2 和 RedPajama 数据集。有关完整的代码和逐步说明，请参见 AWSome 分布式训练 GitHub 仓库。

bashgit clone https//githubcom/awssamples/awsomedistributedtraininggit

按照 2testcases/8neuronxnemomegatron/READMEmd 文件中详细说明的步骤进行操作。按照步骤准备环境、准备模型、下载和分词数据集，以及预编译模型后，您应该编辑 6pretrainmodelsh 脚本和 sbatch 作业提交命令，以包括一个允许您利用 SageMaker HyperPod 自动恢复特性的参数。

将 sbatch 行编辑如下：

试用机场节点

bashsbatch nodes 4 autoresume=1 runslurm /llama27bsh

提交作业后，您将获得一个 JobID，可以使用以下代码检查作业状态：

bashsqueue ltjobidgt

此外，您可以通过以下代码跟踪作业输出日志来监控作业：

bashtail f slurmrunslurmltjobidgtout

清理

要删除您的 SageMaker HyperPod 集群，可以使用 SageMaker 控制台或以下 AWS CLI 命令：

bashaws sagemaker deletecluster clustername ltclusternamegt