Tutorial sobre PEFT-LoRa com clusters do Slurm no HyperPod (GPU) - Amazon SageMaker AI
Configurar o ambiente de GPU do Slurm no HyperPodIniciar a tarefa de treinamento

Tutorial sobre PEFT-LoRa com clusters do Slurm no HyperPod (GPU)

O tutorial a seguir configura o ambiente do Slurm e inicia um trabalho de ajuste fino com eficiência de parâmetros (PEFT) em um modelo Llama com 8 bilhões de parâmetros.

Pré-requisitos

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

  • Configurar cluster de GPU do Slurm no HyperPod

    • O cluster do Slurm no HyperPod deve ter o NVIDIA Enroot e o Pyxis habilitados (eles estão habilitados por padrão).

  • 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) Se você precisar dos pesos pré-treinados do HuggingFace ou se estiver treinando um modelo Llama 3.2, deverá obter o token do HuggingFace antes de começar a treinar. Para ter mais informações sobre como obter o token, consulte User access tokens.

Configurar o ambiente de GPU do Slurm no HyperPod

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. Você deve usar o Python 3.9 ou posterior.

    #set up a virtual environment
python3 -m venv ${PWD}/venv
source venv/bin/activate

  • Clone as fórmulas do SageMaker HyperPod e os repositórios do adaptador do SageMaker HyperPod em um local de armazenamento compartilhado. 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 https://github.com/aws/sagemaker-hyperpod-training-adapter-for-nemo.git
git clone --recursive https://github.com/aws/sagemaker-hyperpod-recipes.git
cd sagemaker-hyperpod-recipes
pip3 install -r requirements.txt

  • Crie um arquivo squash usando o Enroot. Para encontrar a versão mais recente do contêiner da SMP, consulteNotas de lançamento da biblioteca de paralelismo de modelos do SageMaker. Para ter mais informações sobre como usar o arquivo Enroot, consulte Build AWS-optimized Nemo-Launcher image.

    REGION="<region>"
IMAGE="658645717510.dkr.ecr.${REGION}.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311-cu121"
aws ecr get-login-password --region ${REGION} | docker login --username AWS --password-stdin 658645717510.dkr.ecr.${REGION}.amazonaws.com
enroot import -o $PWD/smdistributed-modelparallel.sqsh dockerd://${IMAGE}
mv $PWD/smdistributed-modelparallel.sqsh "/fsx/<any-path-in-the-shared-filesystem>"

  • Para usar o arquivo squash do Enroot para iniciar o treinamento, use o exemplo a seguir para modificar o arquivo recipes_collection/config.yaml.

    container: /fsx/path/to/your/smdistributed-modelparallel.sqsh

Iniciar a tarefa de treinamento

Para iniciar uma tarefa de PEFT para o modelo Llama com 8 bilhões de parâmetros com um comprimento de sequência de 8.192 em um único nó de computação do Slurm, defina o script de execução launcher_scripts/llama/run_hf_llama3_8b_seq8k_gpu_lora.sh como o seguinte:

  • IMAGE: o contêiner da seção de configuração do ambiente.

  • HF_MODEL_NAME_OR_PATH: defina o nome ou o caminho dos pesos pré-treinados no parâmetro hf_model_name_or_path da fórmula.

  • (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=${HF_ACCESS_TOKEN}
#!/bin/bash
IMAGE="${YOUR_IMAGE}"
SAGEMAKER_TRAINING_LAUNCHER_DIR="${SAGEMAKER_TRAINING_LAUNCHER_DIR:-${PWD}}"

TRAIN_DIR="${YOUR_TRAIN_DIR}" # Location of training dataset
VAL_DIR="${YOUR_VAL_DIR}" # Location of validation dataset

# experiment output directory
EXP_DIR="${YOUR_EXP_DIR}"
HF_ACCESS_TOKEN="${YOUR_HF_TOKEN}"
HF_MODEL_NAME_OR_PATH="${YOUR_HF_MODEL_NAME_OR_PATH}"

# Add hf_model_name_or_path and turn off synthetic_data
HYDRA_FULL_ERROR=1 python3 ${SAGEMAKER_TRAINING_LAUNCHER_DIR}/main.py \
    recipes=fine-tuning/llama/hf_llama3_8b_seq8k_gpu_lora \
    base_results_dir=${SAGEMAKER_TRAINING_LAUNCHER_DIR}/results \
    recipes.run.name="hf_llama3_lora" \
    recipes.exp_manager.exp_dir="$EXP_DIR" \
    recipes.model.data.train_dir="$TRAIN_DIR" \
    recipes.model.data.val_dir="$VAL_DIR" \
    recipes.model.hf_model_name_or_path="$HF_MODEL_NAME_OR_PATH" \
    container="${IMAGE}" \
    +cluster.container_mounts.0="/fsx:/fsx" \
    recipes.model.hf_access_token="${HF_ACCESS_TOKEN}"

Depois de configurar todos os parâmetros necessários no script anterior, você pode começar a executar a tarefa de treinamento.

bash launcher_scripts/llama/run_hf_llama3_8b_seq8k_gpu_lora.sh

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