Esegui l'addestramento di reinforcement learning su una VM TPU utilizzando MaxText

Questo tutorial fornisce una guida passo passo per eseguire l'addestramento di reinforcement learning (RL) su una singola v6e-8 istanza di macchina virtuale (VM) di Tensor Processing Unit (TPU) Google Cloud utilizzando MaxText, uno stack di addestramento basato su JAX ad alte prestazioni per modelli linguistici di grandi dimensioni (LLM).

Riepilogo del carico di lavoro

L'esecuzione di questo carico di lavoro di post-addestramento RL su Cloud TPU comporta i seguenti componenti:

  • Modello: Llama 3.1 8B Instruct (llama3.1-8b-Instruct), un LLM con pesi aperti ottimizzato per le attività di dialogo e di follow-up delle istruzioni. In questo tutorial, i pesi del modello vengono convertiti dal formato Hugging Face al formato compatibile con MaxText (mt-format) e preparati per l'esecuzione distribuita su una TPU v6e-8 utilizzando JAX.

  • Algoritmo di reinforcement learning (RL): Group Relative Policy Optimization (GRPO), in cui il carico di lavoro esegue una pipeline di post-addestramento scalabile utilizzando il trainer RL integrato di MaxText. GRPO elimina la necessità di un modello di valore o critico separato valutando un gruppo di output generati per lo stesso prompt, utilizzando ricompense relative per ottimizzare in modo efficiente la policy dell'attore.

  • Set di dati e modello: Grade School Math 8K (GSM8K), un set di dati di alta qualità composto da problemi di matematica di scuola elementare linguisticamente diversi. Questo carico di lavoro lo utilizza per generare rollout, calcolare le ricompense e migliorare le funzionalità di ragionamento multi-step durante la fase di allineamento RL.

  • Infrastruttura e framework: un singolo nodo Cloud TPU v6e-8 (ct6e-standard-8t) che utilizza MaxText (maxtext[tpu-post-train]), uno stack di addestramento ad alte prestazioni basato su JAX.

Obiettivi

  • Configurare un'istanza VM Cloud TPU.
  • Installare MaxText e le relative dipendenze.
  • Convertire un modello Hugging Face in formato MaxText.
  • Eseguire un carico di lavoro di ottimizzazione della policy relativa al gruppo RL (GRPO) sulla TPU.
  • Convertire di nuovo il modello addestrato in formato Hugging Face per la pubblicazione.

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili di Google Cloud:

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il calcolatore prezzi.

I nuovi Google Cloud utenti potrebbero avere diritto a una prova senza costi.

Al termine delle attività descritte in questo documento, puoi evitare l'addebito di ulteriori costi eliminando le risorse che hai creato. Per saperne di più, consulta Esegui la pulizia.

Prima di iniziare

  • Per utilizzare questo tutorial, devi disporre di un token di accesso a Hugging Face. Puoi registrarti per un account senza costi su Hugging Face. Dopo aver creato un account, genera un token di accesso:

    1. Nella pagina Benvenuto in Hugging Face, fai clic sull'avatar del tuo account e seleziona Token di accesso.
    2. Nella pagina Token di accesso, fai clic su Crea nuovo token.
    3. Seleziona il tipo di token Lettura e inserisci un nome per il token.
    4. Viene visualizzato il token di accesso. Salva il token in un luogo sicuro.

  • Sul sito web di Hugging Face, accetta il contratto di licenza per il modello che intendi addestrare. Questo tutorial utilizza il modello llama3.1-8b-Instruct.

Per ottenere le autorizzazioni necessarie per completare questo tutorial, chiedi all'amministratore di concederti i seguenti ruoli IAM nel tuo progetto:

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Configura l'ambiente

Imposta le variabili di ambiente eseguendo il seguente script:

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export TPU_NAME="YOUR_TPU_NAME"

Sostituisci quanto segue:

  • YOUR_PROJECT_ID: l' Google Cloud ID progetto
  • YOUR_ZONE: la zona che vuoi utilizzare
  • YOUR_RESERVATION_NAME: la prenotazione della capacità
  • YOUR_TPU_NAME: il nome dell'istanza VM Cloud TPU

Esegui l'autenticazione con Google Cloud eseguendo il seguente comando:

gcloud auth login

Crea la VM Cloud TPU

Crea un'istanza VM Cloud TPU con 8 chip TPU v6e, associata alla prenotazione della capacità.

gcloud compute instances create "${TPU_NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --machine-type=ct6e-standard-8t \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --boot-disk-size=200GB \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=RESERVATION_BOUND \
    --reservation-affinity=specific \
    --reservation="${RESERVATION}"

Dopo aver creato l'istanza VM, connettiti utilizzando SSH.

gcloud compute ssh "${TPU_NAME}" --zone "${ZONE}" --project "${PROJECT}"

Completa i passaggi seguenti all'interno dell'istanza VM TPU.

Installa MaxText

