在 TPU v6e 上執行 Gemma 4 26B 的多主機強化學習訓練

本教學課程說明如何使用 MaxText 和 Cluster Toolkit，在 Tensor Processing Unit (TPU) v6e-32 叢集上執行多主機強化學習 (RL) 訓練。您可以使用 Cluster Toolkit 執行多主機訓練工作負載，並將結果匯出回 Hugging Face 格式，以便提供服務。

工作負載摘要

在 Cloud TPU 和 GKE 上執行這項 RL 後訓練工作負載，涉及下列元件：

• 模型：Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it)，這是先進的開放權重指令調整模型，在本教學課程中，模型權重會從 Hugging Face 下載，轉換為 MaxText 相容格式，並部署至採用 Cloud TPU v6e 基礎架構的 GKE 叢集。
• 增強學習 (RL) 演算法：使用 MaxText 和 Pathways 的群組相對策略最佳化 (GRPO)。這項工作負載會使用 MaxText 內建的強化學習訓練器 (maxtext.trainers.post_train.rl.train_rl)，搭配 vLLM 取樣器後端，執行可擴充的後訓練管道。GRPO 會為每個提示生成一組輸出內容，並評估相對獎勵，不需使用獨立的評論模型，即可最佳化政策。
• 資料集和範本：OpenMathInstruct-2 (openmathinstruct2_rl.json)，這是大規模的數學指令微調資料集。在強化學習推出階段，工作負載會使用結構化對話範本，強化多步驟數學推論能力。
• 基礎架構和框架：GKE，使用 Cluster Toolkit 管理，在 Cloud TPU v6e (ct6e-standard-4t、v6e-32 配量，採用 4x8 拓撲) 上執行。框架為 MaxText (tpu_post_training:0.2.4)，搭配 Pathways 和 JAX 後端。

目標

• 安裝 Cluster Toolkit 及其依附元件。
• 部署 Cluster Toolkit 叢集。
• 將 Hugging Face 模型轉換為 MaxText 格式。
• 在 TPU v6e 叢集執行強化學習訓練工作負載。
• 將微調後的模型轉換回 Hugging Face 格式，以便提供。

費用

在本文件中，您會使用下列 Google Cloud計費元件：

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Compute Engine

您可以使用 Pricing Calculator，根據預測用量估算費用。

初次使用 Google Cloud 的使用者可能符合免費試用期資格。

完成本文所述工作後，您可以刪除建立的資源，避免繼續計費，詳情請參閱「清除所用資源」一節。

重要事項： 如要完成本教學課程，您必須能存取 v6e-32 TPU。本教學課程假設您已在Google Cloud 專案中預留這些 TPU，確保在需要時有可用資源。如要進一步瞭解如何預留 Cloud TPU 容量，以便日後使用，請參閱「關於 Cloud TPU 預留項目」一文。

事前準備

• 您需要 Hugging Face 存取權杖，才能使用本教學課程。您可以前往 Hugging Face 註冊免費帳戶。擁有帳戶後，請產生存取權杖：

  1. 在「Welcome to Hugging Face」頁面，點選帳戶顯示圖片，然後選取「存取權杖」。
  2. 在「存取權杖」頁面，按一下「建立新權杖」。
  3. 選取「Read」權杖類型，並輸入權杖名稱。
  4. 系統會顯示存取權杖。將權杖儲存在安全的地方。

• 在 Hugging Face 的 google/gemma-4-26b-a4b-it 模型資訊卡上，接受模型的授權協議。

如要取得完成本教學課程所需的權限，請要求管理員在專案中授予您下列 IAM 角色：

• TPU 管理員 (roles/tpu.admin)
• Kubernetes Engine 管理員 (roles/container.admin)
• Compute 管理員 (roles/compute.admin)
• 儲存空間管理員 (roles/storage.admin)
• 服務帳戶使用者 (roles/iam.serviceAccountUser)
• 服務帳戶管理員 (roles/iam.serviceAccountAdmin)
• 專案 IAM 管理員 (roles/resourcemanager.projectIamAdmin)
• 服務使用情形管理員 (roles/serviceusage.serviceUsageAdmin)

如要進一步瞭解如何授予角色，請參閱「管理專案、資料夾和組織的存取權」。

您或許也能透過自訂角色或其他預先定義的角色，取得必要權限。

設定環境變數

設定環境變數：

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"

更改下列內容：

• YOUR_PROJECT_ID： Google Cloud 專案的 ID。
• YOUR_REGION：TPU 預留項目所在的區域，也是您部署叢集的位置。
• YOUR_ZONE：TPU 預留項目和叢集部署作業所在的可用區。
• YOUR_CLUSTER_NAME：GKE 叢集的名稱。名稱開頭須為小寫英文字母，只能使用小寫英文字母、數字和連字號，且長度不得超過 20 個字元。
• YOUR_BUCKET_NAME：Cloud Storage bucket 的全域不重複名稱。
• YOUR_RESERVATION_NAME：TPU v6e 預留項目的名稱。
• YOUR_HF_TOKEN：您的 Hugging Face 存取權杖。

