亚马逊AWS官方博客

面向 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

本篇作者

唐清原

亚马逊云科技高级解决方案架构师,负责 Data Analytic & AIML 产品服务架构设计以及解决方案。10+数据领域研发及架构设计经验,历任 IBM 咨询顾问,Oracle 高级咨询顾问,澳新银行数据部领域架构师职务。在大数据 BI,数据湖,推荐系统,MLOps 等平台项目有丰富实战经验。

龙斌

亚马逊云科技解决方案架构师,负责协助客户业务系统上云的解决方案架构设计和咨询,现致力于容器和机器学习相关领域的研究。