Reinforcement Learning-Training mit mehreren Hosts für Gemma 4 26B auf TPU v6e ausführen

In dieser Anleitung erfahren Sie, wie Sie mit MaxText und dem Cluster Toolkit ein Multi-Host-Reinforcement-Learning-Training (RL) auf einem Tensor Processing Unit (TPU) v6e-32-Cluster ausführen. Sie verwenden das Cluster Toolkit, um eine Trainingsarbeitslast mit mehreren Hosts auszuführen und die Ergebnisse zur Bereitstellung wieder in das Hugging Face-Format zu exportieren.

Zusammenfassung der Arbeitslast

Für die Ausführung dieser RL-Arbeitslast nach dem Training auf Cloud TPU und GKE sind die folgenden Komponenten erforderlich:

  • Modell: Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it), ein fortschrittliches, für Anweisungen optimiertes Open-Weights-Modell. In dieser Anleitung werden die Modellgewichtungen von Hugging Face heruntergeladen, in ein MaxText-kompatibles Format konvertiert und in einem GKE-Cluster mit Cloud TPU v6e-Infrastruktur bereitgestellt.
  • Algorithmus für bestärkendes Lernen (RL): Group Relative Policy Optimization (GRPO) mit MaxText und Pathways. Die Arbeitslast führt eine skalierbare Post-Training-Pipeline mit dem integrierten RL-Trainer von MaxText (maxtext.trainers.post_train.rl.train_rl) in Kombination mit einem vLLM-Sampler-Backend aus. GRPO optimiert die Richtlinie, indem pro Prompt eine Gruppe von Ausgaben generiert und die relativen Belohnungen bewertet werden, ohne dass ein separates Kritikermodell erforderlich ist.
  • Dataset und Vorlage: OpenMathInstruct-2 (openmathinstruct2_rl.json), ein umfangreiches Dataset zur Feinabstimmung von mathematischen Anweisungen. Die Arbeitslast verwendet während der RL-Roll-out-Phase strukturierte Chatvorlagen, um die Fähigkeiten zur mehrstufigen mathematischen Argumentation zu verbessern.
  • Infrastruktur und Framework: GKE, verwaltet mit Cluster Toolkit, auf Cloud TPU v6e (ct6e-standard-4t, v6e-32-Slice mit einer 4x8-Topologie). Das Framework ist MaxText (tpu_post_training:0.2.4) mit Pathways und einem JAX-Backend.

Ziele

  • Cluster Toolkit und seine Abhängigkeiten installieren
  • Cluster Toolkit-Cluster bereitstellen
  • Hugging Face-Modell in das MaxText-Format konvertieren.
  • Führen Sie eine RL-Trainingsarbeitslast auf dem TPU v6e-Cluster aus.
  • Konvertieren Sie das optimierte Modell für die Bereitstellung zurück in das Hugging Face-Format.

Kosten

In diesem Dokument verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen.

Neuen Nutzern von Google Cloud steht möglicherweise eine kostenlose Testversion zur Verfügung.

Nach Abschluss der in diesem Dokument beschriebenen Aufgaben können Sie weitere Kosten vermeiden, indem Sie die erstellten Ressourcen löschen. Weitere Informationen finden Sie unter Bereinigen.

Hinweis

  • Für diese Anleitung benötigen Sie ein Hugging Face-Zugriffstoken. Sie können sich bei Hugging Face für ein kostenloses Konto registrieren. Nachdem Sie ein Konto erstellt haben, generieren Sie ein Zugriffstoken:

    1. Klicken Sie auf der Seite Willkommen bei Hugging Face auf Ihren Konto-Avatar und wählen Sie Zugriffstokens aus.
    2. Klicken Sie auf der Seite Zugriffstokens auf Neues Token erstellen.
    3. Wählen Sie den Tokentyp Lesen aus und geben Sie einen Namen für das Token ein.
    4. Ihr Zugriffstoken wird angezeigt. Speichern Sie das Token an einem sicheren Ort.

Bitten Sie Ihren Administrator, Ihnen die folgenden IAM-Rollen für Ihr Projekt zuzuweisen, um die Berechtigungen zu erhalten, die Sie zum Ausführen dieser Anleitung benötigen:

Weitere Informationen zum Zuweisen von Rollen finden Sie unter Zugriff auf Projekte, Ordner und Organisationen verwalten.

Sie können die erforderlichen Berechtigungen auch über benutzerdefinierte Rollen oder andere vordefinierte Rollen erhalten.

Umgebungsvariablen einrichten

Richten Sie die Umgebungsvariablen ein:

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"

Ersetzen Sie Folgendes:

  • YOUR_PROJECT_ID: die ID Ihres Projekts in Google Cloud .
  • YOUR_REGION: die Region, in der sich Ihre TPU-Reservierung befindet und in der Sie Ihren Cluster bereitstellen.
  • YOUR_ZONE: die Zone, in der sich Ihre TPU-Reservierung befindet und in der Sie Ihren Cluster bereitstellen.
  • YOUR_CLUSTER_NAME: Name Ihres GKE-Clusters. Der Name muss mit einem Kleinbuchstaben beginnen, darf nur Kleinbuchstaben, Ziffern und Bindestriche enthalten und darf maximal 20 Zeichen lang sein.
  • YOUR_BUCKET_NAME: ein global eindeutiger Name für einen Cloud Storage-Bucket.
  • YOUR_RESERVATION_NAME: der Name Ihrer TPU v6e-Reservierung.
  • YOUR_HF_TOKEN: Ihr Hugging Face-Zugriffstoken.