安裝 Cluster Toolkit 依附元件

如要透過 Linux 或 macOS 用戶端或工作站完成本教學課程，請按照 Cluster Toolkit 說明文件中「安裝依附元件」的相關步驟操作。

如果您使用的是 Cloud Shell，則可略過這個部分。

安裝 Cluster Toolkit

下載並解壓縮預先建構的套件，適用於 CLUSTER_TOOLKIT_VERSION (v1.103.0) 中定義的特定 Cluster Toolkit 版本，並將其放入目前使用的目錄：

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

注意： 安裝 Cluster Toolkit 說明文件說明如何下載最新版本。不過，由於藍圖變數可能因 Cluster Toolkit 版本而異，因此在本教學課程中，您必須安裝 CLUSTER_TOOLKIT_VERSION (v1.103.0) 中指定的確切版本，而非最新版本。將套件解壓縮至目前使用的目錄後，您會取得 gcluster 二進位檔和 examples/ 藍圖，以便在後續步驟中使用。

驗證環境

執行部署指令碼前，您必須驗證 Google Cloud CLI，並設定應用程式預設憑證 (ADC)。這樣一來，Cluster Toolkit 就能代您佈建資源，並將映像檔推送至 Artifact Registry。

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 叢集

如要建立及部署具有 32 個 v6e TPU 晶片的 Cluster Toolkit 叢集，請完成下列步驟：

1. 建立 Cloud Storage bucket：

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

2. 使用 gcluster deploy 指令，透過藍圖 gke-tpu-v6e-advanced.yaml 部署 Cluster Toolkit 叢集，並使用 --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

將模型轉換為 MaxText 格式

如要以 MaxText 格式訓練模型，必須將模型從 Hugging Face 格式轉換為 MaxText 格式。

1. 如要簡化後續指令，請使用 gcluster job config 指令設定預設專案、叢集和位置：

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

2. 使用 gcluster job submit 指令，將模型從 Hugging Face 格式轉換為 MaxText 格式，並儲存至 Cloud Storage bucket：

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

   指令包含 --await-job-completion 旗標，代表指令會在目前的終端機中等待轉換完成。

3. 選用：如要在工作執行期間串流轉換記錄，請在擷取 gcluster 的目前使用的目錄中開啟另一個終端機，然後執行 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. 在原始終端機中，確認轉換後的模型檔案位於 Cloud Storage bucket：

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

啟動訓練工作負載

1. 轉換程序完成後，請啟動強化學習訓練工作負載：

   ./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. 檢查狀態，等待協調器 Pod 準備就緒，並串流訓練工作的記錄：

   # 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

   提交後，./gcluster job list 指令的輸出內容可能會短暫顯示 Suspended，同時 GKE 會擴充 Pathways 協調器 CPU 節點集區 (cpu-np)。工作開始時，會在步驟 0 (Pre RL Training: ... mean_reward=...) 輸出初始預先訓練評估，在初始 XLA 編譯期間暫停記錄輸出幾分鐘，然後訓練 50 個步驟，在每個步驟輸出 Reshard finished 和推出生成內容，不含步驟編號。

   請等到記錄輸出 Post RL Training: ... mean_reward=...，且 ./gcluster job list 指令的輸出內容顯示 gemma4-training 的狀態為 Succeeded，再繼續下一個步驟。

   注意： GKE 上已完成的 JobSet 會在設定的存留時間 (TTL) 到期後自動清除，上述訓練工作的 TTL 為 24 小時 (--gke-ttl-after-finished="24h")。已完成的工作負載在存留時間到期後，不會列在 ./gcluster job list 指令的輸出內容中，但所有檢查點和模型權重都會儲存在 Cloud Storage bucket。

3. 確認訓練檢查點已在 Cloud Storage bucket 中生成：

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

將訓練好的模型轉換回 Hugging Face 格式

1. 訓練工作負載完成後，請將模型轉換回 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. 選用：如要在第一個終端機等待上述 --await-job-completion 指令時，監控即時轉換記錄，請在擷取 gcluster 的目前使用的目錄中開啟另一個終端機，並使用 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. 在原始終端機中，確認訓練好的 Hugging Face 模型權重和設定檔，都已儲存在 Cloud Storage bucket：

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

清除所用資源

為避免產生額外費用，請使用 gcluster destroy 指令刪除在本教學課程中建立的資源：

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

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

後續步驟

• 如要進一步瞭解 Cloud TPU，請參閱「Cloud TPU 簡介」。
• 如要瞭解 v6e-32 TPU 的架構和設定詳細資料，請參閱「TPU v6e」。
• 如要進一步瞭解 Cluster Toolkit，請參閱「Cluster Toolkit 總覽」。