在 Slurm 计算节点上运行 Docker 容器 SageMaker HyperPod

要在开启 Slurm 的情况下运行 Docker 容器 SageMaker HyperPod，你需要使用 En root 和 Pyxis。Enroot 软件包可帮助将 Docker 镜像转换为 Slurm 可以理解的运行时，而 Pyxis 则允许通过命令将运行时作为 Slurm 作业进行调度。srun srun --container-image=docker/image:tag

提示

应在集群创建期间安装 Docker、Enroot 和 Pyxis 软件包，作为运行生命周期脚本的一部分，如中所述。从提供的基本生命周期脚本开始 HyperPod创建集群时，请使用 HyperPod 服务团队提供的基本生命周期脚本。 HyperPod 默认情况下，这些基本脚本设置为安装软件包。在config.py脚本中，有一个用于安装软件包的布尔型参数设置为 True (enable_docker_enroot_pyxis=True) 的Config类。它由脚本调用并在脚本中解析，lifecycle_script.py脚本会从utils文件夹中调用install_docker.shinstall_enroot_pyxis.sh脚本。安装脚本是实际安装软件包的地方。此外，安装脚本还会确定它们能否从运行它们的实例中检测到 NVMe 存储路径，并设置 Docker 和 Enroot 的根路径。/opt/dlami/nvme任何新实例的默认根卷都/tmp只能装入 100GB 的 EBS 卷，如果您计划运行的工作负载涉及训练 LLM 以及大型 Docker 容器，则该卷将耗尽。如果您将实例系列（例如 P 和 G）与本地 NVMe 存储一起使用，则需要确保使用连接在的 NVMe 存储/opt/dlami/nvme，并且安装脚本负责配置过程。

检查根路径设置是否正确

