Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

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

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

Obiettivi

Configura un'istanza VM Cloud TPU.
Installa MaxText e le relative dipendenze.
Converti un modello Hugging Face nel formato MaxText.
Esegui un carico di lavoro di ottimizzazione relativa delle policy di gruppo (GRPO) RL sulla TPU.
Converti il modello addestrato nel formato Hugging Face per la pubblicazione.

Costi

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

Cloud TPU

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

I nuovi utenti di Google Cloud 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, è necessario un token di accesso a Hugging Face. Puoi registrarti per un account senza costi su Hugging Face. Una volta 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 prevedi di 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 progetto:

Amministratore TPU (roles/tpu.admin)
Utente Service Account (roles/iam.serviceAccountUser)
Editor Compute (roles/compute.editor)

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 lo script seguente:

export PROJECT="YOUR_PROJECT_ID"
export ZONE="ZONE_NAME"
export RESERVATION="RESERVATION_NAME"
export TPU_NAME="TPU_MACHINE_NAME"

Sostituisci quanto segue:

YOUR_PROJECT_ID: il tuo Google Cloud ID progetto
ZONE_NAME: la zona che vuoi utilizzare
RESERVATION_NAME: la prenotazione di capacità
TPU_MACHINE_NAME: il nome dell'istanza VM Cloud TPU

Autenticati con Google Cloud eseguendo questo 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 alpha compute tpus tpu-vm create $TPU_NAME \
    --zone=$ZONE \
    --project=$PROJECT \
    --accelerator-type=v6e-8 \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

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

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone $ZONE --project $PROJECT

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

Installare 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 dell'ambiente virtuale.

sudo apt install -y python3.12 python3.12-venv

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 venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

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

uv pip install maxtext[tpu-post-train]==0.2.2 --resolution=lowest

Installa le dipendenze obbligatorie rimanenti eseguendo questo comando:

install_tpu_post_train_extra_deps

Converti il modello nel 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 di caricamento e 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 lo script seguente. 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?}

Inizia il carico di allenamento

Una volta completato il processo di conversione, puoi avviare il carico di lavoro RL.

Configura i parametri di addestramento del workload RL.

# -- MaxText configuration --
export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
export CHIPS_PER_VM=8
export NUM_BATCHES=50
export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items

Avvia il job di addestramento. L'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 il modello addestrato di nuovo nel formato Hugging Face

Al termine del carico di lavoro di addestramento, converti nuovamente il modello nel formato Hugging Face.

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

Esegui la conversione nel 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 archiviato in /dev/shm/$MODEL_NAME/hf-trained è pronto per l'uso. Poiché perdi l'accesso ai contenuti della cartella /dev/shm quando la VM viene riavviata, devi spostare il modello ottimizzato in uno spazio di archiviazione permanente o caricarlo su 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 alpha compute tpus tpu-vm 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 TPU v6e-8, vedi TPU v6e.