在 SageMaker HyperPod 叢集上排程 Slurm 任務 - Amazon SageMaker

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

在 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 叢集上執行任務，以使用全線式資料平行處理 （FSDP） 技術訓練 Llama-2。您也可以從 Awsome 分散式訓練 GitHub 儲存庫 找到更多範例。

提示

所有 SageMaker HyperPod 範例都可在 Awsome 分散式訓練 GitHub 儲存庫 的 3.test_cases 資料夾中取得。

  1. 複製 Awsome 分散式訓練 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