面向 GPU 服务器的 SageMaker 无痛使用指南（三）—SageMaker HyperPod 集群

SageMaker HyperPod 集群简介

在进行大规模机器学习训练时，计算资源和数据量是两大关键挑战。传统的单机单卡训练方式已无法满足现代 GenAI 的 LLM/Stable Diffusion 等大模型对计算力的极高需求，也难以高效处理 TB 甚至 PB 级别的海量训练数据。这就需要利用 SageMaker Hyperpod 集群这样的分布式训练平台。

SageMaker Hyperpod 集群采用了基于 Slum 的 HPC 高性能弹性计算集群，能够实现跨机器跨 GPU 的大规模并行训练。它提供了原生的 IaaS 基础设施服务器，与普通 EC2 实例一样，可以自由地操控和部署客户所需要的模型和框架，并且充分发挥亚马逊云端可伸缩的计算能力，线性扩展训练吞吐量，显著缩短了大模型在海量数据集上的训练时间。

并且，SageMaker HyperPod 预配置了 SageMaker 的分布式训练库，使客户能够自动将训练工作负载拆分到数千个 GPU 服务器芯片上，因此可以并行处理工作负载以提高模型性能。它还通过定期保存检查点，帮助客户在分布式设置中进行数周或数月的无中断训练，当训练期间出现硬件故障时，SageMaker HyperPod 会自动检测故障，修复或更换有故障的实例，并从上次保存的检查点恢复训练，从而使客户无需手动管理此过程，

以下我们详细介绍 SageMaker Hyperpod 集群启动的步骤和使用方法。

Hyperpod 集群启动

和 EC2 服务器类似，SageMaker Hyperpod 集群启动需要设置你的 VPC 和 IAM 权限。

• IAM 权限设置
  • 启动 Hyperpod 集群的用户或者角色需要有相关的 policy，如 sagemaker cluster/cloudwatch/s3…etc
  • HyperPod 集群创建及运行所需要的 policy 整理如下

    {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": [
                    "logs:PutLogEvents",
                    "logs:CreateLogStream",
                    "logs:DescribeLogStreams"
                ],
                "Resource": [
                    "arn:aws:logs:*:*:log-group:/aws/sagemaker/Clusters/*:log-stream:*"
                ]
            },
            {
                "Effect": "Allow",
                "Action": [
                    "logs:CreateLogGroup"
                ],
                "Resource": [
                    "arn:aws:logs:*:*:log-group:/aws/sagemaker/Clusters/*"
                ]
            },
            {
                "Effect": "Allow",
                "Action": [
                    "cloudwatch:PutMetricData"
                ],
                "Resource": [
                    "*"
                ],
                "Condition": {
                    "StringEquals": {
                        "cloudwatch:namespace": "/aws/sagemaker/Clusters"
                    }
                }
            },
            {
                "Effect": "Allow",
                "Action": [
                    "s3:ListBucket",
                    "s3:GetObject"
                ],
                "Resource": [
                    "arn:aws:s3:::sagemaker-*"
                ]
            },
            {
                "Effect": "Allow",
                "Action": [
                    "ssmmessages:CreateControlChannel",
                    "ssmmessages:CreateDataChannel",
                    "ssmmessages:OpenControlChannel",
                    "ssmmessages:OpenDataChannel"
                ],
                "Resource": "*"
            },
            {
                "Sid": "AdditionToEnableVpcConfig",
                "Effect": "Allow",
                "Action": [
                    "ec2:CreateNetworkInterface",
                    "ec2:CreateNetworkInterfacePermission",
                    "ec2:DeleteNetworkInterface",
                    "ec2:DeleteNetworkInterfacePermission",
                    "ec2:DescribeNetworkInterfaces",
                    "ec2:DescribeVpcs",
                    "ec2:DescribeDhcpOptions",
                    "ec2:DescribeSubnets",
                    "ec2:DescribeSecurityGroups",
                    "ec2:DetachNetworkInterface"
                ],
                "Resource": "*"
            },
            {
                "Sid": "Addition2ToEnableVpcConfig",
                "Effect": "Allow",
                "Action": "ec2:CreateTags",
                "Resource": [
                    "arn:aws:ec2:*:*:network-interface/*"
                ]
            }
        ]
    }
    

    可以把以上 policy 直接加到之前已有的 IAM role 中，比如 Amazon SageMaker 笔记本实例上的 ExecutionRole。

