Programación de un trabajo de Slurm en un clúster de SageMaker HyperPod - Amazon SageMaker AI

Programación de un trabajo de Slurm en un clúster de SageMaker HyperPod

Puede iniciar trabajos de entrenamiento utilizando los comandos sbatch o srun estándar de Slurm. Por ejemplo, para iniciar un trabajo de entrenamiento de 8 nodos, puede ejecutar srun -N 8 --exclusive train.sh. SageMaker HyperPod admite el entrenamiento en una variedad de entornos, incluidos conda, venv, docker y enroot. Puede configurar un entorno de ML ejecutando scripts de ciclo de vida en los clústeres de SageMaker HyperPod. También tiene la opción de asociar un sistema de archivos compartidos, como Amazon FSx, que también se puede utilizar como entorno virtual.

En el siguiente ejemplo, se muestra cómo ejecutar un trabajo para entrenar a Llama-2 con la técnica de Paralelismo de datos totalmente particionados (FSDP) en un clúster de SageMaker HyperPod con un sistema de archivos compartidos de Amazon FSx. También puede encontrar más ejemplos en el repositorio de GitHub de Awsome Distributed Training.

sugerencia

Todos los ejemplos de SageMaker HyperPod están disponibles en la carpeta 3.test_cases del repositorio de GitHub de Awsome Distributed Training.

  1. Clone el repositorio de GitHub de Awsome Distributed Training y copie los ejemplos de trabajos de entrenamiento en el sistema de archivos de Amazon FSx. 

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

  2. Ejecute el script create_conda_env.sh. Con ello se crea un entorno conda en el sistema de archivos de Amazon FSx. Asegúrese de que todos los nodos del clúster puedan acceder al sistema de archivos.

  3. Cree el entorno virtual de Conda iniciando un trabajo de Slurm de un solo nodo de la siguiente manera.

    $ srun -N 1 /path_to/create_conda_env.sh

  4. Una vez creado el entorno, puede iniciar un trabajo de entrenamiento apuntando a la ruta del entorno en el volumen compartido. Puede iniciar trabajos de entrenamiento de un solo nodo y de varios nodos con la misma configuración. Para iniciar un trabajo, cree un script de inicio de trabajos (también llamado script de punto de entrada) de la siguiente manera.

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

    Si desea que su trabajo de entrenamiento sea más resiliente a los fallos de hardware mediante el uso de la capacidad de reanudación automática de SageMaker HyperPod, debe configurar correctamente la variable de entorno MASTER_ADDR en el script de punto de entrada. Para obtener más información, consulte Reanudación automática.

    En este tutorial, se presupone que este script se guarda como /fsx/users/my_user/train.sh.

  5. Con este script en el volumen compartido en /fsx/users/my_user/train.sh, ejecute el siguiente comando srun para programar el trabajo de Slurm.

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