Eseguire l'addestramento di apprendimento per rinforzo multi-host per Gemma 4 26B su TPU v6e

Questo tutorial mostra come eseguire l'addestramento di apprendimento per rinforzo (RL) multi-host su un cluster di Tensor Processing Unit (TPU) v6e-32 utilizzando MaxText e Cluster Toolkit. Utilizza Cluster Toolkit per eseguire un workload di addestramento multi-host ed esportare i risultati nel formato Hugging Face per la pubblicazione.

Riepilogo del workload

L'esecuzione di questo workload di post-addestramento RL su Cloud TPU e GKE prevede i seguenti componenti:

  • Modello: Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it), un modello open-weights avanzato ottimizzato per le istruzioni. In questo tutorial, i pesi del modello vengono scaricati da Hugging Face, convertiti in un formato compatibile con MaxText ed è stato eseguito il deployment su un cluster GKE con l'infrastruttura Cloud TPU v6e.
  • Algoritmo di apprendimento per rinforzo (RL): Group Relative Policy Optimization (GRPO) utilizzando MaxText e Pathways. Il workload esegue una pipeline di post-addestramento scalabile utilizzando l'addestratore RL integrato di MaxText (maxtext.trainers.post_train.rl.train_rl) accoppiato a un backend del campionatore vLLM. GRPO ottimizza la policy generando un gruppo di output per prompt e valutando i reward relativi senza bisogno di un modello critic separato.
  • Set di dati e template: OpenMathInstruct-2 (openmathinstruct2_rl.json), un set di dati di ottimizzazione delle istruzioni matematiche su larga scala. Il workload utilizza modelli di chat strutturati durante la fase di implementazione dell'RL per migliorare le capacità di ragionamento matematico in più passaggi.
  • Infrastruttura e framework: GKE, gestito utilizzando Cluster Toolkit, su Cloud TPU v6e (slice ct6e-standard-4t, v6e-32 con una topologia 4x8). Il framework è MaxText (tpu_post_training:0.2.4) con Pathways e un backend JAX.

Obiettivi

  • Installa Cluster Toolkit e le relative dipendenze.
  • Esegui il deployment di un cluster Cluster Toolkit.
  • Converti un modello Hugging Face nel formato MaxText.
  • Esegui un workload di addestramento RL sul cluster TPU v6e.
  • Converti il modello ottimizzato nel 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 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, devi disporre di un token di accesso a Hugging Face. Puoi creare un account senza costi su Hugging Face. Dopo aver creato un account, genera un token di accesso:

    1. Nella pagina Welcome to 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.

Per ottenere le autorizzazioni necessarie per completare questo tutorial, chiedi all'amministratore di concederti i seguenti ruoli IAM nel 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 le variabili di ambiente

Imposta le variabili di ambiente:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_post_training:0.2.4"
export COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-32"
export TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"
export HF_TOKEN="YOUR_HF_TOKEN"

Sostituisci quanto segue:

  • YOUR_PROJECT_ID: l'ID del tuo progetto Google Cloud .
  • YOUR_REGION: la regione in cui si trova la prenotazione della TPU e in cui esegui il deployment del cluster.
  • YOUR_ZONE: la zona in cui si trova la prenotazione della TPU e in cui esegui il deployment del cluster.
  • YOUR_CLUSTER_NAME: il nome del tuo cluster GKE. Il nome deve iniziare con una lettera minuscola, contenere solo lettere minuscole, numeri e trattini e avere una lunghezza massima di 20 caratteri.
  • YOUR_BUCKET_NAME: un nome univoco a livello globale per un bucket Cloud Storage.
  • YOUR_RESERVATION_NAME: il nome della prenotazione TPU v6e.
  • YOUR_HF_TOKEN: il tuo token di accesso a Hugging Face.

Installa le dipendenze di Cluster Toolkit

Per completare questo tutorial da un client o una workstation Linux o macOS, segui i passaggi pertinenti in Installa le dipendenze nella documentazione di Cluster Toolkit.

Se utilizzi Cloud Shell, puoi saltare questa sezione.

Installa Cluster Toolkit

Scarica ed estrai il bundle predefinito per la versione specifica di Cluster Toolkit definita in CLUSTER_TOOLKIT_VERSION (v1.103.0) nella directory di lavoro attuale:

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64)              BUNDLE="gcluster_bundle_linux_amd64.tgz" ;;
  Linux-aarch64|Linux-arm64) BUNDLE="gcluster_bundle_linux_arm64.tgz" ;;
  Darwin-x86_64)             BUNDLE="gcluster_bundle_mac_amd64.tgz" ;;
  Darwin-arm64)              BUNDLE="gcluster_bundle_mac_arm64.tgz" ;;
  *) BUNDLE=""; echo "Unsupported platform: $(uname -s)-$(uname -m)" >&2 ;;
