在 Slurm 計算節點上運行碼頭容器 SageMaker HyperPod - Amazon SageMaker

本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。

在 Slurm 計算節點上運行碼頭容器 SageMaker HyperPod

要使用 Slurm 運行碼頭容器 SageMaker HyperPod，您需要使用恩魯特和 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.sh和install_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 容器了。

用泥漿測試碼頭工人

  1. 在您的計算節點上，嘗試以下命令來檢查 Docker 和 Enroot 是否已正確安裝。

    $ docker --help
$ enroot --help

  2. 通過運行 NVIDIA CU DA Ubuntu 映像之一測試 Pyxis 和恩魯特正確安裝。

    $ 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它顯示了與模型 parallel 性（SMP）庫的駱駝 2 模型的模型平行培訓的示例工作。 SageMaker

  1. 如果您想要使用其中一個由 SageMaker 或 DLC 發佈的預先建置 ECR 映像檔，請確定您授與 HyperPod 叢集的權限，以便透過. 適用於的 IAM 角色 SageMaker HyperPod 如果您使用自己的圖像或開源 Docker 映像，則可以跳過此步驟。將下列權限新增至適用於的 IAM 角色 SageMaker HyperPod。在本教學課程中，我們會使用預先封裝在 SMP 程式庫中的 SMP 泊塢視窗映像檔。

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

  2. 在計算節點上，複製存放庫，然後移至提供 SMP 訓練範例指令碼的資料夾。

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

  3. 在此教學課程中，請執行提取 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

  4. 使用建立批次指令碼以啟動訓練工作sbatch。在本教學課程中，提供的範例指令碼會launch_training_enroot.sh啟動 70 億個參數 Lamama 2 模型的模型 parallel 訓練工作，並在 8 個運算節點上使用合成資料集。在提供了一組培訓腳本 3.test_cases/17.SM-modelparallelv2/scripts，並launch_training_enroot.sh作train_external.py為入口點腳本。

    重要

    若要在上使用 Docker 容器 SageMaker HyperPod，您必須將/var/log目錄從主機 (在本例中為 HyperPod 計算節點) 掛載到容器中的/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 模型 parallel 處理原則程式庫、Docker 和具有 Slurm 的 Enroot 執行模型平行訓練工作。 GitHub 如需有關上使用 Slurm 叢集進行分散式訓練的詳細資訊 SageMaker HyperPod，請繼續執行下一個主題，網址為。使用 Slurm on 執行分散式訓練工作負載 SageMaker HyperPod