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

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

提示

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。任何全新執行個體的預設根磁碟區只會掛載到具有 100GB EBS 磁碟區的 /tmp，如果您計劃執行的工作負載涉及 LLM 的訓練以及大型 Docker 容器，則該磁碟區將耗盡。如果您使用 P 和 G 等執行個體系列搭配本機 NVMe 儲存體，則需要確保您使用 /opt/dlami/nvme 中連接的 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 執行訓練任務。它顯示如何使用 SageMaker AI 模型平行化 (SMP) 程式庫以模型平行化方式訓練 Llama 2 模型的範例任務。

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

在運算節點上，複製儲存庫並前往資料夾，其中提供使用 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/scripts 提供一組訓練指令碼，且 launch_training_enroot.sh 會將 train_external.py 當作進入點指令碼。

重要

若要在 SageMaker HyperPod 上使用 Docker 容器，您必須從主機電腦將 /var/log 目錄掛載到容器中的 /var/log 目錄，在此案例中，該主機電腦是 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 搭配 Slurm 執行模型平行化訓練任務。如需在 SageMaker HyperPod 上使用 Slurm 叢集進行分散式訓練的詳細資訊，請前往繼續進行在 HyperPod 上使用 Slurm 執行分散式訓練工作負載中的下一個主題。

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

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

文件慣用形式

排程 Slurm 任務

執行分散式訓練工作負載