• 部署 SageMaker Hyperpod 集群的 VPC 设置
  • Hyperpod 集群需要配置为 VPC 内，且 VPC 包含私有子网+公有子网，私有子网通过 nat 走 IGW 能够访问外网，这在训练任务需要下载公网数据，访问公网链接时是必要的。
  • 相应的 VPC 需要为 S3 和 DDB 建对应子网的终端节点，以便 Hyperpod 集群能在 DDB 中存储集群节点元数据等信息。
• 集群启动脚本
  在 Hyperpod 集群每个节点拉起时，会执行配置的生命周期脚本，可以定制自己需要的随集群启动的软硬件安装，如 pip，apt-get 第三方 lib 安装包等。
  Amazon 官方提供了示例的 Hyperpod 集群启动脚本，可以通过官方 github sample 链接下载直接上传到你的 S3 路径使用，官方示例启动脚本链接见附录。
  以官方链接上传自己的 S3 桶路径示例：

  BUCKET=<你的s3桶>
  s5cmd sync ./LifecycleScripts/base-config s3://${BUCKET}/LifeCycleScripts/
  aws s3 cp --recursive LifeCycleScripts/base-config s3://${BUCKET}/LifeCycleScripts/base-config
  s3://$BUCKET/LifeCycleScripts/

• Launch cluster 启动 Hyperpod 集群
  按上文章节配置好集群的 IAM 和 VPC 后，我们就可以配置 Hyperpod 集群所需的 GPU 服务器资源，及集群节点分组和队列。
  Hyperpod 集群的算力机资源通过一个 json 格式文件进行配置，cluster-config.json 配置示例如下：

  cluster-config.json
  [
    {
      "InstanceGroupName": "controller-machine",
      "InstanceType": "ml.c5.xlarge",
      "InstanceCount": 1,
      "LifeCycleConfig": {
        "SourceS3Uri": "s3://${BUCKET}/LifeCycleScripts/base-config/",
        "OnCreate": "on_create.sh"
      },
      "ExecutionRole": "${ROLE}",
      "ThreadsPerCore": 1
    },
    {
      "InstanceGroupName": "compute-nodes",
      "InstanceType": "ml.g5.2xlarge",
      "InstanceCount": 1,
      "LifeCycleConfig": {
        "SourceS3Uri": "s3://${BUCKET}/LifeCycleScripts/base-config/",
        "OnCreate": "on_create.sh"
      },
      "ExecutionRole": "${ROLE}",
      "ThreadsPerCore": 1
    }
  ]
  

  如上：

• LifeCycleConfig 为上文章节中的生命周期启动脚本的位置
• InstanceGroupName 为节点分组名
• Hyperpod 集群需要至少 2 个分组，一个为 controller 分组，做集群管理主节点，一个为 worker 分组，进行模型训练的主要算力机资源节点
• InstanceType：节点算力机 EC2 实例类型，如 g5（A10），p4d（A100）
• InstanceCount：节点实例数量，Hyperpod 集群通常可以管理数十到数百台的大规模节点

集群节点分组队列配置，也通过一个 json 格式配置文件进行设置，node provision 配置示例如下：

{
  "version": "1.0.0",
  "workload_manager": "slurm",
  "controller_group": "controller-machine",
  "worker_groups": [
    {
      "instance_group_name": "compute-nodes",
      "partition_name": "dev"
    }
  ]
}

• 如上，workload_manager 默认为 slum，未来 Hyperpod 集群还可以支持 k8s，关于 slum 资源管理的相关内容可以参阅附录中资料。
• controller_group 管理节点组的名字和 worker_group 计算节点组的名字需要与上一章节节点资源配置中 group 名字一致。
• 在 node provision 配置中还可以支持亚马逊云共享存储等其他相关配置，如 lustre fsx，这里不再赘述，感兴趣的小伙伴可以参阅附录中 SageMaker Hyperpod 配置文档。

完成以上配置后，我们即可以通过 cli 启动脚本拉起集群。拉起集群的命令示例如下：

aws sagemaker create-cluster \
    --cluster-name ml-hyperPod-cluster-2 \
    --instance-groups file://cluster-config.json \
    --region us-west-2

HyperPod 集群登陆

建立 SageMaker Hyperpod 集群之后，我们即可以像 EC2 实例一样登陆到集群的算力机节点上。登陆集群算力机节点有两种方式：

