在集群上安排 Slurm 作业 SageMaker HyperPod - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

在集群上安排 Slurm 作业 SageMaker HyperPod

您可以使用标准 Slurm sbatch 或srun命令启动训练作业。例如，要启动 8 节点训练作业，可以在各种环境中运行srun -N 8 --exclusive train.sh SageMaker HyperPod 支持训练，包括conda、venvdocker、和enroot。您可以通过在 SageMaker HyperPod 集群上运行生命周期脚本来配置 ML 环境。您还可以选择附加共享文件系统，例如 Amazon FSx，它也可以用作虚拟环境。

以下示例说明如何在具有 Amazon FSx 共享文件系统的集群上使用完全分片数据并行 (FSDP) 技术运行训练 Llama-2 的 SageMaker HyperPod 作业。您还可以从 Awsome 分布式训练 GitHub 资料库中找到更多示例。

提示

所有 SageMaker HyperPod 示例都可在 Awsome 分布式训练 GitHub 存储库的3.test_cases文件夹中找到。

  1. 克隆 Awsome 分布式训练 GitHub 存储库，并将训练作业示例复制到您的 Amazon FSx 文件系统。

    $ TRAINING_DIR=/fsx/users/my-user/fsdp
$ git clone https://github.com/aws-samples/awsome-distributed-training/

  2. 运行 create_conda_env.sh 脚本。这将在您的 Amazon FSx 文件系统上创建一个conda环境。确保群集中的所有节点都可以访问文件系统。

  3. 通过启动单节点 slurm 作业来构建虚拟 Conda 环境，如下所示。

    $ srun -N 1 /path_to/create_conda_env.sh

  4. 环境构建完成后，您可以通过指向共享卷上的环境路径来启动训练作业。您可以使用相同的设置启动单节点和多节点训练作业。要启动作业，请按如下方式创建任务启动器脚本（也称为入口点脚本）。

    #!/usr/bin/env bash
set -ex

ENV_PATH=/fsx/users/my_user/pytorch_env
TORCHRUN=$ENV_PATH/bin/torchrun
TRAINING_SCRIPT=/fsx/users/my_user/pt_train.py

WORLD_SIZE_JOB=$SLURM_NTASKS
RANK_NODE=$SLURM_NODEID
PROC_PER_NODE=8
MASTER_ADDR=(`scontrol show hostnames \$SLURM_JOB_NODELIST | head -n 1`)
MASTER_PORT=$(expr 10000 + $(echo -n $SLURM_JOBID | tail -c 4))

DIST_ARGS="--nproc_per_node=$PROC_PER_NODE \
           --nnodes=$WORLD_SIZE_JOB \
           --node_rank=$RANK_NODE \
           --master_addr=$MASTER_ADDR \
           --master_port=$MASTER_PORT \
          "
          
$TORCHRUN $DIST_ARGS $TRAINING_SCRIPT
    提示

    如果要使用的自动恢复功能提高训练作业抵御硬件故障的能力 SageMaker HyperPod，则需要在 entrypoint 脚本MASTER_ADDR中正确设置环境变量。要了解更多信息，请参阅自动恢复。

    本教程假设此脚本另存为/fsx/users/my_user/train.sh。

  5. 将此脚本位于共享卷中/fsx/users/my_user/train.sh，运行以下srun命令来安排 Slurm 作业。

    $ cd /fsx/users/my_user/
$ srun -N 8 train.sh