在 Slurm 集群的计算节点上 SageMaker HyperPod，运行以下命令以确保生命周期脚本正常运行并且每个节点的根卷设置为。/opt/dlami/nvme/*以下命令显示了检查 Slurm 集群中 8 个计算节点的 Enroot 运行时路径和数据根路径的示例。


$ srun -N 8 cat /etc/enroot/enroot.conf | grep "ENROOT_RUNTIME_PATH"
ENROOT_RUNTIME_PATH        /opt/dlami/nvme/tmp/enroot/user-$(id -u)
... // The same or similar lines repeat 7 times


$ srun -N 8 cat /etc/docker/daemon.json
{
    "data-root": "/opt/dlami/nvme/docker/data-root"
}
... // The same or similar lines repeat 7 times

确认运行时路径已正确设置为后/opt/dlami/nvme/*，就可以使用 Enroot 和 Pyxis 构建和运行 Docker 容器了。

使用 Slurm 测试 Docker

在您的计算节点上，尝试使用以下命令来检查 Docker 和 Enroot 是否已正确安装。
```
$ docker --help
$ enroot --help
```

通过运行其中一个 NVIDIA CUDA U buntu 镜像来测试 Pyxis 和 Enroot 是否正确安装。


$ srun --container-image=nvidia/cuda:XX.Y.Z-base-ubuntuXX.YY nvidia-smi
pyxis: importing docker image: nvidia/cuda:XX.Y.Z-base-ubuntuXX.YY
pyxis: imported docker image: nvidia/cuda:XX.Y.Z-base-ubuntuXX.YY
DAY MMM DD HH:MM:SS YYYY
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.141.03   Driver Version: 470.141.03   CUDA Version: XX.YY    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:1E.0 Off |                    0 |
| N/A   40C    P0    27W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

您也可以通过创建脚本并运行sbatch命令来对其进行测试，如下所示。


$ cat <<EOF >> container-test.sh
#!/bin/bash
#SBATCH --container-image=nvidia/cuda:XX.Y.Z-base-ubuntuXX.YY
nvidia-smi
EOF

$ sbatch container-test.sh
pyxis: importing docker image: nvidia/cuda:XX.Y.Z-base-ubuntuXX.YY
pyxis: imported docker image: nvidia/cuda:XX.Y.Z-base-ubuntuXX.YY
DAY MMM DD HH:MM:SS YYYY
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.141.03   Driver Version: 470.141.03   CUDA Version: XX.YY    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:1E.0 Off |                    0 |
| N/A   40C    P0    27W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

使用 Docker 运行测试 Slurm 作业

使用 Docker 完成 Slurm 的设置后，你可以带上任何预先构建的 Docker 镜像，然后使用 Slurm 开启运行。 SageMaker HyperPod以下是一个示例用例，它将引导您了解如何使用 Docker 和 Slurm 运行训练作业。 SageMaker HyperPod它显示了使用模型并行度 (SMP) 库对 Llama 2 模型进行 SageMaker 模型并行训练的示例作业。

如果您想使用 SageMaker 或 DLC 分发的预构建 ECR 映像之一，请确保授予 HyperPod 集群通过提取 ECR 映像的权限。IAM 角色适用于 SageMaker HyperPod如果您使用自己的镜像或开源 Docker 镜像，则可以跳过此步骤。将以下权限添加到IAM 角色适用于 SageMaker HyperPod。在本教程中，我们将使用与 SMP 库一起预先打包的 SMP Docker 镜像。


{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "ecr:BatchCheckLayerAvailability",
                "ecr:BatchGetImage",
                "ecr-public:*",
                "ecr:GetDownloadUrlForLayer",
                "ecr:GetAuthorizationToken",
                "sts:*"
            ],
            "Resource": "*"
        }
    ]
}

在计算节点上，克隆存储库并转到提供使用 SMP 进行训练的示例脚本的文件夹。


$ git clone https://github.com/aws-samples/awsome-distributed-training/
$ cd awsome-distributed-training/3.test_cases/17.SM-modelparallelv2

在本教程中，运行提取 SMP Docker 镜像的示例脚本 docker_build.sh，构建 Docker 容器，然后将其作为 Enroot 运行时运行。您可以根据需要对其进行修改。


$ cat docker_build.sh
#!/usr/bin/env bash

region=us-west-2
dlc_account_id=658645717510
aws ecr get-login-password --region $region | docker login --username AWS --password-stdin $dlc_account_id.dkr.ecr.$region.amazonaws.com

docker build -t smpv2 .
enroot import -o smpv2.sqsh  dockerd://smpv2:latest


$ bash docker_build.sh

使用创建用于启动训练作业的批处理脚本sbatch。在本教程中，提供的示例脚本launch_training_enroot.sh启动了 700 亿参数的 Llama 2 模型的模型并行训练作业，该模型在 8 个计算节点上使用合成数据集。提供了一组训练脚本 3.test_cases/17.SM-modelparallelv2/scripts，这些脚本train_external.py作为launch_training_enroot.sh入口点脚本。

重要

要在上使用 Docker 容器 SageMaker HyperPod，必须将主机（在本例中为 HyperPod 计算节点）中的/var/log目录挂载到容器中的/var/log目录上。您可以通过为 Enroot 添加以下变量来进行设置。


"${HYPERPOD_PATH:="/var/log/aws/clusters":"/var/log/aws/clusters"}"


$ cat launch_training_enroot.sh
#!/bin/bash

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

#SBATCH --nodes=8 # number of nodes to use, 2 p4d(e) = 16 A100 GPUs
#SBATCH --job-name=smpv2_llama # name of your job
#SBATCH --exclusive # job has exclusive use of the resource, no sharing
#SBATCH --wait-all-nodes=1

set -ex;

###########################
###### User Variables #####
###########################

#########################
model_type=llama_v2
model_size=70b

# Toggle this to use synthetic data
use_synthetic_data=1


# To run training on your own data  set Training/Test Data path  -> Change this to the tokenized dataset path in Fsx. Acceptable formats are huggingface (arrow) and Jsonlines.
# Also change the use_synthetic_data to 0

export TRAINING_DIR=/fsx/path_to_data
export TEST_DIR=/fsx/path_to_data
export CHECKPOINT_DIR=$(pwd)/checkpoints

# Variables for Enroot
: "${IMAGE:=$(pwd)/smpv2.sqsh}"
: "${HYPERPOD_PATH:="/var/log/aws/clusters":"/var/log/aws/clusters"}" # This is needed for validating its hyperpod cluster
: "${TRAIN_DATA_PATH:=$TRAINING_DIR:$TRAINING_DIR}"
: "${TEST_DATA_PATH:=$TEST_DIR:$TEST_DIR}"
: "${CHECKPOINT_PATH:=$CHECKPOINT_DIR:$CHECKPOINT_DIR}"   


###########################
## Environment Variables ##
###########################

#export NCCL_SOCKET_IFNAME=en
export NCCL_ASYNC_ERROR_HANDLING=1

export NCCL_PROTO="simple"
export NCCL_SOCKET_IFNAME="^lo,docker"
export RDMAV_FORK_SAFE=1
export FI_EFA_USE_DEVICE_RDMA=1
export NCCL_DEBUG_SUBSYS=off
export NCCL_DEBUG="INFO"
export SM_NUM_GPUS=8
export GPU_NUM_DEVICES=8
export FI_EFA_SET_CUDA_SYNC_MEMOPS=0

# async runtime error ...
export CUDA_DEVICE_MAX_CONNECTIONS=1


#########################
## Command and Options ##
#########################

if [ "$model_size" == "7b" ]; then
    HIDDEN_WIDTH=4096
    NUM_LAYERS=32
    NUM_HEADS=32
    LLAMA_INTERMEDIATE_SIZE=11008
    DEFAULT_SHARD_DEGREE=8
# More Llama model size options
elif [ "$model_size" == "70b" ]; then
    HIDDEN_WIDTH=8192
    NUM_LAYERS=80
    NUM_HEADS=64
    LLAMA_INTERMEDIATE_SIZE=28672
    # Reduce for better perf on p4de
    DEFAULT_SHARD_DEGREE=64
fi


if [ -z "$shard_degree" ]; then
    SHARD_DEGREE=$DEFAULT_SHARD_DEGREE
else
    SHARD_DEGREE=$shard_degree
fi

if [ -z "$LLAMA_INTERMEDIATE_SIZE" ]; then
    LLAMA_ARGS=""
else
    LLAMA_ARGS="--llama_intermediate_size $LLAMA_INTERMEDIATE_SIZE "
fi


if [ $use_synthetic_data == 1 ]; then
    echo "using synthetic data"
    declare -a ARGS=(
    --container-image $IMAGE
    --container-mounts $HYPERPOD_PATH,$CHECKPOINT_PATH
    )
else
    echo "using real data...."
    declare -a ARGS=(
    --container-image $IMAGE
    --container-mounts $HYPERPOD_PATH,$TRAIN_DATA_PATH,$TEST_DATA_PATH,$CHECKPOINT_PATH
    )
fi


declare -a TORCHRUN_ARGS=(
    # change this to match the number of gpus per node:
    --nproc_per_node=8 \
    --nnodes=$SLURM_JOB_NUM_NODES \
    --rdzv_id=$SLURM_JOB_ID \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$(hostname) \
)

srun -l "${ARGS[@]}" torchrun "${TORCHRUN_ARGS[@]}" /path_to/train_external.py \
            --train_batch_size 4 \
            --max_steps 100 \
            --hidden_width $HIDDEN_WIDTH \
            --num_layers $NUM_LAYERS \
            --num_heads $NUM_HEADS \
            ${LLAMA_ARGS} \
            --shard_degree $SHARD_DEGREE \
            --model_type $model_type \
            --profile_nsys 1 \
            --use_smp_implementation 1 \
            --max_context_width 4096 \
            --tensor_parallel_degree 1 \
            --use_synthetic_data $use_synthetic_data \
            --training_dir $TRAINING_DIR \
            --test_dir $TEST_DIR \
            --dataset_type hf \
            --checkpoint_dir $CHECKPOINT_DIR \
            --checkpoint_freq 100 \

$ sbatch launch_training_enroot.sh

要查找可下载的代码示例，请参阅在 Awsome 分布式训练存储库中使用模型并行度库运行 SageMaker 模型并行训练作业、Docker 和 Enroot with Slurm。 GitHub 有关在 Slurm 集群开启的情况下进行分布式训练的更多信息 SageMaker HyperPod，请继续阅读下一个主题。在 Slurm 开启的情况下运行分布式训练工作负载 SageMaker HyperPod

Javascript 在您的浏览器中被禁用或不可用。

要使用 Amazon Web Services 文档，必须启用 Javascript。请参阅浏览器的帮助页面以了解相关说明。

在集群上安排 Slurm 作业 SageMaker HyperPod

在 Slurm 开启的情况下运行分布式训练工作负载 SageMaker HyperPod