Agendar uma tarefa do Slurm em um cluster do SageMaker HyperPod - Amazon SageMaker AI

Agendar uma tarefa do Slurm em um cluster do SageMaker HyperPod

Você pode iniciar trabalhos de treinamento usando os comandos padrão Slurm sbatch ou srun. Por exemplo, para iniciar um trabalho de treinamento de 8 nós, você pode executar o SageMaker srun -N 8 --exclusive train.shHyperPod compatível com treinamento em vários ambientes, incluindo conda, venv, docker, e enroot. Você pode configurar um ambiente de ML executando scripts de ciclo de vida em seus clusters do SageMaker HyperPod. Você também tem a opção de anexar um sistema de arquivos compartilhado, como o Amazon FSx, que também pode ser usado como um ambiente virtual.

O exemplo a seguir mostra como executar um trabalho para treinar o Llama-2 com a técnica Fully Sharded Data Parallelism (FSDP) em um cluster do SageMaker HyperPod com um sistema de arquivos compartilhado Amazon FSx. Você também pode encontrar mais exemplos no repositório do Awsome Distributed Training no GitHub.

dica

Todos os exemplos do SageMaker HyperPod estão disponíveis na pasta 3.test_cases do repositório do Awsome Distributed Training no GitHub.

  1. Clone o repositório GitHub do Awsome Distributed Training e copie os exemplos de trabalhos de treinamento para o seu sistema de arquivos Amazon FSx. 

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

  2. Execute o script create_conda_env.sh. Isso cria um ambiente conda no seu sistema de arquivos do Amazon FSx. Verifique se o sistema de arquivos está acessível a todos os nós do cluster.

  3. Crie o ambiente virtual Conda iniciando um trabalho de slurm de nó único da seguinte forma:

    $ srun -N 1 /path_to/create_conda_env.sh

  4. Depois que o ambiente for criado, você poderá iniciar um trabalho de treinamento apontando para o caminho do ambiente no volume compartilhado. Você pode iniciar trabalhos de treinamento de nó único e de vários nós com a mesma configuração. Para iniciar uma tarefa, crie um script inicializador de tarefas (também chamado de script de ponto de entrada) da seguinte forma:

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

    Se você quiser tornar seu trabalho de treinamento mais resiliente contra falhas de hardware usando o recurso de retomada automática do SageMaker HyperPod, você precisa configurar adequadamente a variável de ambiente no script do ponto de entrada MASTER_ADDR. Para saber mais, consulte Retomada automática.

    Este tutorial pressupõe que esse script seja salvo como /fsx/users/my_user/train.sh.

  5. Com esse script no volume compartilhado em /fsx/users/my_user/train.sh, execute o comando srun a seguir para agendar o trabalho do Slurm.

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