esac
mkdir -p cluster-toolkit
curl -fsSL "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:?}/${BUNDLE:?}" \
  | tar -xz -C cluster-toolkit
cd cluster-toolkit

Autentica il tuo ambiente

Prima di eseguire gli script di deployment, devi autenticare Google Cloud CLI e impostare le Credenziali predefinite dell'applicazione (ADC). In questo modo, Cluster Toolkit può eseguire il provisioning delle risorse e il push delle immagini in Artifact Registry per tuo conto.

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud auth application-default set-quota-project "${PROJECT:?}"

gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION:?}-docker.pkg.dev" --quiet

Crea il tuo cluster Cluster Toolkit

Per creare ed eseguire il deployment di un cluster Cluster Toolkit con 32 chip TPU v6e, completa i seguenti passaggi:

  1. Crea un bucket Cloud Storage:

    gcloud storage buckets create "gs://${GCS_BUCKET:?}" --project="${PROJECT:?}" --location="${REGION:?}" || true

  2. Utilizza il comando gcluster deploy per eseguire il deployment del cluster Cluster Toolkit utilizzando il progetto gke-tpu-v6e-advanced.yaml e passando le variabili richieste utilizzando il flag --vars:

    TMP_DIR=$(mktemp -d)
BLUEPRINT_FILE="${TMP_DIR}/gke-tpu-v6e-advanced.yaml"

cp ./examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml "${BLUEPRINT_FILE}"

# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" "${BLUEPRINT_FILE}"
# Pin system node pool to a single zone to avoid cross-zone stockouts
sed -i '/system_node_pool_machine_type/a \      system_node_pool_zones: [$(vars.zone)]' "${BLUEPRINT_FILE}"
# Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" "${BLUEPRINT_FILE}"

./gcluster deploy "${BLUEPRINT_FILE}" \
    --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=${CLUSTER_NODEPOOL_COUNT:?},tpu_topology=${TOPOLOGY:?},authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    -l IGNORE --auto-approve -w

Converti il modello nel formato MaxText

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

  1. Per semplificare i comandi successivi, utilizza il comando gcluster job config per configurare il progetto, il cluster e la località predefiniti:

    # Configure gcluster Defaults
./gcluster job config set project "${PROJECT:?}"
./gcluster job config set cluster "${CLUSTER_NAME:?}"
./gcluster job config set location "${REGION:?}"

  2. Utilizza il comando gcluster job submit per convertire il modello dal formato Hugging Face al formato MaxText e archiviarlo nel tuo bucket Cloud Storage:

    ./gcluster job submit \
  --name="gemma4-hf-to-mt" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --await-job-completion \
  --command="[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
    python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name=${MODEL_NAME:?} \
    hf_access_token=${HF_TOKEN:?} \
    --hf_model_path='google/gemma-4-26b-a4b-it' \
    base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME}/max-text-format/ \
    scan_layers=False \
    use_multimodal=False \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=0 \
    checkpoint_storage_use_ocdbt=0 \
    hardware=cpu \
    --lazy_load_tensors=True"

    Il comando include il flag --await-job-completion, il che significa che attende nel terminale corrente fino al completamento della conversione.

  3. (Facoltativo) Per trasmettere in streaming i log di conversione mentre il job è in esecuzione, apri un terminale separato nella directory di lavoro in cui hai estratto gcluster ed esegui il comando gcluster job logs:

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs gemma4-hf-to-mt --main-only -f

  4. Nel terminale originale, verifica che i file del modello convertito siano disponibili nel bucket Cloud Storage:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"