方式 1：ssm 客户端登陆

• Hyperpod 集群在每个节点实例都预置了 SSM agent 服务端，因此与 SageMaker Training Job/Inference Endpoint 一样，我们可以通过 ssm 客户端登陆 SageMaker Hyperpod 集群节点主机。
• 可以登陆 cotroller group，也可以登陆 worker group 节点组。
• hyperpod cluster 的 target id 命名规则为：sagemaker-cluster:${CLUSTER_ID}_${CONTROLLER_GROUP}-${INSTANCE_ID}
• 登陆节点的命令行示例如下：

  #CONTROLLER_GROUP=compute-nodes
  #INSTANCE_ID=i-06824a193cef10e69
  
  ##你的集群id，在hyperpod控制台可见
  CLUSTER_ID=frgms7n7l237 
  ##你的节点组名，见上文章节配置文件
  CONTROLLER_GROUP=controller-machine
  ##你的实例id，通过aws sagemaker list-cluster-nodes 可查看每个节点的id
  INSTANCE_ID=i-09bce424e1e96d199
  
  TARGET_ID=sagemaker-cluster:${CLUSTER_ID}_${CONTROLLER_GROUP}-${INSTANCE_ID}
  aws ssm start-session --target $TARGET_ID --region us-west-2
  

方式 2：关口机 ssh 客户端登陆

• 与 SageMaker Notebook 笔记本实例一样，我们可以通过一台同 VPC 网络的跳板机，使用 ssh 密钥文件方式登陆节点，这种方式不需要安装 ssm 客户端，也不需要查看集群 id 和节点 id 信息，通过内网 ip 地址实现 ssh 跳转，更加符合 EC2 使用的习惯。
• 这种方式下，只需要在同一个 VPC，public 子网开一台关口机。由于 VPC 网络安全组已经配置了 public 子网和 private 子网的互联互通，因此从关口机就可以跳转到 hpyperpod worker/controller 节点。
• 我们在启动脚本中加入 ssh pem 密钥登陆的自动初始化，以便关口机能和 Hyperpod 集群主机互信，详细步骤如下：
  • 客户新建 pem 密钥文件，用于在新建集群时加入节点 public key 公钥认证

    ##创建密钥或者使用你已经有的 pem 密钥
    ssh-keygen -t rsa -b 4096
    

  • 新建 add_ssh_pem.sh 脚本，作为生命周期启动脚本的一部分，将 pem 加入 HP 集群每台主机的 ubuntu 用户目录下 .ssh 子目录的 authorized_keys 中，以便能免密登陆

    #!/bin/bash
    
    set -e
    set -x
    
    PEM_FILE="pub.pem"
    PEM_PRIV_FILE="priv.pem"
    
    main() {
      if [[ ! -f $PEM_FILE ]]; then
        echo "Shared user file $PEM_FILE does not exist. Skipping adding ssh pem."
        exit 0
      fi
    
      {
      cat $PEM_FILE|tr '\n' ' '
      } >> /home/ubuntu/.ssh/authorized_keys
      cp $PEM_PRIV_FILE  /home/ubuntu/.ssh/id_rsa
      chmod 400 /home/ubuntu/.ssh/id_rsa
      chown ubuntu:ubuntu /home/ubuntu/.ssh/id_rsa
      #echo "ssh-rsa ">> /home/ubuntu/.ssh/authorized_keys
      #cat $PEM_FILE >> /home/ubuntu/.ssh/authorized_keys
      #echo " 203.pem "  >> /home/ubuntu/.ssh/authorized_keys
    }
    
    main "$@"

  • 将 pem 文件和 add_ssh_pem.sh 脚本放置于集群启动脚本同级目录
    
  • 修改集群生命周期脚本（lifecycle_script.py），增加 ssh pem 操作

    def main(args):
        params = ProvisioningParameters(args.provisioning_parameters)
        resource_config = ResourceConfig(args.resource_config)
    
        ExecuteBashScript("./add_users.sh").run()
    
        fsx_dns_name, fsx_mountname = params.fsx_settings
        if fsx_dns_name and fsx_mountname:
            print(f"Mount fsx: {fsx_dns_name}. Mount point: {fsx_mountname}")
            ExecuteBashScript("./mount_fsx.sh").run(fsx_dns_name, fsx_mountname)
    
        ...省略
        
        ## add ssh key authentication support
        ExecuteBashScript("./add_ssh_pem.sh").run()
    

  • VPC 跳板机准备
    同一 VPC 下 public 子网的 EC2 实例即可，在 ubuntu 用户.ssh 下放置 pem 密钥文件，并修改相应权限

    chmod 400 /home/ubuntu/.ssh/<你的密钥名>.pem

     

  • 登陆 worker 或者 controller
    至此 ssh 配置已经完成，可以像 EC2 一样 ssh 登陆节点。
    通过 scontrol 或者 list node 可以获得 hyperpod 集群 controller/worker 节点主机内网 IP 或者域名。

    scontrol show node
    NodeName=ip-10-1-83-209 Arch=x86_64 CoresPerSocket=8 
       CPUAlloc=0 CPUEfctv=16 CPUTot=16 CPULoad=0.05
       AvailableFeatures=(null)
       ActiveFeatures=(null)
       Gres=(null)
       NodeAddr=10.1.83.209 NodeHostName=ip-10-1-83-209 Version=23.02.3
       OS=Linux 5.15.0-1047-aws #52~20.04.1-Ubuntu SMP Thu Sep 21 10:05:54 UTC 2023 
       RealMemory=32768 AllocMem=0 FreeMem=26556 Sockets=1 Boards=1
       State=IDLE+CLOUD ThreadsPerCore=2 TmpDisk=0 Weight=1 Owner=N/A MCS_label=N/A
       Partitions=dev 
       BootTime=2023-12-28T09:07:28 SlurmdStartTime=2023-12-28T09:14:51
       LastBusyTime=2023-12-28T09:38:06 ResumeAfterTime=None
       CfgTRES=cpu=16,mem=32G,billing=16
       AllocTRES=
       CapWatts=n/a
       CurrentWatts=0 AveWatts=0
       ExtSensorsJoules=n/s ExtSensorsWatts=0 ExtSensorsTemp=n/s
    

    我们试一下在跳板机用 ssh 密钥文件登陆节点

    ssh -vvv  -i  203.pem ubuntu@10.1.83.209

    可以看到正常登陆，ssh 进入节点后是 ubuntu 用户

    The list of available updates is more than a week old.
    To check for new updates run: sudo apt update
    
    Last login: Tue Mar  5 02:23:11 2024 from 10.0.31.208
    ubuntu@ip-10-1-83-209:~$ 
    

