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

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

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

创建和管理集群

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

在接下来的设置中，提供一个集群名称，并根据所选实例类型和要分配给每个实例组的实例数量配置实例组。

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

使用 Amazon CLI

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

// demo-cluster.json[
    {
        "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": "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
    }]

为创建集群，我运行了以下 Amazon CLI 命令：

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

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

您还可以选择附加  Amazon FSx for Lustre 等共享文件系统。欲使用 FSx for Lustre，需要使用  Amazon Virtual Private Cloud (Amazon VPC) 配置来设置集群。这是一个  Amazon CloudFormation 模板，展示了如何 创建 SageMaker VPC 以及如何 部署 FSx for Lustre。

连接您的集群

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

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

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

使用 Slurm 在集群上调度和运行作业

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

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

注意

Llama 2 模型由  Meta 管理。您可以通过  Meta 请求访问页面请求访问权限。

设置集群环境

SageMaker HyperPod 支持在各种环境中进行训练，其中包括  Conda 、 venv 、 Docker 和  enroot 。我按照  README 中的说明，构建了虚拟环境  amazon_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 的弹性功能就会自动启用。SageMaker HyperPod 会自动检测硬件故障，根据需要替换节点，如果设置了  auto-resume 参数，还会从检查点恢复训练，如上一条命令所示。

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

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 现已在以下亚马逊云科技商业服务区域推出：美国东部（俄亥俄州）、美国东部（北弗吉尼亚州）、美国西部（俄勒冈州）、亚太地区（新加坡）、亚太地区（悉尼）、亚太地区（东京）、欧洲（法兰克福）、欧洲（爱尔兰）和欧洲（斯德哥尔摩）地区。

了解更多信息：

● 有关定价信息和支持的集群实例类型列表，请参阅  Amazon SageMaker HyperPod。

● 查看 开发人员指南

● 访问  亚马逊云科技管理控制台，开始使用 SageMaker HyperPod 训练基础模型

文章来源：  https://aws.amazon.com/blogs/aws/introducing-amazon-sagemaker-hyperpod-a-purpose-built-infrastructure-for-distributed-training-at-scale?trk=cndc-detail/

—  Antje

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

作者 Antje Barth