Avvia il workload di addestramento

  1. Al termine del processo di conversione, avvia il workload di addestramento RL:

    ./gcluster job submit \
  --name="gemma4-training" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET:?}/pathways/" \
  --gke-ttl-after-finished="24h" \
  --restarts=0 \
  --env="GRPC_DNS_RESOLVER=native" \
  --env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, re; p = pathlib.Path('/deps/src/maxtext/trainers/post_train/rl/utils_rl.py'); p.write_text(re.sub('optax[.]adamw[(][^)]+[)]', 'optax.adafactor(learning_rate=learning_rate)', p.read_text()))\" && \
    JAX_PLATFORMS=proxy,cpu ENABLE_PATHWAYS_PERSISTENCE=1 \
    python3 -m maxtext.trainers.post_train.rl.train_rl \
    run_name=rl \
    base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME:?}/trained/ \
    model_name=${MODEL_NAME} \
    scan_layers=False \
    load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
    hf_access_token=${HF_TOKEN:?} \
    data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json \
    num_batches=50 \
    batch_size=4 \
    train_micro_batch_size=2 \
    max_target_length=640 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=4 \
    rollout_expert_parallelism=1 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='google/gemma-4-26b-a4b-it' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.7 \
    remat_policy=full \
    async_scheduling=False \
    allow_split_physical_axes=true \
    ragged_gather_reduce_fallback=True \
    enable_dp_attention=False \
    decode_sampling_temperature=0.8 \
    decode_sampling_top_k=50 \
    decode_sampling_nucleus_p=0.95 \
    learning_rate=2e-5 \
    learning_rate_schedule_steps=100 \
    rl.num_generations=4 \
    rl.reshard_chunk_size=4 \
    debug=True \
    vllm_hf_overrides='{\"architectures\": [\"MaxTextForCausalLM\"]}' \
    vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'model_call_mode': 'inference', 'enable_dp_attention': false, 'allow_split_physical_axes': true, 'use_ragged_sort': false, 'ragged_gather_reduce_fallback': true, 'prefuse_moe_weights': true, 'weight_dtype': 'bfloat16'}}\""

  2. Controlla lo stato, attendi che il pod del coordinatore sia pronto e trasmetti in streaming i log del job di addestramento:

    # Use the list command to check status (initially shows "Suspended" briefly while the CPU node pool scales up)
./gcluster job list

# Ensure kubectl is configured
gcloud container clusters get-credentials "${CLUSTER_NAME:?}" \
  --location="${REGION:?}" \
  --project="${PROJECT:?}"

# Wait for the Pathways head pod to be ready before streaming logs
kubectl wait --for=condition=Ready pod \
    -l job-name=gemma4-training-pathways-head-0 \
    --timeout=600s

# Stream training logs (runs for 50 steps; wait for "Post RL Training: ... mean_reward=..." and job status Succeeded before converting to Hugging Face)
kubectl logs -f \
    -l job-name=gemma4-training-pathways-head-0 \
    -c workload-container

    Subito dopo l'invio, l'output del comando ./gcluster job list potrebbe mostrare brevemente Suspended mentre GKE esegue lo scale up del pool di nodi di CPU del coordinatore Pathways (cpu-np). All'avvio del job, viene generata una valutazione iniziale di pre-addestramento al passaggio 0 (Pre RL Training: ... mean_reward=...), l'output del log viene messo in pausa per diversi minuti durante la compilazione XLA iniziale e poi viene eseguito l'addestramento per 50 passaggi, stampando Reshard finished e le generazioni di implementazione in ogni passaggio senza numeri di passaggio.

    Prima di procedere con il passaggio successivo, attendi che l'output dei log sia Post RL Training: ... mean_reward=... e che l'output del comando ./gcluster job list mostri lo stato di gemma4-training come Succeeded.

  3. Verifica che i checkpoint di addestramento siano generati nel tuo bucket Cloud Storage:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/rl/checkpoints/actor/"

Converti il modello addestrato nel formato Hugging Face

  1. Al termine del workload di addestramento, converti il modello nel formato Hugging Face:

    ./gcluster job submit \
  --name="gemma4-mt-to-hf" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --await-job-completion \
  --command="[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
    python3 -m maxtext.checkpoint_conversion.to_huggingface \
      model_name=${MODEL_NAME:?} \
      hf_access_token=${HF_TOKEN:?} \
      load_parameters_path=gs://${GCS_BUCKET:?}/${MODEL_NAME}/trained/rl/checkpoints/actor/50/model_params/ \
      base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ \
      skip_jax_distributed_system=true \
      hardware=cpu \
      scan_layers=False \
      use_multimodal=False \
      weight_dtype=bfloat16 \
      --override_model_architecture"

  2. (Facoltativo) Per monitorare i log di conversione in tempo reale mentre il comando --await-job-completion precedente è in attesa nel primo terminale, apri un terminale separato nella directory di lavoro in cui hai estratto gcluster e utilizza il comando gcluster job logs:

    # Use the list command to check status
./gcluster job list 

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs gemma4-mt-to-hf --main-only -f

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

  3. Nel terminale originale, verifica che i pesi e i file di configurazione del modello Hugging Face addestrato siano presenti nel bucket Cloud Storage:

    gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/hf-trained/"

Esegui la pulizia

Per evitare addebiti aggiuntivi, utilizza il comando gcluster destroy per eliminare le risorse create durante questo tutorial:

./gcluster destroy "${CLUSTER_NAME:?}" --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

# To delete the local deployment folder
rm -rf .ghpc "${CLUSTER_NAME:?}"

Passaggi successivi