Pianificazione di un processo Slurm su un cluster SageMaker HyperPod - Amazon SageMaker AI

Pianificazione di un processo Slurm su un cluster SageMaker HyperPod

Puoi avviare job di addestramento utilizzando i comandi standard sbatch o srun di Slurm. Ad esempio, per avviare un job di addestramento a 8 nodi, puoi eseguire srun -N 8 --exclusive train.sh. SageMaker HyperPod supporta l’addestramento in una vasta gamma di ambienti, tra cui conda, venv, docker e enroot. Puoi configurare un ambiente di ML eseguendo script del ciclo di vita sui cluster SageMaker HyperPod. È inoltre possibile collegare un file system condiviso come Amazon FSx, che può essere utilizzato anche come ambiente virtuale.

L’esempio seguente mostra come eseguire un job per addestrare Llama-2 con la tecnica Fully Sharded Data Parallelism (FSDP) su un cluster SageMaker HyperPod con un file system condiviso Amazon FSx. Puoi trovare altri esempi anche nel repository GitHub Awsome Distributed Training.

Suggerimento

Tutti gli esempi di SageMaker HyperPod sono disponibili nella cartella 3.test_cases del repository GitHub Awsome Distributed Training.

  1. Clona il repository GitHub Awsome Distributed Training e copia i job di addestramento di esempio sul tuo file system Amazon FSx. 

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

  2. Eseguire lo script create_conda_env.sh. Questa operazione crea un ambiente conda sul file system Amazon FSx. Verifica che il file system sia accessibile a tutti i nodi del cluster.

  3. Crea l’ambiente virtuale Conda avviando un processo Slurm a nodo singolo come descritto di seguito.

    $ srun -N 1 /path_to/create_conda_env.sh

  4. Dopo aver creato l’ambiente, puoi avviare un job di addestramento indicando il percorso dell’ambiente sul volume condiviso. Puoi avviare job di addestramento sia a nodo singolo che multinodo con la stessa configurazione. Per avviare un processo, crea uno script di avvio dei processi (chiamato anche script del punto di ingresso) come segue.

    #!/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
    Suggerimento

    Per rendere il job di addestramento più resiliente ai guasti hardware utilizzando la funzionalità di ripresa automatica di SageMaker HyperPod, devi configurare correttamente la variabile di ambiente MASTER_ADDR nello script del punto di ingresso. Per ulteriori informazioni, consulta Ripresa automatica.

    Questo tutorial presuppone che questo script sia salvato come /fsx/users/my_user/train.sh.

  5. Con questo script nel volume condiviso in /fsx/users/my_user/train.sh, esegui il comando srun seguente per pianificare il processo Slurm.

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