集群任务管理

Hyeperpod 集群预置了 Slum，SageMaker sdk 等各种集群管理工具，我们可以 sagemake cli 或者 slum cli 查看集群。

aws sagemaker describe-cluster --cluster-name ml-hyperPod-cluster --region us-west-2
sh-4.2$ aws sagemaker list-clusters
{
    "ClusterSummaries": [
        {
            "ClusterArn": "arn:aws:sagemaker:us-west-2:687912291502:cluster/evrfrz4dddzq",
            "ClusterName": "ml-hyperPod-cluster",
            "CreationTime": 1703677276.055,
            "ClusterStatus": "InService"
        }
    ]
}
i-09bce424e1e96d199
sh-4.2$ aws sagemaker list-cluster-nodes --cluster-name ml-hyperPod-cluster --region us-west-2 
{
    "ClusterNodeSummaries": [
        {
            "InstanceGroupName": "compute-nodes",
            "InstanceId": "i-0002848a9fc9640b7",
            "InstanceType": "ml.c5.2xlarge",
            "LaunchTime": 1703677284.813,
            "InstanceStatus": {
                "Status": "Running",
                "Message": ""
            }
        },
        {
            "InstanceGroupName": "controller-machine",
            "InstanceId": "i-064e83d3fa8e53fc0",
            "InstanceType": "ml.c5.2xlarge",
            "LaunchTime": 1703677283.042,
            "InstanceStatus": {
                "Status": "Running",
                "Message": ""
            }
        }
    ]
}


