亚马逊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 整理如下
可以把以上 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 桶路径示例:
- Launch cluster 启动 Hyperpod 集群
按上文章节配置好集群的 IAM 和 VPC 后,我们就可以配置 Hyperpod 集群所需的 GPU 服务器资源,及集群节点分组和队列。
Hyperpod 集群的算力机资源通过一个 json 格式文件进行配置,cluster-config.json 配置示例如下:如上:
- LifeCycleConfig 为上文章节中的生命周期启动脚本的位置
- InstanceGroupName 为节点分组名
- Hyperpod 集群需要至少 2 个分组,一个为 controller 分组,做集群管理主节点,一个为 worker 分组,进行模型训练的主要算力机资源节点
- InstanceType:节点算力机 EC2 实例类型,如 g5(A10),p4d(A100)
- InstanceCount:节点实例数量,Hyperpod 集群通常可以管理数十到数百台的大规模节点
集群节点分组队列配置,也通过一个 json 格式配置文件进行设置,node provision 配置示例如下:
- 如上,workload_manager 默认为 slum,未来 Hyperpod 集群还可以支持 k8s,关于 slum 资源管理的相关内容可以参阅附录中资料。
- controller_group 管理节点组的名字和 worker_group 计算节点组的名字需要与上一章节节点资源配置中 group 名字一致。
- 在 node provision 配置中还可以支持亚马逊云共享存储等其他相关配置,如 lustre fsx,这里不再赘述,感兴趣的小伙伴可以参阅附录中 SageMaker Hyperpod 配置文档。
完成以上配置后,我们即可以通过 cli 启动脚本拉起集群。拉起集群的命令示例如下:
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}
- 登陆节点的命令行示例如下:
方式 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 公钥认证
- 新建 add_ssh_pem.sh 脚本,作为生命周期启动脚本的一部分,将 pem 加入 HP 集群每台主机的 ubuntu 用户目录下 .ssh 子目录的 authorized_keys 中,以便能免密登陆
- 将 pem 文件和 add_ssh_pem.sh 脚本放置于集群启动脚本同级目录
- 修改集群生命周期脚本(lifecycle_script.py),增加 ssh pem 操作
- VPC 跳板机准备
同一 VPC 下 public 子网的 EC2 实例即可,在 ubuntu 用户.ssh 下放置 pem 密钥文件,并修改相应权限 - 登陆 worker 或者 controller
至此 ssh 配置已经完成,可以像 EC2 一样 ssh 登陆节点。
通过 scontrol 或者 list node 可以获得 hyperpod 集群 controller/worker 节点主机内网 IP 或者域名。我们试一下在跳板机用 ssh 密钥文件登陆节点
可以看到正常登陆,ssh 进入节点后是 ubuntu 用户
集群任务管理
Hyeperpod 集群预置了 Slum,SageMaker sdk 等各种集群管理工具,我们可以 sagemake cli 或者 slum cli 查看集群。
我们可以用 slum cli 进行多机集群的训练任务提交:
HyperPod 集群远程调试
和 SageMaker Training Job 类似,当生产环境需要 trouble shooting/perf turning 等运维优化工作时,同样可以通过远程调试方式,在 IDE 中连接到 Hyperpod 节点主机上的训练脚本和应用程序,进行快速的定位和断点调试。
具体方法如下:
- HyperPod 集群创建时增加上文章节中提到的 ssh 登陆初始化脚本
- 创建同一 VPC 集群的 EC2 关口机(与上文章节一样,此处只用于 ssh 隧道,因此可以选择小机型如 c5.large)
- 在本机打 ssh 隧道
- 其中 10.1.124.244 为 HyperPod 集群实例的私网 ip
- ec2-35-87-74-61.us-west-2.compute.amazonaws.com 为 EC2 关口机公网域名
- 在 IDE(vscode 为例)配置远程 ssh 主机
- 在 vscode 里面使用刚才配置的 hyperpod 的远程主机,点击连接
HyperPod 集群模型训练
如上文中所述,Hyperpod 集群特别适合大规模集群分布式训练,由于其提供了底层 IaaS 基础设施的接入,因此可以方便的使用业界流行的各种分布式框架,如 accelerate,Deepspeed…etc。
以下我们详细介绍在 Hyperpod 集群上进行模型分布式训练的步骤方法。
训练脚本放共享存储
与 EC2 实例一样,Hyperpod 集群实例上可以挂载各种共享存储,如 EFS,Lustre,S3 等,此处我们以 mount-s3 为例。
mount-s3 共享存储安装及挂载脚本示例:
运行环境创建
- 很多客户喜好使用 conda/venv 虚拟环境,以便不同版本的对比测试和依赖隔离
- Hyperpod 集群可以使用各种虚拟环境,conda env/python venv/anaconda…etc 均可
- 以 miniconda 环境创建为例
- miniconda 环境创建脚本示例 create_env.sh
使用 slum srun 命令安装该 conda env 在所有 node 上:
以下我们看一看主流的各种分布式训练框架在 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 命令行即可拉起该训练任务:
Accelerate(多机多卡)训练运行方式示例:
- accelerate 框架是 Stable Diffusion/Transformers 类大模型训练常用的分布式训练框架,支持 pipeline 流水线并行。
- 在 Hyperpod 集群上进行 Accelerate 框架的多机多卡分布式训练,我们可以通过 slum 提供的 nodeid 等环境变量方便的设置 accelerate yaml 配置中的 node_rank,从而实现在不同节点实例上 accelerate 的配置,以便拉起 accelerate 的分布式训练集群。
以 2 台 P4d(A100)机器的分布式训练为例,我们可以编写如下 run_demo.sh 训练脚本示例:
如上 SLURM_NODEID 即为 hyperpod 集群预置的 slum 环境变量,accelerate_config_${node_id}.yaml 即为不同节点上 accelerate 配置的 yaml 文件,通过不同的节点,可以区分 accelerate 集群中的 master node,和 node_rank 等配置。
分布式运行启动脚本示例:
DeepSpeed(多机多卡)训练运行示例:
- DeepSpeed 是提供了业界广受好评的 ZeRO stage 的分布式训练框架,特别适合模型参数规模庞大的 LLM,比如 Llama 70B。
- 在 Hyperpod 集群上,如上文章节所述,我们已经能够拿到 GPU 服务器节点主机的内网 ip,且节点之间已经实现 ssh 免密互信,因此我们可以直接使用 Deepspeed —hostfile 方式,在一台计算节点上自动拉起 DeepSpeed 多机集群,不需要在每一台节点启动指定 master node,node id 变量的 torch run 脚本。
- 多机 deepspeed 拉起脚本 deepspeed_test.sh 示例:
deepspeed hostfile 配置文件示例(g5.12xlarge 4 卡机 2 台)
运行示例(在一台计算节点上拉起即可)
总结
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