Führe Multi-Host-RL-Training für Qwen3-30b-a3b auf TPU v6e durch.

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

Ziele

  • Installieren Sie das Cluster Toolkit und seine Abhängigkeiten.
  • Cluster Toolkit-Cluster bereitstellen
  • Hugging Face-Modell in das MaxText-Format konvertieren
  • Führen Sie eine RL-Trainingsarbeitslast im TPU v6e-Cluster aus.
  • Konvertieren Sie das feinabgestimmte 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 unter Hugging Face für ein kostenloses Konto registrieren. Nachdem Sie ein Konto erstellt haben, können Sie ein Zugriffstoken generieren:

    1. Klicken Sie auf der Seite Welcome to Hugging Face (Willkommen bei Hugging Face) auf Ihren Konto-Avatar und wählen Sie Access tokens (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 den Token ein.
    4. Ihr Zugriffstoken wird angezeigt. Speichern Sie das Token an einem sicheren Ort.

  • Akzeptieren Sie auf der Hugging Face-Website die Lizenzvereinbarung für das Modell, das Sie trainieren möchten. In dieser Anleitung wird das Modell qwen3-30b-a3b verwendet.

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 TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-30b-a3b"
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 Sie Ihren Cluster bereitstellen möchten.
  • YOUR_ZONE: die Zone, in der Sie Ihren Cluster bereitstellen möchten.
  • YOUR_CLUSTER_NAME: Der Name Ihres Google Kubernetes Engine-Clusters (bis zu 20 Zeichen).
  • YOUR_BUCKET_NAME: ein global eindeutiger Name für einen Cloud Storage-Bucket.
  • YOUR_RESERVATION_NAME: der Name Ihrer 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 Workstation ausführen möchten, folgen Sie der Anleitung unter Abhängigkeiten installieren in der Cluster Toolkit-Dokumentation.

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

Cluster Toolkit installieren

Installieren Sie das vorgefertigte Bundle für Cluster Toolkit in Ihrem aktuellen Arbeitsverzeichnis. Folgen Sie dazu der Anleitung unter Cluster Toolkit installieren.

Sie können das Bundle beispielsweise so herunterladen und in Ihr aktuelles Arbeitsverzeichnis extrahieren:

wget -qO- "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:-v1.103.0}/gcluster_bundle_linux_amd64.tgz" | tar -xz

Wenn Sie das Bundle in Ihr aktuelles Arbeitsverzeichnis extrahieren, erhalten Sie die Binärdatei gcluster und die Blaupausen examples/, die Sie in den folgenden Schritten verwenden.

Cluster Toolkit-Cluster erstellen

So erstellen und stellen Sie einen Cluster mit 32 v6e-TPU-Chips mit dem Cluster Toolkit bereit:

  1. Erstellen Sie einen Cloud Storage-Bucket:

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

  2. Kopieren Sie den Cluster Toolkit-Blueprint in Ihr aktuelles Arbeitsverzeichnis:

    cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml .

  3. Standardmäßig hat das Dienstkonto Ihres Cluster-Knotenpools nicht die erforderlichen Berechtigungen zum Schreiben in Ihren Cloud Storage-Bucket. Damit das Dienstkonto des Knotenpools in Ihren Cloud Storage-Bucket schreiben darf, müssen Sie ihm die Rolle Storage Admin zuweisen. Um diese Rolle zuzuweisen, bearbeiten Sie die Datei gke-tpu-v6e-advanced.yaml, indem Sie das Modul service-account mit dem Namen node_pool_service_account aktualisieren:

    - id: node_pool_service_account
  source: modules/project/service-account
  settings:
    name: gke-np-sa
    project_roles:
    - logging.logWriter
    - monitoring.metricWriter
    - monitoring.viewer
    - stackdriver.resourceMetadata.writer
    - storage.admin
    - artifactregistry.reader

  4. Stellen Sie Ihren Cluster Toolkit-Cluster mit dem gcluster deploy-Befehl bereit. Verwenden Sie dazu den Blueprint gke-tpu-v6e-advanced.yaml und übergeben Sie die erforderlichen Variablen mit dem Flag --vars:

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

  5. Konfigurieren Sie die Container Registry-Authentifizierung und weisen Sie Ihren GKE-Dienstkonten die Rolle „Storage-Administrator“ (roles/storage.admin) zu:

    gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION}-docker.pkg.dev" --quiet
gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet
gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-np-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet

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 qwen-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='Qwen/Qwen3-30B-A3B-Instruct-2507' \
  base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  hardware=cpu \
  --lazy_load_tensors=True"

  3. Verwenden Sie den Befehl gcluster job logs, um den Status des Konvertierungsjobs zu prüfen:

    # 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 qwen-hf-to-mt --main-only -f

  4. Prüfen Sie, 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

Nachdem die Konvertierung abgeschlossen ist, starten Sie die RL-Trainingsarbeitslast:

./gcluster job submit \
  --name="qwen-rl" \
  --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="(echo 190G > /sys/fs/cgroup/memory.max || echo 190G > /sys/fs/cgroup/memory/memory.limit_in_bytes) 2>/dev/null || true && \
    export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, tunix.generate.vllm_sampler as vs; p = pathlib.Path(vs.__file__); p.write_text(p.read_text().replace('reshard_chunk_size: Optional[int] = None', 'reshard_chunk_size: Optional[int] = 4').replace('reshard_chunk_size=self.config.reshard_chunk_size', 'reshard_chunk_size=4'))\" && \
    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} \
    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 \
    num_test_batches=0 \
    batch_size=8 \
    train_micro_batch_size=8 \
    max_target_length=512 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=1 \
    rollout_expert_parallelism=4 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='Qwen/Qwen3-30B-A3B-Instruct-2507' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.55 \
    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'}}\""

  • So prüfen Sie den Status des Trainingsjobs:

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

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

# Check progress of the job
kubectl logs -f \
    -l jobset.sigs.k8s.io/replicatedjob-name=pathways-head \
    -c workload-container

  • So prüfen Sie, ob die Trainingsprüfpunkte in Ihrem Cloud Storage-Bucket generiert wurden:

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

Trainiertes Modell zurück in das Hugging Face-Format konvertieren

Nach Abschluss des Trainingsjobs konvertieren Sie das Modell zurück in das Hugging Face-Format:

./gcluster job submit \
  --name="qwen-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=True \
  use_multimodal=False \
  weight_dtype=bfloat16 \
  --override_model_architecture"

  • So prüfen Sie den Status des Konvertierungsjobs:

    # 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 qwen-mt-to-hf --main-only -f
# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

  • So prüfen Sie, ob die trainierten Hugging Face-Modellgewichte 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 gcluster destroy-Befehl, um die in dieser Anleitung erstellten Ressourcen zu löschen und zusätzliche Gebühren zu vermeiden:

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

# To delete the local deployment folder and copied blueprint
rm -rf .ghpc "${CLUSTER_NAME}" gke-tpu-v6e-advanced.yaml

Nächste Schritte