介绍Amazon SageMaker HyperPod以大规模训练基础模型

建立基础模型（Foundation Models，简称FMs）需要构建、维护和优化大型集群，以在大量数据上训练数十亿到数百亿参数的模型。创建一个能够处理故障和环境变化、而不会丧失数天或数周的模型训练进展的弹性环境，是一个需要您实施集群扩缩容、主动健康监控、作业检查点和自动恢复训练能力的操作挑战。

我们很高兴地宣布，Amazon SageMaker HyperPod现已全面推出，可让您以高度弹性的训练环境，以最多达40%的加速器速度更快地训练基础模型，同时消除在操作大规模训练集群中所需的不具有差异化的繁重工作。借助SageMaker HyperPod，机器学习（ML）从业者可以连续数周到数月地训练FMs，而不会受到干扰，并且无需处理硬件故障问题。

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

“作为领先的开源生成式人工智能公司，我们的目标是最大限度地提高现代人工智能的可访问性。我们正在构建拥有数百亿参数的基础模型，这些模型要求基础设施能够自动扩展训练性能。借助SageMaker HyperPod的托管基础设施和优化库，我们可以将训练时间和成本减少50%以上。这使得我们的模型训练更加弹性和高性能，更快地构建最先进的模型。”

– Emad Mostaque，Stability AI创始人兼首席执行官。

为了使开发FMs的完整周期对硬件故障具备弹性，SageMaker HyperPod帮助您创建集群，监控集群健康状况，在运行中修复和替换故障节点，保存频繁检查点，并在不丢失进展的情况下自动恢复训练。此外，SageMaker HyperPod已预配置了Amazon SageMaker分布式训练库，包括SageMaker数据并行库（SMDDP）和SageMaker模型并行库（SMP），通过简化将训练数据和模型分成更小的块并在集群节点上并行处理它们，充分利用集群的计算和网络基础设施，提高FM训练性能。SageMaker HyperPod集成了Slurm工作负载管理器以进行集群和训练作业编排。

Slurm工作负载管理器概述

Slurm，以前称为Simple Linux Utility for Resource Management，是用于在分布式计算集群上运行作业的作业调度器。它还提供了使用NVIDIA Collective Communications Library (NCCL)或Message Passing Interface (MPI)标准运行并行作业的框架。Slurm是一种流行的开源集群资源管理系统，广泛应用于高性能计算（HPC）和生成式人工智能及FM训练工作负载。SageMaker HyperPod提供一种简单快速搭建Slurm集群的方法。

下面是用户与SageMaker HyperPod交互以及各个集群组件如何相互交互以及与其他AWS服务（例如Amazon FSx for Lustre和Amazon Simple Storage Service（Amazon S3））的高级架构图。

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

有关其他Slurm命令和配置的信息，请参阅Slurm Workload Manager文档。

自动恢复和修复功能

SageMaker HyperPod的一项新功能是具备作业的自动恢复功能。之前，在训练或微调作业运行期间，如果工作节点失败，用户需要检查作业状态，从最新的检查点重新启动作业，并在整个运行过程中持续监控作业。由于训练作业或微调作业需要运行数天、数周甚至数月，这将导致额外的管理开销，用户需要耗费时间来监控和维护作业，以防节点崩溃，还会产生昂贵的空闲加速计算实例的成本。

SageMaker HyperPod通过使用自动健康检查、节点替换和作业恢复来处理作业的可靠性。在SageMaker HyperPod中，使用SageMaker自定义的Slurm插件和SPANK框架来监控Slurm作业。当训练作业失败时，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）的最新版本和<如果Session Manager插件为AWS系统管理器的一项功能。

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

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

要设置集群生命周期脚本，即在新集群实例等事件发生时运行的脚本，您需要创建一个S3存储桶，并复制并可选地自定义默认的生命周期脚本。在本示例中，我们将所有生命周期脚本存储在“lifecycle-scripts”存储桶前缀中。

首先，您从GitHub存储库下载示例生命周期脚本。您应该根据需要自定义这些脚本以适应您的集群行为。

接下来，创建一个S3存储桶来存储自定义的生命周期脚本。

aws s3 mb s3://<your_bucket_name>

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

aws s3 sync . s3://<your_bucket_name>/lifecycle-scripts

最后，为了方便连接到集群的头节点，您应该安装或更新AWS CLI并安装AWS Session Manager CLI插件以允许在终端交互式连接以管理集群和运行训练作业。

您可以使用可用的按需资源或通过向SageMaker请求容量预留来创建SageMaker HyperPod集群。要创建容量预留，您需要在服务限额控制台上创建一个配额增加请求，以在特定的计算实例类型和容量分配上保留。

设置训练集群

完成以下步骤来创建SageMaker HyperPod集群：

