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

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

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

若要在 上執行具有 Slurm 的 Docker 容器 SageMaker HyperPod,您需要使用 EnrootPyxis 。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.shinstall_enroot_pyxis.sh指令碼。安裝指令碼是套件實際安裝的位置。此外,安裝指令碼會識別是否可以從執行的執行個體中偵測NVMe儲存區路徑,並將 Docker 和 Enroot 的根路徑設定為 /opt/dlami/nvme。任何新執行個體的預設根磁碟區/tmp只會以 100GB EBS磁碟區掛載到 ,如果您計劃執行的工作負載涉及 和大型 Docker 容器的訓練LLMs,則會執行。如果您將 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

  1. 在運算節點上,請嘗試下列命令,以檢查 Docker 和 Enroot 是否已正確安裝。

    $ docker --help $ enroot --help
  2. 執行其中一個 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 映像,並在 上使用 Slurm 執行 SageMaker HyperPod。以下是範例使用案例,說明如何在 上使用 Docker 和 Slurm 執行訓練任務 SageMaker HyperPod。它顯示具有模型平行處理 (SMP) 程式庫的 Llama 2 模型 SageMaker 模型平行訓練的範例任務。

  1. 如果您想要使用由 SageMaker 或 分發的其中一個預先建置ECR映像DLC,請確定您授予 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": "*" } ] }
  2. 在運算節點上,複製儲存庫,然後前往提供使用 訓練範例指令碼的資料夾SMP。

    $ git clone https://github.com/aws-samples/awsome-distributed-training/ $ cd awsome-distributed-training/3.test_cases/17.SM-modelparallelv2
  3. 在本教學課程中,執行提取 Docker SMP 映像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啟動 700 億參數 Llama 2 模型的模型平行訓練任務,並在 8 個運算節點上搭配合成資料集。一組訓練指令碼提供於 3.test_cases/17.SM-modelparallelv2/scriptslaunch_training_enroot.sh並以 train_external.py作為進入點指令碼。

    重要

    若要在 上使用 Docker 容器 SageMaker HyperPod,您必須從主機機器將/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 分散式訓練儲存庫中的使用模型平行處理程式庫、Docker 和 Enroot with Slurm 執行 SageMaker 模型平行訓練任務 GitHub 如需在 上使用 Slurm 叢集進行分散式訓練的詳細資訊 SageMaker HyperPod,請前往 繼續進行下一個主題使用 Slurm on 執行分散式訓練工作負載 HyperPod