在 HyperPod 上的 Slurm 運算節點上執行 Docker 容器

若要在 SageMaker HyperPod 上執行具有 Slurm 的 Docker 容器，您需要使用 Enroot 和 Pyxis。 HyperPod Enroot 套件有助於將 Docker 映像轉換為 Slurm 可以了解的執行時間，而 Pyxis 則可透過 srun命令將執行時間排程為 Slurm 任務srun --container-image=docker/image:tag。

提示

在叢集建立期間，應安裝 Docker、Enroot 和 Pyxis 套件，做為執行生命週期指令碼的一部分，如中的指引從 HyperPod 提供的基本生命週期指令碼開始。建立 HyperPod 叢集時，請使用 HyperPod 服務團隊提供的基本生命週期指令碼。這些基本指令碼會設定為預設安裝套件。在config.py指令碼中，有 Config類別，其中包含布林值類型參數，用於安裝設定為 True() 的套件enable_docker_enroot_pyxis=True。這由呼叫，並在lifecycle_script.py指令碼中剖析，從 utils 資料夾呼叫 install_docker.sh和 install_enroot_pyxis.sh指令碼。安裝指令碼是套件實際安裝的位置。此外，安裝指令碼會識別是否可以從執行的執行個體偵測 NVMe 存放區路徑，並設定 Docker 和 Enroot 的根路徑至 /opt/dlami/nvme。任何新執行個體的預設根磁碟區/tmp只會以 100GB EBS 磁碟區掛載到，如果您計劃執行的工作負載涉及 LLMs訓練，因此大型 Docker 容器就會執行。如果您使用 P 和 G 等執行個體系列搭配本機 NVMe 儲存體，則需要確保使用連接到的 NVMe 儲存體/opt/dlami/nvme，且安裝指令碼會處理組態程序。

檢查根路徑是否已正確設定

在 SageMaker HyperPod 上 Slurm 叢集的運算節點上，執行下列命令，以確保生命週期指令碼正常運作，且每個節點的根磁碟區設定為 /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 Ubuntu 映像，測試 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 映像，並在 SageMaker HyperPod 上使用 Slurm 執行。以下是範例使用案例，引導您如何在 SageMaker HyperPod 上使用 Docker 和 Slurm 執行訓練任務。 HyperPod 它顯示使用 SageMaker AI 模型平行處理 (SMP) 程式庫的 Llama 2 模型模型模型平行訓練的範例任務。

如果您想要使用 SageMaker AI 或 DLC 分發的其中一個預先建置 ECR 映像，請確定您授予 HyperPod 叢集透過提取 ECR 映像的許可SageMaker HyperPod 的 IAM 角色。如果您使用自己的或開放原始碼 Docker 映像，您可以略過此步驟。將下列許可新增至 SageMaker HyperPod 的 IAM 角色。在本教學課程中，我們使用預先封裝的 SMP Docker 映像與 SMP 程式庫。


{
    "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

在本教學課程中，請執行範例指令碼docker_build.sh，以提取 SMP Docker 映像、建置 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/scriptslaunch_training_enroot.sh並以 train_external.py做為進入點指令碼。

重要

若要在 SageMaker HyperPod 上使用 Docker 容器，您必須將 /var/log目錄從主機機器掛載到容器的/var/log目錄上，此機器是 HyperPod 運算節點。 HyperPod 您可以為 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 分散式訓練 GitHub 儲存庫中的使用 SageMaker AI 模型平行處理程式庫、Docker 和 Enroot with Slurm 執行模型平行訓練任務。 GitHub 如需在 SageMaker HyperPod 上使用 Slurm 叢集進行分散式訓練的詳細資訊，請前往繼續進行下一個主題在 HyperPod 上使用 Slurm 執行分散式訓練工作負載。

您的瀏覽器已停用或無法使用 Javascript。

您必須啟用 Javascript，才能使用 AWS 文件。請參閱您的瀏覽器說明頁以取得說明。

文件慣用形式

在 SageMaker HyperPod 叢集上排程 Slurm 任務

在 HyperPod 上使用 Slurm 執行分散式訓練工作負載