Aggiorna i pacchetti di sistema all'interno dell'istanza VM TPU.

sudo apt update && sudo apt upgrade -y --fix-missing

Installa Python 3.12, richiesto da MaxText, e il relativo pacchetto di ambiente virtuale.

sudo apt install -y build-essential cmake ninja-build

Utilizza uv per velocizzare l'installazione del pacchetto Python.

curl -LsSf https://astral.sh/uv/install.sh | sh
source "${HOME}/.local/bin/env"

Crea un ambiente virtuale denominato maxtext_venv e attivalo.

uv python install 3.12
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

Installa MaxText e le dipendenze richieste per le attività di post-addestramento.

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.3" --resolution=lowest

Installa le dipendenze rimanenti richieste eseguendo il seguente comando:

install_tpu_post_train_extra_deps

Converti il modello in formato MaxText

Per addestrare il modello in formato MaxText, devi convertirlo dal formato Hugging Face al formato MaxText.

Fornisci i seguenti valori:

  • Il token di accesso a Hugging Face
  • Il nome del modello che vuoi utilizzare
  • La directory in cui vuoi salvare il modello in formato MaxText
  • Opzioni per il caricamento e l'archiviazione
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='llama3.1-8b-Instruct'
export MODEL_CHECKPOINT_DIRECTORY=/dev/shm/$MODEL_NAME/mt-format/
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX
export LAZY_LOAD_TENSORS=False # True to use lazy load, False to use eager load.

Sostituisci YOUR_HF_TOKEN con il token di accesso a Hugging Face che hai creato in precedenza.

Per convertire il modello dal formato Hugging Face al formato MaxText, esegui il seguente script. La conversione richiede circa cinque minuti.

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name="${MODEL_NAME?}" \
    hf_access_token="${HF_TOKEN?}" \
    base_output_directory="${MODEL_CHECKPOINT_DIRECTORY?}" \
    scan_layers=True \
    use_multimodal=False \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --lazy_load_tensors="${LAZY_LOAD_TENSORS?}"

Avvia il carico di lavoro di addestramento

Al termine del processo di conversione, puoi avviare il carico di lavoro RL.

  1. Configura i parametri di addestramento del carico di lavoro RL.

    # -- MaxText configuration --
export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
export RUN_NAME
export CHIPS_PER_VM=8
export NUM_BATCHES=50
export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items
export TPU_ACCELERATOR_TYPE=v6e-8
export TPU_WORKER_ID=0
TPU_NAME=$(hostname)
export TPU_NAME
export TPU_SKIP_MDS_QUERY=1
export TPU_WORKER_HOSTNAMES=localhost
export TPU_TOPOLOGY=2x4
export TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
export TPU_HOST_BOUNDS=1,1,1

  2. Avvia il job di addestramento. Questa operazione richiede circa 10 minuti su un'istanza VM v6e-8.

    python3 -m maxtext.trainers.post_train.rl.train_rl \
    model_name="${MODEL_NAME?}" \
    load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
    run_name="${RUN_NAME?}" \
    base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
    chips_per_vm="${CHIPS_PER_VM?}" \
    num_batches="${NUM_BATCHES?}" \
    num_test_batches=10 \
    rollout_data_parallelism=1 \
    rollout_tensor_parallelism=-1

Converti di nuovo il modello addestrato in formato Hugging Face

Al termine del carico di lavoro di addestramento, converti di nuovo il modello in formato Hugging Face.

  1. Imposta i percorsi per l'esportazione e i parametri addestrati.

    export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
export HF_MODEL_NAME=llama3.1-8b
export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/actor/$NUM_BATCHES/model_params

  2. Esegui la conversione di nuovo in formato Hugging Face.

    python3 -m maxtext.checkpoint_conversion.to_huggingface \
    model_name="${HF_MODEL_NAME?}" \
    load_parameters_path="${POST_TRAIN_PATH?}" \
    base_output_directory="${HF_EXPORT?}" \
    scan_layers=True \
    use_multimodal=False \
    weight_dtype=bfloat16

Al termine della conversione, il modello ottimizzato memorizzato in /dev/shm/$MODEL_NAME/hf-trained è pronto per l'uso. Poiché perdi l'accesso ai contenuti della cartella /dev/shm al riavvio della VM, devi spostare il modello ottimizzato nello spazio di archiviazione permanente o caricarlo in Hugging Face Hub.

Esegui la pulizia

Per evitare addebiti aggiuntivi, elimina le risorse create durante questo tutorial.

Elimina l'istanza VM TPU

Elimina l'istanza VM Cloud TPU.

gcloud compute instances delete "${TPU_NAME}" --zone="${ZONE}" --project="${PROJECT}" --quiet

Passaggi successivi

  • Per saperne di più su Cloud TPU, consulta Introduzione a Cloud TPU.
  • Per i dettagli sull'architettura e sulla configurazione della v6e-8 TPU, consulta TPU v6e.