Tutorial de pré-treinamento de clusters do Slurm no Trainium - Amazon SageMaker AI
Configurar o ambiente Trainium no cluster do SlurmIniciar a tarefa de treinamento no Trainium

Tutorial de pré-treinamento de clusters do Slurm no Trainium

O tutorial a seguir configura um ambiente Trainium em um cluster do Slurm e inicia uma tarefa de treinamento em um modelo Llama de 8 bilhões de parâmetros.

Pré-requisitos

Antes de começar a configurar seu ambiente, você deve:

  • Configurar um cluster do Slurm no Trainium para o SageMaker HyperPod.

  • Ter um local de armazenamento compartilhado. Pode ser um sistema de arquivos do Amazon FSx ou um sistema do NFS acessível por meio dos nós do cluster.

  • Ter dados em um dos seguintes formatos:

    • JSON

    • JSONGZ (JSON compactado)

    • ARROW

  • (Opcional) Você deve obter um token do HuggingFace se estiver usando os pesos do modelo do HuggingFace para pré-treinamento ou ajuste fino. Para ter mais informações sobre como obter o token, consulte User access tokens.

Configurar o ambiente Trainium no cluster do Slurm

Para iniciar uma tarefa de treinamento em um cluster do Slurm, faça o seguinte:

  • Use SSH no nó cabeça do cluster.

  • Depois de fazer login, configure o ambiente Neuron. Para ter informações sobre como configurar o Neuron, consulte as etapas de configuração do Neuron. Recomendamos contar com as AMIs de deep learning que vêm pré-instaladas com os drivers do Neuron, como o Ubuntu 20 com DLAMI do Pytorch.

  • Clone o repositório de fórmulas do SageMaker HyperPod em um local de armazenamento compartilhado no cluster. O local de armazenamento compartilhado pode ser um sistema de arquivos do Amazon FSx ou um sistema do NFS acessível por meio dos nós do cluster.

    git clone --recursive https://github.com/aws/sagemaker-hyperpod-recipes.git
cd sagemaker-hyperpod-recipes
pip3 install -r requirements.txt

  • Siga este tutorial:: HuggingFace Llama3-8B Pretraining.

  • Prepare uma configuração do modelo. As configurações do modelo estão disponíveis no repositório do Neuron. Com relação à configuração do modelo usada neste tutorial, consulte a configuração do modelo llama3 8b.

Iniciar a tarefa de treinamento no Trainium

Para iniciar uma tarefa de treinamento no Trainium, especifique uma configuração de cluster e uma fórmula do Neuron. Por exemplo, para iniciar uma tarefa de pré-treinamento do llama3 8b no Trainium, defina o script de execução launcher_scripts/llama/run_hf_llama3_8b_seq8k_trn1x4_pretrain.sh como o seguinte:

  • MODEL_CONFIG: a configuração do modelo na seção de configuração do ambiente.

  • (Opcional) Você pode fornecer o token do HuggingFace se precisar de pesos pré-treinados do HuggingFace definindo o seguinte par de chave-valor:

    recipes.model.hf_access_token=<your_hf_token>
#!/bin/bash

#Users should set up their cluster type in /recipes_collection/config.yaml

SAGEMAKER_TRAINING_LAUNCHER_DIR=${SAGEMAKER_TRAINING_LAUNCHER_DIR:-"$(pwd)"}

COMPILE=0
TRAIN_DIR="${TRAIN_DIR}" # Location of training dataset
MODEL_CONFIG="${MODEL_CONFIG}" # Location of config.json for the model

HYDRA_FULL_ERROR=1 python3 "${SAGEMAKER_TRAINING_LAUNCHER_DIR}/main.py" \
    base_results_dir="${SAGEMAKER_TRAINING_LAUNCHER_DIR}/results" \
    instance_type="trn1.32xlarge" \
    recipes.run.compile="$COMPILE" \
    recipes.run.name="hf-llama3-8b" \
    recipes.trainer.num_nodes=4 \
    recipes=training/llama/hf_llama3_8b_seq8k_trn1x4_pretrain \
    recipes.data.train_dir="$TRAIN_DIR" \
    recipes.model.model_config="$MODEL_CONFIG"

Para iniciar a tarefa de treinamento, execute o seguinte comando:

bash launcher_scripts/llama/run_hf_llama3_8b_seq8k_trn1x4_pretrain.sh

Para ter mais informações sobre como configurar um cluster do Slurm, consulte Executar uma tarefa de treinamento do Slurm no HyperPod.