scontrol show node
NodeName=ip-10-1-83-209 Arch=x86_64 CoresPerSocket=8 
   CPUAlloc=0 CPUEfctv=16 CPUTot=16 CPULoad=0.05
   AvailableFeatures=(null)
   ActiveFeatures=(null)
   Gres=(null)
   NodeAddr=10.1.83.209 NodeHostName=ip-10-1-83-209 Version=23.02.3
   OS=Linux 5.15.0-1047-aws #52~20.04.1-Ubuntu SMP Thu Sep 21 10:05:54 UTC 2023 
   RealMemory=32768 AllocMem=0 FreeMem=26556 Sockets=1 Boards=1
   State=IDLE+CLOUD ThreadsPerCore=2 TmpDisk=0 Weight=1 Owner=N/A MCS_label=N/A
   Partitions=dev 
   BootTime=2023-12-28T09:07:28 SlurmdStartTime=2023-12-28T09:14:51
   LastBusyTime=2023-12-28T09:38:06 ResumeAfterTime=None
   CfgTRES=cpu=16,mem=32G,billing=16
   AllocTRES=
   CapWatts=n/a
   CurrentWatts=0 AveWatts=0
   ExtSensorsJoules=n/s ExtSensorsWatts=0 ExtSensorsTemp=n/s

aws sagemaker delete-cluster --cluster-name ml-hyperPod-cluster --region us-west-2 

我们可以用 slum cli 进行多机集群的训练任务提交：

# 查看queue状态
sinfo
# PARTITION AVAIL TIMELIMIT NODES STATE NODELIST
# dev* up infinite 1 idle ip-10-1-82-20
# dev* up infinite 1 idle ip-10-1-82-21

# 提交作业1 - 交互式提交作业，等待结束
srun hostname -p
# ip-10-1-82-20
srun -N1  -p dev "df" "-h"

# 提交作业2 - sbatch 后台提交作业，
sbatch -p dev -N4 myscript

# 提交作业3 - salloc 分配模式作业
salloc

# 查看作业队列
squeue

# JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
#    14       dev    sleep   ubuntu  R       0:05      1 ip-10-1-82-20

HyperPod 集群远程调试

和 SageMaker Training Job 类似，当生产环境需要 trouble shooting/perf turning 等运维优化工作时，同样可以通过远程调试方式，在 IDE 中连接到 Hyperpod 节点主机上的训练脚本和应用程序，进行快速的定位和断点调试。

具体方法如下：

• HyperPod 集群创建时增加上文章节中提到的 ssh 登陆初始化脚本
• 创建同一 VPC 集群的 EC2 关口机（与上文章节一样，此处只用于 ssh 隧道，因此可以选择小机型如 c5.large）
• 在本机打 ssh 隧道

  ssh -i /Users/tangqy/AMAZON/keypairs/ec203.pem -N -L 8422:10.1.124.244:22 ubuntu@ec2-35-87-74-61.us-west-2.compute.amazonaws.com

• 其中 10.1.124.244 为 HyperPod 集群实例的私网 ip
• ec2-35-87-74-61.us-west-2.compute.amazonaws.com 为 EC2 关口机公网域名
• 在 IDE（vscode 为例）配置远程 ssh 主机

  Host hyperpod_forward
   HostName localhost2
   IdentityFile /Users/tangqy/AMAZON/keypairs/ec203.pem
   User ubuntu
  

• 在 vscode 里面使用刚才配置的 hyperpod 的远程主机，点击连接
  

• 显示连接后，点击左侧资源文件夹，可以看到 hyperpod 节点主机上的目录文件
• 点击 python 脚本文件，即可以单步调试（注意：需要安装对应的 python remote debug 插件）
  

HyperPod 集群模型训练

如上文中所述，Hyperpod 集群特别适合大规模集群分布式训练，由于其提供了底层 IaaS 基础设施的接入，因此可以方便的使用业界流行的各种分布式框架，如 accelerate，Deepspeed…etc。

以下我们详细介绍在 Hyperpod 集群上进行模型分布式训练的步骤方法。

训练脚本放共享存储

与 EC2 实例一样，Hyperpod 集群实例上可以挂载各种共享存储，如 EFS，Lustre，S3 等，此处我们以 mount-s3 为例。

mount-s3 共享存储安装及挂载脚本示例：

###下载 s3mount
wget  https://s3.amazonaws.com/mountpoint-s3-release/latest/x86_64/mount-s3.deb
sudo apt-get install -y  ./mount-s3.deb
srun -N2 "mkdir" "<你的挂载目录>"
## 在所有节点上挂载
srun -N2 "sudo" "mount-s3"  "—allow-other" "—allow-overwrite"  "sagemaker-us-west-2-687912291502" "/mnt/sm_bucket"
## unmount s3
# srun -N2 "sudo" "umount" "/mnt/sm_bucket"

