Exécuter un entraînement d'apprentissage par renforcement multi-hôte pour Gemma 4 26B sur TPU v6e

Ce tutoriel vous explique comment exécuter un entraînement d'apprentissage par renforcement (RL) multi-hôte sur un cluster de Tensor Processing Unit (TPU) v6e-32 à l'aide de MaxText et de Cluster Toolkit. Vous utilisez Cluster Toolkit pour exécuter une charge de travail d'entraînement multi-hôte et exporter les résultats au format Hugging Face pour la diffusion.

Résumé de la charge de travail

L'exécution de cette charge de travail de post-entraînement par RL sur Cloud TPU et GKE implique les composants suivants :

  • Modèle : Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it), un modèle avancé à pondération ouverte et réglé pour les instructions. Dans ce tutoriel, les pondérations du modèle sont téléchargées depuis Hugging Face, converties dans un format compatible avec MaxText et déployées sur un cluster GKE avec une infrastructure Cloud TPU v6e.
  • Algorithme d'apprentissage par renforcement (RL) : Group Relative Policy Optimization (GRPO) à l'aide de MaxText et Pathways. La charge de travail exécute un pipeline de post-entraînement évolutif à l'aide de l'outil d'entraînement RL intégré de MaxText (maxtext.trainers.post_train.rl.train_rl) associé à un backend d'échantillonneur vLLM. GRPO optimise la stratégie en générant un groupe de sorties par requête et en évaluant les récompenses relatives sans avoir besoin d'un modèle critique distinct.
  • Ensemble de données et modèle : OpenMathInstruct-2 (openmathinstruct2_rl.json), un ensemble de données de réglage d'instructions mathématiques à grande échelle. La charge de travail utilise des modèles de chat structurés pendant la phase de déploiement de l'apprentissage par renforcement pour améliorer les capacités de raisonnement mathématique en plusieurs étapes.
  • Infrastructure et framework : GKE, géré à l'aide de Cluster Toolkit, sur Cloud TPU v6e (ct6e-standard-4t, tranche v6e-32 avec une topologie 4x8). Le framework est MaxText (tpu_post_training:0.2.4) avec Pathways et un backend JAX.

Objectifs

  • Installez Cluster Toolkit et ses dépendances.
  • Déployer un cluster Cluster Toolkit
  • Convertissez un modèle Hugging Face au format MaxText.
  • Exécutez une charge de travail d'entraînement par RL sur le cluster TPU v6e.
  • Convertissez le modèle affiné au format Hugging Face pour la diffusion.

Coûts

Dans ce document, vous utilisez les composants facturables suivants de Google Cloud :

Obtenez une estimation des coûts en fonction de votre utilisation prévue à l'aide du simulateur de coût.

Les nouveaux utilisateurs de Google Cloud peuvent bénéficier d'un essai sans frais.

Une fois que vous avez terminé les tâches décrites dans ce document, supprimez les ressources que vous avez créées pour éviter que des frais vous soient facturés. Pour en savoir plus, consultez la section Effectuer un nettoyage.

