推出适合大规模分布式训练的专用基础设施 Amazon SageMaker HyperPod

今天，我们将介绍 Amazon SageMaker HyperPod，它通过为大规模分布式训练提供专门构建的基础设施，帮助缩短基础模型（FM）的训练时间。现在，您可以使用 SageMaker HyperPod 对 FM 进行数周甚至数月的训练，同时 SageMaker 会主动监控集群运行状况，并通过替换故障节点和从检查点恢复模型训练来提供自动节点和作业弹性。

集群预配置了 SageMaker 的分布式训练库，可帮助您在所有节点中拆分训练数据和模型，从而进行并行处理并充分利用集群的计算和网络基础设施。您可以通过安装其他框架、调试工具和优化库来进一步自定义训练环境。

下面，我将介绍如何开始使用 SageMaker HyperPod。在下面的演示中，我创建了一个 SageMaker HyperPod，并使用在 AWS 机器学习训练参考架构 GitHub 存储库中分享的示例展示如何训练 Llama 2 7B 模型。

创建和管理集群
作为 SageMaker HyperPod 管理员，您可以使用 AWS 管理控制台或 AWS 命令行界面（AWS CLI）创建和管理集群。在控制台中，导航到 Amazon SageMaker，在左侧菜单的 HyperPod 集群下选择集群管理，然后选择创建集群。

在随后的设置中，提供集群名称并使用您选择的实例类型和要分配给每个实例组的实例数量配置实例组。

您还需要准备一个或多个生命周期脚本并将其上传到您的 Amazon Simple Storage Service（Amazon S3）桶，以便在集群创建期间在每个实例组中运行此脚本。使用生命周期脚本，您可以自定义集群环境并安装所需的库和软件包。您可以在 GitHub 存储库中找到 SageMaker HyperPod 的示例生命周期脚本。

使用 AWS CLI
您还可以使用 AWS CLI 来创建和管理集群。在演示中，我在 JSON 文件中指定集群配置。我选择创建两个实例组，一个用于名为“controller-group”的集群控制器节点，另一个用于名为“worker-group”的集群 Worker 节点。对于要执行模型训练的 Worker 节点，我指定由 AWS Trainium 芯片提供支持的 Amazon EC2 Trn1 实例。

// demo-cluster.json
{
   "InstanceGroups": [
        {
            "InstanceGroupName": "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",
            "InstanceType": "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 CLI 命令：

aws sagemaker create-cluster \
--cluster-name antje-demo-cluster \
--instance-groups file://demo-cluster.json

创建后，可以使用 aws sagemaker describe-cluster 和 aws sagemaker list-cluster-nodes 来查看集群和节点详细信息。记下控制器节点的集群 ID 和实例 ID。您需要这些信息才能连接到集群。

您还可以选择连接共享文件系统，例如 Amazon FSx for Lustre。要使用 FSx for Lustre，您需要利用 Amazon Virtual Private Cloud（Amazon VPC）配置来设置集群。这是一个 AWS CloudFormation 模板，展示了如何创建 SageMaker VPC以及如何部署 FSx for Lustre。

连接您的集群
作为集群用户，您需要对集群管理员配置的集群的访问权限。获得访问权限后，您可以使用 SSH 连接到集群，并计划和运行作业。您可以使用预安装的适用于 AWS Systems Manager 的 AWS CLI 插件连接到集群的控制器节点。

在演示中，我运行了以下命令，将控制器节点的集群 ID 和实例 ID 指定为目标。

aws ssm start-session \
--target sagemaker-cluster:ntg44z9os8pn_i-05a854e0d4358b59c \
--region us-west-2

使用 Slurm 在集群上计划和运行作业
在发布之初，SageMaker HyperPod 便支持使用 Slurm 进行工作负载编排。Slurm 是热门开源集群管理和作业调度系统。在集群创建过程中，您可以通过生命周期脚本安装和设置 Slurm。示例生命周期脚本展示了操作方法。然后，您可以使用标准 Slurm 命令来计划和启动作业。查看 Slurm 快速入门用户指南，了解架构详细信息和有用命令。

在本演示中，我使用的是 AWS 机器学习训练参考架构 GitHub 存储库中的这个示例，该示例展示了如何使用 Trn1 实例在 Slurm 上训练 Llama 2 7B。我的集群已经使用 Slurm 设置好了，并且已经挂载了 FSx for Lustre 文件系统。

注意
Llama 2 模型由 Meta 控制。您可以通过 Meta 请求访问页面请求访问权限。

设置集群环境
SageMaker HyperPod 支持在各种环境中进行训练，包括 Conda、venv、Docker 和 enroot。按照自述文件中的说明，我构建了虚拟环境 aws_neuron_venv_pytorch，并设置了 torch_neuronx 和 neuronx-nemo-megatron 库，用于在 Trn1 实例上训练模型。

准备模型、分词器和数据集
我按照说明下载了 Llama 2 模型和分词器，然后将模型转换为 Hugging Face 格式。接着，我下载并标记了 RedPajama 数据集。准备的最后一步，我使用预先（AOT）编译来预编译 Llama 2 模型，以加快模型训练。

在集群上启动作业
现在，我可以使用 sbatch 命令开始模型训练作业。

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

您可以使用 squeue 命令来查看作业队列。训练作业开始运行后，SageMaker HyperPod 的弹性功能将自动启用。如果设置了 auto-resume 参数，SageMaker HyperPod 将自动检测硬件故障，根据需要替换节点，并从检查点恢复训练，如前述命令所示。

您可以在以下文件中查看模型训练作业的输出：

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

表示模型训练已开始的示例输出如下所示：

Epoch 0:  22%|██▏       | 4499/20101 [22:26:14<77:48:37, 17.95s/it, loss=2.43, v_num=5563, reduced_train_loss=2.470, gradient_norm=0.121, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.40]
Epoch 0:  22%|██▏       | 4500/20101 [22:26:32<77:48:18, 17.95s/it, loss=2.43, v_num=5563, reduced_train_loss=2.470, gradient_norm=0.121, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.40]
Epoch 0:  22%|██▏       | 4500/20101 [22:26:32<77:48:18, 17.95s/it, loss=2.44, v_num=5563, reduced_train_loss=2.450, gradient_norm=0.120, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.50]

要进一步监控和分析模型训练作业，您可以使用 SageMaker 托管 TensorBoard 或您选择的任何其他工具。

现已推出
SageMaker HyperPod 现已在以下 AWS 区域推出：美国东部（俄亥俄州）、美国东部（弗吉尼亚州北部）、美国西部（俄勒冈州）、亚太地区（新加坡）、亚太地区（悉尼）、亚太地区（东京）、欧洲地区（法兰克福）、欧洲地区（爱尔兰）和欧洲地区（斯德哥尔摩）。

了解更多信息：

有关定价信息和支持的集群实例类型列表，请参阅 Amazon SageMaker HyperPod
查看开发人员指南
访问 AWS 管理控制台，开始使用 SageMaker HyperPod 训练 FM

– Antje

附注：即使文章标题下只有一个署名，在 AWS 撰写博客文章也是团队的共同努力。因此，我要感谢 Brad Doran、Justin Pirtle、Ben Snyder、Pierre-Yves Aquilanti、Keita Watanabe 和 Verdi March 在示例代码方面给予的慷慨帮助，以及不吝分享他们在管理大规模模型训练基础设施、Slurm 和 SageMaker HyperPod 方面的专业知识。

亚马逊AWS官方博客

推出适合大规模分布式训练的专用基础设施 Amazon SageMaker HyperPod