运行环境创建

• 很多客户喜好使用 conda/venv 虚拟环境，以便不同版本的对比测试和依赖隔离
• Hyperpod 集群可以使用各种虚拟环境，conda env/python venv/anaconda…etc 均可
• 以 miniconda 环境创建为例
• miniconda 环境创建脚本示例 create_env.sh

  #!/usr/bin/env bash
  set -ex
   
  wget  https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
  chmod +x  Miniconda3-latest-Linux-x86_64.sh
  ./Miniconda3-latest-Linux-x86_64.sh  -b -f -p ./miniconda3
   
  source  ./miniconda3/bin/activate
   
  conda create -y -n py310  python=3.10
   
  source activate py310
   
  # Install PyTorch
  pip install torch torchvision  torchaudio
  pip install packaging  transformers accelerate ninja tensorboard h5py datasets
  pip uninstall -y deepspeed  && pip3 install deepspeed
  pip install  transformers==4.28.0
  pip install lightning
  sudo apt install jq
   
  # create output dir
  mkdir tensorboard
  mkdir checkpoints
  

使用 slum srun 命令安装该 conda env 在所有 node 上：

srun -N2  "bash" "/mnt/sm_bucket/scripts/create_env.sh"

以下我们看一看主流的各种分布式训练框架在 Hyperpod 上如何方便地拉起训练任务。

Pytorch DDP（单机多卡）训练运行方式示例

我们以一台 g5.12xlarge 4 卡机器，运行一个 pytorch ddp 的数据分布式，单机多卡的模型训练为例。

如上文所述，数据通过 mount-s3 挂载在节点目录（e.g: /mnt/sm_bucket/），因此无论哪台节点均可访问 S3 桶路径下的 dataset 数据集（e.g： /mnt/sm_bucket/dataset/），也可以访问到共享存储上的脚本文件（e.g：/mnt/sm_bucket/scripts/）。

conda 环境中已经安装 pytorch 2.1 + python310 的依赖库，因此我们直接通过 torch distrubute 命令行即可拉起该训练任务：

conda activate py310
python3 -m torch.distributed.launch \
        —nproc_per_node=4  \
        /mnt/sm_bucket/scripts/parallel_syncnet_tanh_ddp.py  \
        —gradient_accumulation_steps 16 \
        —data_root /mnt/sm_bucket/std_dataset  \
        —checkpoint_dir /tmp/trained_syncnet/

Accelerate（多机多卡）训练运行方式示例：

• accelerate 框架是 Stable Diffusion/Transformers 类大模型训练常用的分布式训练框架，支持 pipeline 流水线并行。
• 在 Hyperpod 集群上进行 Accelerate 框架的多机多卡分布式训练，我们可以通过 slum 提供的 nodeid 等环境变量方便的设置 accelerate yaml 配置中的 node_rank，从而实现在不同节点实例上 accelerate 的配置，以便拉起 accelerate 的分布式训练集群。

以 2 台 P4d（A100）机器的分布式训练为例，我们可以编写如下 run_demo.sh 训练脚本示例：

#!/bin/bash
 
 
# Get the slum parameter
echo "$((SLURM_NODEID + 1))"
echo "$SLURM_NPROCS"
echo "$SLURM_NTASKS"
node_name=$(hostname -s)
node_id=$((SLURM_NODEID + 1))
 
# Define the config file based on the master node
if [ "$node_name" == "ip-10-1-124-244" ]; then
     config_file="/mnt/sm_bucket/accelerate/examples/accelerate_config_${node_id}.yaml"
else
     config_file="/mnt/sm_bucket/accelerate/examples/accelerate_config_${node_id}.yaml"
fi
 
# Launch the accelerate script with the appropriate config file
/home/ubuntu/.local/bin/accelerate launch —config_file=$config_file  /mnt/sm_bucket/accelerate/examples/nlp_example.py

如上 SLURM_NODEID 即为 hyperpod 集群预置的 slum 环境变量，accelerate_config_${node_id}.yaml 即为不同节点上 accelerate 配置的 yaml 文件，通过不同的节点，可以区分 accelerate 集群中的 master node，和 node_rank 等配置。

分布式运行启动脚本示例：

sbatch -N2 "bash" "/mnt/sm_bucket/accelerate/examples/run_demo.sh"

DeepSpeed（多机多卡）训练运行示例：

