SageMaker HyperPod クラスターでの Slurm ジョブのスケジューリング - Amazon SageMaker AI

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

SageMaker HyperPod クラスターでの Slurm ジョブのスケジューリング

Slurm 標準の sbatch コマンドまたは srun コマンドを使用してトレーニングジョブを起動することができます。例えば、8 ノードトレーニングジョブを起動するには、srun -N 8 --exclusive train.sh SageMaker HyperPod を実行すると、conda、venv、docker、enroot などのさまざまな環境でのトレーニングがサポートされます。SageMaker HyperPod クラスターでライフサイクルスクリプトを実行することにより、ML 環境を設定できます。さらに、仮想環境としても使用できる、Amazon FSx などの共有ファイルシステムをアタッチするオプションもあります。

次の例は、Amazon FSx 共有ファイルシステムを持つ SageMaker HyperPod クラスターで Fully Sharded Data Parallelism (FSDP) 手法を使用して Llama-2 をトレーニングするためのジョブを実行する方法を示しています。Awsome Distributed Training GitHub リポジトリでは、その他の例も確認できます。

ヒント

すべての SageMaker HyperPod の例は、Awsome Distributed Training GitHub リポジトリの 3.test_cases フォルダにあります。

  1. Awsome Distributed Training GitHub リポジトリをクローンし、トレーニングジョブのサンプルを Amazon FSx ファイルシステムにコピーします。

    $ TRAINING_DIR=/fsx/users/my-user/fsdp
$ git clone https://github.com/aws-samples/awsome-distributed-training/

  2. create_conda_env.sh スクリプトを実行します。これにより、Amazon FSx ファイルシステムに conda 環境が作成されます。クラスター内のすべてのノードからファイルシステムにアクセスできることを確認してください。

  3. 次のように、単一ノードの Slurm ジョブを起動することにより、仮想 Conda 環境を構築します。

    $ srun -N 1 /path_to/create_conda_env.sh

  4. 環境を構築したら、共有ボリュームの環境パスを指してトレーニングジョブを起動できます。同じセットアップで、単一ノードトレーニングジョブとマルチノードトレーニングジョブの両方を起動することができます。ジョブを起動するには、次のようにジョブランチャースクリプト (エントリポイントスクリプトとも呼ばれます) を作成します。

    #!/usr/bin/env bash
set -ex

ENV_PATH=/fsx/users/my_user/pytorch_env
TORCHRUN=$ENV_PATH/bin/torchrun
TRAINING_SCRIPT=/fsx/users/my_user/pt_train.py

WORLD_SIZE_JOB=$SLURM_NTASKS
RANK_NODE=$SLURM_NODEID
PROC_PER_NODE=8
MASTER_ADDR=(`scontrol show hostnames \$SLURM_JOB_NODELIST | head -n 1`)
MASTER_PORT=$(expr 10000 + $(echo -n $SLURM_JOBID | tail -c 4))

DIST_ARGS="--nproc_per_node=$PROC_PER_NODE \
           --nnodes=$WORLD_SIZE_JOB \
           --node_rank=$RANK_NODE \
           --master_addr=$MASTER_ADDR \
           --master_port=$MASTER_PORT \
          "
          
$TORCHRUN $DIST_ARGS $TRAINING_SCRIPT
    ヒント

    SageMaker HyperPod の自動再開機能を使用することにより、ハードウェア障害に対するトレーニングジョブの回復性を高めるには、エントリポイントスクリプトで環境変数 MASTER_ADDR を適切に設定する必要があります。詳細については自動再開を参照してください。

    このチュートリアルでは、このスクリプトが /fsx/users/my_user/train.sh という名前で保存されていることを前提としています。

  5. このスクリプトが /fsx/users/my_user/train.sh の共有ボリュームに保存されている状態で、次の srun コマンドを実行して Slurm ジョブをスケジュールします。

    $ cd /fsx/users/my_user/
$ srun -N 8 train.sh