で Slurm コンピューティングノードで Docker コンテナを実行する HyperPod - Amazon SageMaker

翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

で Slurm コンピューティングノードで Docker コンテナを実行する HyperPod

で Slurm で Docker コンテナを実行するには SageMaker HyperPod、Enroot と Pyxis を使用する必要があります。Enroot パッケージは、Docker イメージを Slurm が理解できるランタイムに変換するのに役立ちます。Pyxis では、 srun コマンド を使用してランタイムを Slurm ジョブとしてスケジュールできますsrun --container-image=docker/image:tag。

ヒント

Docker、Enroot、Pyxis パッケージは、 で説明されているライフサイクルスクリプトの実行の一環として、クラスターの作成中にインストールする必要がありますが提供するベースライフサイクルスクリプトから開始する HyperPod。 HyperPod クラスターを作成するときは、 HyperPod サービスチームが提供するベースライフサイクルスクリプトを使用します。これらのベーススクリプトは、デフォルトでパッケージをインストールするように設定されています。config.py スクリプトには、 True () に設定されたパッケージをインストールするためのブール型パラメータを持つ Config クラスがありますenable_docker_enroot_pyxis=True。これは によって呼び出され、 lifecycle_script.py スクリプトで解析されます。このスクリプトは、 utilsフォルダから install_docker.shと install_enroot_pyxis.shスクリプトを呼び出します。インストールスクリプトは、パッケージの実際のインストールが行われる場所です。さらに、インストールスクリプトは、実行されるインスタンスからNVMeストアパスを検出できるかどうかを特定し、Docker と Enroot のルートパスを にセットアップします/opt/dlami/nvme。新しいインスタンスのデフォルトのルートボリュームは、100GB ボリューム/tmpでのみ にマウントされます。このEBSボリュームは、実行する予定のワークロードに のトレーニングが含まれている場合LLMs、つまり Docker コンテナのサイズが大きい場合に使い果たされます。ローカルNVMeストレージで P や G などのインスタンスファミリーを使用する場合は、 にアタッチされた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。このチュートリアルでは、提供されたサンプルスクリプトが、8 つのコンピューティングノードに合成データセットを持つ 700 億パラメータの Llama 2 モデルのモデル並列トレーニングジョブlaunch_training_enroot.shを起動します。トレーニングスクリプトのセットは で提供され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 Distributed Training リポジトリの「モデル並列処理ライブラリ、Docker および Enroot with Slurm を使用して SageMaker モデル並列トレーニングジョブを実行する」を参照してください。 GitHub の Slurm クラスターを使用した分散トレーニングの詳細については SageMaker HyperPod、「」の次のトピックを参照してくださいで Slurm を使用して分散トレーニングワークロードを実行する HyperPod。