• DeepSpeed 是提供了业界广受好评的 ZeRO stage 的分布式训练框架，特别适合模型参数规模庞大的 LLM，比如 Llama 70B。
• 在 Hyperpod 集群上，如上文章节所述，我们已经能够拿到 GPU 服务器节点主机的内网 ip，且节点之间已经实现 ssh 免密互信，因此我们可以直接使用 Deepspeed —hostfile 方式，在一台计算节点上自动拉起 DeepSpeed 多机集群，不需要在每一台节点启动指定 master node，node id 变量的 torch run 脚本。
• 多机 deepspeed 拉起脚本 deepspeed_test.sh 示例：

  #!/bin/bash
  MODEL_S3_BUCKET="sagemaker-us-west-2-687912291502"    
   
  export NCCL_IB_DISABLE=1
  export NCCL_SOCKET_IFNAME=eth0
   
  chmod +x /home/ubuntu/s5cmd
  /home/ubuntu/s5cmd  sync s3://$MODEL_S3_BUCKET/llama/pretrain/pinkmanlove/llama-7b-hf/*  /tmp/llama_pretrain/
   
  #cd FastChat && pip  install -e . && cd ..
   
  DEEPSPEED_OPTS="""
     /mnt/sm_bucket/FastChat/fastchat/train/train_mem.py 
      —deepspeed ds.json 
      —model_name_or_path  "/tmp/llama_pretrain/" 
      —data_path  data/dummy_conversation.json 
      —output_dir "/tmp/llama_out"  
      —num_train_epochs 1 
      —per_device_train_batch_size 1 
      —per_device_eval_batch_size  1 
      —gradient_accumulation_steps 4 
      —evaluation_strategy "no" 
      —save_strategy "no" 
      —save_steps 2000 
      —save_total_limit 1 
      —learning_rate 2e-5 
      —weight_decay 0. 
      —warmup_ratio 0.03 
      —lr_scheduler_type "cosine" 
      —logging_steps 1 
      —cache_dir '/tmp' 
      —model_max_length 2048 
      —gradient_checkpointing True 
      —lazy_preprocess True 
      —fp16 True 
      —tf32 True 
      —report_to "wandb"
  """    
   
  ## for mutiple node
  CMD="/home/ubuntu/.local/bin/deepspeed  —hostfile hostfile ${DEEPSPEED_OPTS}“
  ## for single node
  CMD="python3  ${DEEPSPEED_OPTS}"
  echo ${CMD}
  ${CMD} 2>&1 
  echo "begin to upload  trained model"
  bash /home/ubuntu/s5cmd sync  /tmp/llama_out s3://$MODEL_S3_BUCKET/llama/output/$(date  +%Y-%m-%d-%H-%M-%S)/
  

  deepspeed hostfile 配置文件示例（g5.12xlarge 4 卡机 2 台）

  10.1.119.249  slots=4
  10.1.124.244  slots=4 
  

  运行示例（在一台计算节点上拉起即可）

  bash /mnt/sm_bucket/FastChat/deepspeed_test.sh

总结

SageMaker HyperPod 集群是一种基于 Slurm 的高性能弹性计算集群，可以实现跨机器跨 GPU 的大规模并行训练。

它提供了原生的 IaaS 基础设施服务器，可以自由操控和部署各种模型和框架，充分发挥亚马逊云端可伸缩的计算能力，显著缩短了大模型在海量数据集上的训练时间。

本文详细介绍了 SageMaker Hyperpod 集群的启动、配置、连接、管理和在其上进行分布式训练的方法，包括:

• 配置和启动 Hyperpod 集群
• 使用 ssm 客户端和 ssh 密钥登陆集群节点
• 配置 Hyperpod 远程调试环境
• 挂载共享存储创建运行环境
• 使用 Pytorch DDP、Accelerate、DeepSpeed 等框架进行分布式训练示例

总的来说，SageMaker Hyperpod 提供了云端可扩展的大规模分布式训练能力，结合各种流行的分布式训练框架，使得训练大模型更加便捷高效。

附录

Hyperpod 官方文档：https://docs.aws.amazon.com/sagemaker/latest/dg/cluster.html

Hyperpod 官方示例启动脚本：https://github.com/aws-samples/sagemaker-hyperpod-lifecycle-configs

Slurm 资源管理介绍：https://slurm.schedmd.com/documentation.htm

本篇作者