1. 在SageMaker控制台上，在导航窗格中选择超级Pod集群下的集群管理。
2. 选择创建集群。
3. 提供集群名称和可选的任何标签以应用于集群资源，然后选择下一步。
4. 选择创建实例组并指定实例组名称、所需实例类型、所需实例数量以及您之前复制集群生命周期脚本的S3存储桶和前缀路径。

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

1. 首先创建控制器实例组，其中包括集群头节点。
2. 对于此实例组的AWS身份和访问管理（IAM）角色，选择创建新角色并指定您希望实例组中的集群实例可以访问的任何S3存储桶。

生成的角色默认将被授予对指定存储桶的只读访问权限。

1. 选择创建角色。
2. 在创建上述每个实例时在on-create脚本提示中输入要运行的脚本名称。在本示例中，on-create脚本称为on_create.sh。
3. 选择保存。
4. 选择创建实例组以创建您的工作实例组。
5. 提供所有请求的详细信息，包括所需实例类型和数量。

此示例使用四个ml.trn1.32xl加速实例来执行我们的训练作业。您可以使用与之前相同的IAM角色，也可以为工作实例自定义角色。类似地，您可以为这个工作实例组使用不同的on-create生命周期脚本。

1. 选择 下一步 继续。
2. 为您的群集实例选择所需的 VPC、子网和安全组。

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

请注意，如果您经常访问 S3 数据，则建议创建与私有子网路由表相关联的 VPC 终端节点，以减少任何潜在的数据传输费用。

1. 选择 下一步。
2. 查看群集详情摘要，然后选择 提交。

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

// create-cluster-slurm-default-vpc.json{   "ClusterName": "sagemaker-demo-cluster",   "InstanceGroups": [        {            "InstanceGroupName": "my-controller-group",            "InstanceType": "ml.m5.xlarge",            "InstanceCount": 1,            "lifecycleConfig": {                "SourceS3Uri": "s3://<your-s3-bucket>/<lifecycle-script-directory>/",                "OnCreate": "on_create.sh"            },            "ExecutionRole": "arn:aws:iam::111122223333:role/my-role-for-cluster",            "ThreadsPerCore": 1        },         {            "InstanceGroupName": "worker-group-1",            "InstanceType": "ml.trn1.32xlarge",            "InstanceCount": 4,            "lifecycleConfig": {                "SourceS3Uri": "s3://<your-s3-bucket>/<lifecycle-script-directory>/",                "OnCreate": "on_create.sh"            },            "ExecutionRole": "arn:aws:iam::111122223333:role/my-role-for-cluster",            "ThreadsPerCore": 1        }    ]}

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

aws sagemaker create-cluster create-cluster-slurm-default-vpc.json

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

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

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

aws sagemaker describe-cluster –cluster-name <cluster_name>

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

“ClusterArn”: “arn:aws:sagemaker:us-west-2:111122223333:cluster/<cluster_id>”

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

aws sagemaker list-cluster-nodes --cluster-name <cluster_name>

请注意响应中的 InstanceGroupName 和 InstanceId，这些将用于使用会话管理器连接到实例。

现在，您可以使用会话管理器登录到控制节点或登录节点，并运行训练作业：

aws ssm start-session —target sagemaker-cluster:<cluster_id>_<instance_group_name>-<instance_id>

接下来，我们要准备环境并下载Llama 2和RedPajama数据集。如需完整的代码和逐步操作指南，请按照AWSome Distributed Training GitHub存储库上的说明进行。

git clone https://github.com/aws-samples/awsome-distributed-training.git

按照2.test_cases/8.neuronx-nemo-megatron/README.md文件中的步骤处理。在准备环境、准备模型、下载和标记化数据集以及预编译模型的步骤之后，您应编辑6.pretrain-model.sh脚本和sbatch作业提交命令，以包含一个参数，以便您可以利用SageMaker HyperPod的自动恢复功能。

将sbatch行编辑为以下方式：

sbatch --nodes 4 --auto-resume=1 run.slurm ./llama2_7b.sh

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

squeue <jobid>

此外，您可以使用以下代码来监控作业的输出日志：

tail -f slurm-run.slurm-<jobid>.out

清理

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

aws sagemaker delete-cluster --cluster-name <cluster_name>

结论

本文向您展示了如何准备AWS环境、部署第一个SageMaker HyperPod集群，并训练一个70亿参数的Llama 2模型。SageMaker HyperPod现已在美国（弗吉尼亚州、俄亥俄州和俄勒冈州）、亚太地区（新加坡、悉尼和东京）以及欧洲地区（法兰克福、爱尔兰和斯德哥尔摩）提供。它们可以通过SageMaker控制台、AWS CLI和AWS SDK进行部署，并支持p4d、p4de、p5、trn1、inf2、g5、c5、c5n、m5和t3实例类型。

要了解更多关于SageMaker HyperPod的信息，请访问Amazon SageMaker HyperPod。