Cluster Toolkit-Abhängigkeiten installieren

Wenn Sie diese Anleitung auf einem Linux- oder macOS-Client oder einer Linux- oder macOS-Workstation ausführen möchten, folgen Sie den entsprechenden Schritten unter Abhängigkeiten installieren in der Cluster Toolkit-Dokumentation.

Wenn Sie Cloud Shell verwenden, können Sie diesen Abschnitt überspringen.

Cluster Toolkit installieren

Laden Sie das vordefinierte Bundle für die in CLUSTER_TOOLKIT_VERSION (v1.103.0) definierte Version des Cluster Toolkits herunter und extrahieren Sie es in Ihr aktuelles Arbeitsverzeichnis:

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

Umgebung authentifizieren

Bevor Sie die Bereitstellungsskripts ausführen, müssen Sie die Google Cloud CLI authentifizieren und Ihre Standardanmeldedaten für Anwendungen (Application Default Credentials, ADC) festlegen. So kann das Cluster Toolkit in Ihrem Namen Ressourcen bereitstellen und Images per Push an Artifact Registry senden.

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

Cluster Toolkit-Cluster erstellen

Führen Sie die folgenden Schritte aus, um einen Cluster Toolkit-Cluster mit 32 v6e-TPU-Chips zu erstellen und bereitzustellen:

  1. Erstellen Sie einen Cloud Storage-Bucket:

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

  2. Verwenden Sie den Befehl gcluster deploy, um Ihren Cluster Toolkit-Cluster mit dem Blueprint gke-tpu-v6e-advanced.yaml bereitzustellen und die erforderlichen Variablen mit dem Flag --vars zu übergeben:

    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

Modell in das MaxText-Format konvertieren

Wenn Sie das Modell im MaxText-Format trainieren möchten, müssen Sie es vom Hugging Face-Format in das MaxText-Format konvertieren.

  1. Verwenden Sie den gcluster job config-Befehl, um Ihr Standardprojekt, Ihren Standardcluster und Ihren Standardstandort zu konfigurieren und nachfolgende Befehle zu vereinfachen:

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

  2. Verwenden Sie den Befehl gcluster job submit, um das Modell vom Hugging Face-Format in das MaxText-Format zu konvertieren und in Ihrem Cloud Storage-Bucket zu speichern:

    ./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"

    Der Befehl enthält das Flag --await-job-completion. Das bedeutet, dass er in Ihrem aktuellen Terminal wartet, bis die Konvertierung abgeschlossen ist.

  3. Optional: Wenn Sie die Konvertierungslogs streamen möchten, während der Job ausgeführt wird, öffnen Sie ein separates Terminal im Arbeitsverzeichnis, in dem Sie gcluster extrahiert haben, und führen Sie den Befehl gcluster job logs aus:

    # 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. Prüfen Sie im ursprünglichen Terminal, ob die konvertierten Modelldateien in Ihrem Cloud Storage-Bucket verfügbar sind:

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

Trainingsarbeitslast starten

  1. Starten Sie nach Abschluss des Konvertierungsvorgangs die RL-Trainingsarbeitslast:

    ./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. Prüfen Sie den Status, warten Sie, bis der Koordinator-Pod bereit ist, und streamen Sie die Logs des Trainingsjobs:

    # 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

    Unmittelbar nach dem Senden wird in der Ausgabe des Befehls ./gcluster job list möglicherweise kurz Suspended angezeigt, während GKE den CPU-Knotenpool des Pathways-Koordinators (cpu-np) hochskaliert. Wenn der Job beginnt, wird bei Schritt 0 (Pre RL Training: ... mean_reward=...) eine erste Bewertung vor dem Training ausgegeben. Die Logausgabe wird während der ersten XLA-Kompilierung für einige Minuten pausiert. Anschließend wird der Job für 50 Schritte trainiert, wobei bei jedem Schritt Reshard finished und Rollout-Generierungen ohne Schrittnummern ausgegeben werden.

    Bevor Sie mit dem nächsten Schritt fortfahren, warten Sie, bis die Logs Post RL Training: ... mean_reward=... ausgeben und die Ausgabe des Befehls ./gcluster job list den Status von gemma4-training als Succeeded anzeigt.

  3. Prüfen Sie, ob die Trainingsprüfpunkte in Ihrem Cloud Storage-Bucket generiert werden:

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

Trainiertes Modell wieder in das Hugging Face-Format konvertieren

  1. Nachdem die Trainingsarbeitslast abgeschlossen ist, konvertieren Sie das Modell wieder in das Hugging Face-Format:

    ./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. Optional: Wenn Sie die Live-Konvertierungslogs überwachen möchten, während der vorherige --await-job-completion-Befehl im ersten Terminal wartet, öffnen Sie ein separates Terminal im Arbeitsverzeichnis, in dem Sie gcluster extrahiert haben, und verwenden Sie den gcluster job logs-Befehl:

    # 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. Prüfen Sie im ursprünglichen Terminal, ob die trainierten Hugging Face-Modellgewichtungen und ‑Konfigurationsdateien in Ihrem Cloud Storage-Bucket vorhanden sind:

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

Bereinigen

Verwenden Sie den Befehl gcluster destroy, um die in dieser Anleitung erstellten Ressourcen zu löschen und zusätzliche Kosten zu vermeiden:

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

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

Nächste Schritte