SageMaker HyperPod 集群简介
在进行大规模机器学习训练时,计算资源和数据量是两大关键挑战。传统的单机单卡训练方式已无法满足现代生成式 AI 的 LLM 或 Stable Diffusion 等大模型对计算力的极高需求,也难以高效处理 TB 甚至 PB 级别的海量训练数据。
这就需要利用 Amazon SageMaker Hyperpod 集群这样的分布式训练平台。
Amazon SageMaker Hyperpod 集群采用了基于 Slum 的 HPC 高性能弹性计算集群,能够实现跨机器跨 GPU 的大规模并行训练。
它提供了原生的 IaaS 基础设施服务器,与普通 Amazon EC2 实例一样,可以自由地操控和部署客户所需要的模型和框架,并且充分发挥亚马逊云科技云端可伸缩的计算能力,线性扩展训练吞吐量,显著缩短了大模型在海量数据集上的训练时间。
并且,Amazon SageMaker HyperPod 预配置了 Amazon SageMaker 的分布式训练库,使客户能够自动将训练工作负载拆分到数千个 GPU 服务器芯片上,因此可以并行处理工作负载以提高模型性能。
它还通过定期保存检查点,帮助客户在分布式设置中进行数周或数月的无中断训练,当训练期间出现硬件故障时,Amazon SageMaker HyperPod 会自动检测故障,修复或更换有故障的实例,并从上次保存的检查点恢复训练,从而使客户无需手动管理此过程,
以下我们详细介绍 Amazon SageMaker Hyperpod 集群启动的步骤和使用方法。
Hyperpod 集群启动
和 Amazon EC2 服务器类似,Amazon SageMaker Hyperpod 集群启动需要设置你的 Amazon VPC 和 Amazon IAM 权限。
Amazon 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 直接加到之前已有的 Amazon IAM role 中,比如 Amazon SageMaker 笔记本实例上的 Execution Role。
部署 Amazon SageMaker Hyperpod 集群的 Amazon VPC 设置
Hyperpod 集群需要配置为 Amazon VPC 内,且 Amazon VPC 包含私有子网+公有子网,私有子网通过 NAT 走 IGW 能够访问外网,这在训练任务需要下载公网数据,访问公网链接时是必要的。
相应的 Amazon VPC 需要为 Amazon S3 和 DDB 建对应子网的终端节点,以便 Hyperpod 集群能在 DDB 中存储集群节点元数据等信息。
集群启动脚本
在 Hyperpod 集群每个节点拉起时,会执行配置的生命周期脚本,可以定制自己需要的随集群启动的软硬件安装,如 pip,apt-get 第三方 lib 安装包等。
亚马逊云科技官方提供了示例的 Hyperpod 集群启动脚本,可以通过官方 github sample 链接下载直接上传到你的 Amazon S3 路径使用,官方示例启动脚本链接见附录。
以官方链接上传自己的 Amazon 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 集群
按上文章节配置好集群的 Amazon IAM 和 Amazon 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:节点算力机 Amazon 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,这里不再赘述,感兴趣的小伙伴可以参阅附录中 Amazon SageMaker Hyperpod 配置文档。
完成以上配置后,我们即可通过 cli 启动脚本拉起集群。拉起集群的命令示例如下:
aws sagemaker create-cluster \
--cluster-name ml-hyperPod-cluster-2 \
--instance-groups file://cluster-config.json \
--region us-west-2
左右滑动查看完整示意
HyperPod 集群登陆
建立 Amazon SageMaker Hyperpod 集群之后,我们即可像 Amazon EC2 实例一样登陆到集群的算力机节点上。登陆集群算力机节点有两种方式:
方式 1:SSM 客户端登陆
Hyperpod 集群在每个节点实例都预置了 SSM agent 服务端,因此与 Amazon SageMaker Training Job/Inference Endpoint 一样,我们可以通过 SSM 客户端登陆 Amazon 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 客户端登陆
与 Amazon SageMaker Notebook 笔记本实例一样,我们可以通过一台同 Amazon VPC 网络的跳板机,使用 .ssh 密钥文件方式登陆节点,这种方式不需要安装 ssm 客户端,也不需要查看集群 ID 和节点 ID 信息,通过内网 IP 地址实现 .ssh 跳转,更加符合 Amazon EC2 使用的习惯。
这种方式下,只需要在同一个 Amazon VPC,public 子网开一台关口机。由于 Amazon 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()
左右滑动查看完整示意
Amazon VPC 跳板机准备
同一 Amazon VPC 下 public 子网的 Amazon EC2 实例即可,在 ubuntu 用户 .ssh 下放置 pem 密钥文件,并修改相应权限
chmod 400 /home/ubuntu/.ssh/<你的密钥名>.pem
左右滑动查看完整示意
登陆 worker 或者 controller
至此 .ssh 配置已经完成,可以像 Amazon 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、Amazon 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 集群远程调试
和 Amazon SageMaker Training Job 类似,当生产环境需要 trouble shooting/perf turning 等运维优化工作时,同样可以通过远程调试方式,在 IDE 中连接到 Hyperpod 节点主机上的训练脚本和应用程序,进行快速的定位和断点调试。
具体方法如下:
HyperPod 集群创建时增加上文章节中提到的 ssh 登陆初始化脚本
创建同一 Amazon VPC 集群的 Amazon 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 为 Amazon 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 集群上进行模型分布式训练的步骤方法。
训练脚本放共享存储
与 Amazon EC2 实例一样,Hyperpod 集群实例上可以挂载各种共享存储,如 Amazon EFS、Lustre、Amazon 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/),因此无论哪台节点均可访问 Amazon 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
左右滑动查看完整示意
总结
Amazon SageMaker HyperPod 集群是一种基于 Slurm 的高性能弹性计算集群,可以实现跨机器跨 GPU 的大规模并行训练。
它提供了原生的 IaaS 基础设施服务器,可以自由操控和部署各种模型和框架,充分发挥亚马逊云科技云端可伸缩的计算能力,显著缩短了大模型在海量数据集上的训练时间。
本文详细介绍了 Amazon SageMaker Hyperpod 集群的启动、配置、连接、管理和在其上进行分布式训练的方法,包括:
配置和启动 Hyperpod 集群
使用 SSM 客户端和 .ssh 密钥登陆集群节点
配置 Hyperpod 远程调试环境
挂载共享存储创建运行环境
使用 Pytorch DDP、Accelerate、DeepSpeed 等框架进行分布式训练示例
总的来说,Amazon 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.html
本篇作者
唐清原
亚马逊云科技高级解决方案架构师,负责数据分析、人工智能与机器学习、产品服务架构设计以及解决方案。拥有超过10年数据领域研发及架构设计经验。在大数据 BI、数据湖、推荐系统、MLOps 等平台项目有丰富实战经验。
龙斌
亚马逊云科技解决方案架构师,负责协助客户业务系统上云的解决方案架构设计和咨询,现致力于容器和机器学习相关领域的研究。
星标不迷路,开发更极速!
关注后记得星标「亚马逊云开发者」
听说,点完下面4个按钮
就不会碰到bug了!
点击阅读原文查看博客!获得更详细内容!