Avant de commencer

  • Vous avez besoin d'un jeton d'accès Hugging Face pour suivre ce tutoriel. Vous pouvez créer un compte sans frais sur Hugging Face. Une fois que vous disposez d'un compte, générez un jeton d'accès :

    1. Sur la page Welcome to Hugging Face (Bienvenue sur Hugging Face), cliquez sur l'avatar de votre compte, puis sélectionnez Access tokens (Jetons d'accès).
    2. Sur la page Access tokens (Jetons d'accès), cliquez sur Create new token (Créer un jeton).
    3. Sélectionnez le type de jeton Read (Lecture) et saisissez un nom pour votre jeton.
    4. Votre jeton d'accès s'affiche. Enregistrez le jeton dans un endroit sûr.

Pour obtenir les autorisations nécessaires pour suivre ce tutoriel, demandez à votre administrateur de vous accorder les rôles IAM suivants sur votre projet :

Pour en savoir plus sur l'attribution de rôles, consultez la page Gérer l'accès aux projets, aux dossiers et aux organisations.

Vous pouvez également obtenir les autorisations requises avec des rôles personnalisés ou d'autres rôles prédéfinis.

Configurer vos variables d'environnement

Configurez vos variables d'environnement :

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"

Remplacez les éléments suivants :

  • YOUR_PROJECT_ID : ID de votre projet Google Cloud .
  • YOUR_REGION : région où se trouve votre réservation de TPU et où vous déployez votre cluster.
  • YOUR_ZONE : zone où se trouve votre réservation de TPU et où vous déployez votre cluster.
  • YOUR_CLUSTER_NAME : nom de votre cluster GKE. Le nom doit commencer par une lettre minuscule, ne contenir que des lettres minuscules, des chiffres et des traits d'union, et comporter 20 caractères ou moins.
  • YOUR_BUCKET_NAME : nom unique d'un bucket Cloud Storage.
  • YOUR_RESERVATION_NAME : nom de votre réservation de TPU v6e.
  • YOUR_HF_TOKEN : votre jeton d'accès Hugging Face.

Installer les dépendances de Cluster Toolkit

Pour suivre ce tutoriel à partir d'un client ou d'un poste de travail Linux ou macOS, suivez les étapes pertinentes de la section Installer les dépendances de la documentation de Cluster Toolkit.

Si vous utilisez Cloud Shell, vous pouvez ignorer cette section.

Installer Cluster Toolkit

Téléchargez et extrayez le bundle prédéfini pour la version spécifique de Cluster Toolkit définie dans CLUSTER_TOOLKIT_VERSION (v1.103.0) dans votre répertoire de travail actuel :

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

Authentifier votre environnement

Avant d'exécuter les scripts de déploiement, vous devez authentifier la Google Cloud CLI et définir vos Identifiants par défaut de l'application (ADC). Cela permet à Cluster Toolkit de provisionner des ressources et de transférer des images vers Artifact Registry en votre nom.

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

Créer votre cluster Cluster Toolkit

Pour créer et déployer un cluster Cluster Toolkit avec 32 puces TPU v6e, procédez comme suit :

  1. Créez un bucket Cloud Storage :

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

  2. Utilisez la commande gcluster deploy pour déployer votre cluster Cluster Toolkit à l'aide du plan gke-tpu-v6e-advanced.yaml et en transmettant les variables requises à l'aide de l'option --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

Convertir le modèle au format MaxText

Pour entraîner le modèle au format MaxText, vous devez le convertir du format Hugging Face au format MaxText.

  1. Pour simplifier les commandes suivantes, utilisez la commande gcluster job config pour configurer votre projet, votre cluster et votre emplacement par défaut :

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

  2. Utilisez la commande gcluster job submit pour convertir le modèle du format Hugging Face au format MaxText et le stocker dans votre 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"

    La commande inclut l'option --await-job-completion, ce qui signifie qu'elle attend dans votre terminal actuel que la conversion se termine.

  3. Facultatif : Pour diffuser les journaux de conversion pendant l'exécution du job, ouvrez un terminal distinct dans le répertoire de travail où vous avez extrait gcluster, puis exécutez la commande 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. Dans votre terminal d'origine, vérifiez que les fichiers du modèle converti sont disponibles dans votre bucket Cloud Storage :

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

Démarrer la charge de travail d'entraînement

  1. Une fois le processus de conversion terminé, démarrez la charge de travail d'entraînement par 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. Vérifiez l'état, attendez que le pod coordinateur soit prêt et diffusez les journaux du job d'entraînement :

    # 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

    Immédiatement après l'envoi, la sortie de la commande ./gcluster job list peut afficher brièvement Suspended pendant que GKE effectue un scaling à la hausse du pool de nœuds de processeur du coordinateur Pathways (cpu-np). Lorsque le job commence, il génère une évaluation initiale de pré-entraînement à l'étape 0 (Pre RL Training: ... mean_reward=...), suspend la sortie du journal pendant plusieurs minutes lors de la compilation XLA initiale, puis s'entraîne pendant 50 étapes, en affichant Reshard finished et les générations de déploiement à chaque étape sans numéro d'étape.

    Avant de passer à l'étape suivante, attendez que les journaux affichent Post RL Training: ... mean_reward=... et que le résultat de la commande ./gcluster job list indique l'état Succeeded pour gemma4-training.

  3. Vérifiez que les points de contrôle d'entraînement sont générés dans votre bucket Cloud Storage :

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

Convertir le modèle entraîné au format Hugging Face

  1. Une fois la charge de travail d'entraînement terminée, reconvertissez le modèle au format 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. Facultatif : Pour surveiller les journaux de conversion en direct pendant que la commande --await-job-completion précédente attend dans votre premier terminal, ouvrez un terminal distinct dans le répertoire de travail où vous avez extrait gcluster, puis utilisez la commande 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. Dans votre terminal d'origine, vérifiez que les pondérations et les fichiers de configuration du modèle Hugging Face entraîné sont présents dans votre bucket Cloud Storage :

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

Effectuer un nettoyage

Pour éviter des frais supplémentaires, utilisez la commande gcluster destroy pour supprimer les ressources créées au cours de ce tutoriel :

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

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

Étapes suivantes