Google 會運用 AI 技術將內容翻譯成你偏好的語言，但可能會出錯。

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

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

目標

安裝 Cluster Toolkit 及其依附元件。
部署 Cluster Toolkit 叢集。
將 Hugging Face 模型轉換為 MaxText 格式。
在 TPU v6e 叢集上執行 RL 訓練工作負載。
將微調模型轉換回 Hugging Face 格式，以供使用。

費用

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

如要根據預測用量估算費用，請使用 Pricing Calculator。

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

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

事前準備

如要使用本教學課程，您需要 Hugging Face 存取權杖。你可以在 Hugging Face 申請免費帳戶。建立帳戶後，請產生存取權杖：

在「Welcome to Hugging Face」(歡迎使用 Hugging Face) 頁面中，按一下帳戶顯示圖片，然後選取「Access tokens」(存取權杖)。
在「存取權杖」頁面，按一下「建立新權杖」。
選取「讀取」權杖類型，然後輸入權杖名稱。
畫面上會顯示存取權杖。將權杖儲存在安全的地方。

在 Hugging Face 網站上，接受要訓練模型的授權協議。本教學課程使用 gemma4-26b 模型。

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

如要完成本教學課程：
- TPU 管理員 (roles/tpu.admin)
- 服務帳戶使用者 (roles/iam.serviceAccountUser)
- 運算編輯者 (roles/compute.editor)
- 儲存空間管理員 (roles/storage.admin)

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

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

設定環境變數

設定環境變數：

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-64"
export TOPOLOGY="8x8"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export HF_TOKEN="YOUR_HF_TOKEN"

更改下列內容：

YOUR_PROJECT_ID：專案的 ID。 Google Cloud
YOUR_REGION：要部署叢集的區域。
YOUR_ZONE：要部署叢集的可用區。
YOUR_CLUSTER_NAME：Google Kubernetes Engine 叢集的名稱。
YOUR_BUCKET_NAME：Cloud Storage bucket 的全域不重複名稱。
YOUR_RESERVATION_NAME：預訂名稱。
YOUR_HF_TOKEN：您的 Hugging Face 存取權杖。

安裝 Cluster Toolkit 依附元件

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

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

安裝 Cluster Toolkit

按照「安裝 Cluster Toolkit」一文中的說明，安裝 Cluster Toolkit 的預先建構套件。

建立 Cluster Toolkit 叢集

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

建立 Cloud Storage bucket：

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

依預設，叢集節點集區服務帳戶不具備寫入 Cloud Storage bucket 的必要權限。如要允許節點集區服務帳戶寫入 Cloud Storage bucket，請授予 Storage Admin 角色。如要授予這個角色，請編輯 gke-tpu-v6e-advanced.yaml 檔案，更新名為 node_pool_service_account 的 service-account 模組：

- 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

使用藍圖 gke-tpu-v6e-advanced.yaml 部署 Cluster Toolkit 叢集，並使用 --vars 旗標傳遞必要變數：

./gcluster deploy examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml \
    --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 格式。

如要簡化後續指令，請設定預設專案、叢集和位置：

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

將模型從 Hugging Face 格式轉換為 MaxText 格式，並儲存在 Cloud Storage bucket 中：

./gcluster job submit \
  --name="gemma4-hf-to-mt" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --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"

查看轉換工作的狀態：

# Use the list command to check status
./gcluster job list \
    --cluster "${CLUSTER_NAME}" \
    --project "${PROJECT}" \
    --location "${REGION}"

# 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 \
    --cluster "${CLUSTER_NAME}" \
    --project "${PROJECT}" \
    --location "${REGION}"

啟動訓練工作負載

轉換程序完成後，請啟動 RL 訓練工作負載：

./gcluster job submit \
  --name="gemma4-training" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET}/pathways/" \
  --env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}'); \
      python3 -c \"import pathlib, tpu_inference.layers.common.fused_moe_gmm as f; p = pathlib.Path(f.__file__); p.write_text(p.read_text().replace('onehot_moe_permute_threshold: int = 0,', 'onehot_moe_permute_threshold: int = 100000,'))\"; \
      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} \
      num_batches=50 \
      batch_size=8 \
      rollout_tensor_parallelism=2 \
      rollout_expert_parallelism=4 \
      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 \
      hbm_utilization_vllm=0.55 \
      remat_policy=full \
      async_scheduling=False \
      allow_split_physical_axes=true \
      ragged_gather_reduce_fallback=True \
      vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \
      vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'allow_split_physical_axes': 'true', 'use_ragged_sort': 'false', 'ragged_gather_reduce_fallback': 'true', 'prefuse_moe_weights': 'true', 'weight_dtype': 'bfloat16'}}\""

查看訓練工作的狀態：

# Use the list command to check status
./gcluster job list \
    --cluster "${CLUSTER_NAME}" \
    --project "${PROJECT}" \
    --location "${REGION}"

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

# Check progress of the job
kubectl logs -f \
    -l job-name=gemma4-training-pathways-head-0 \
    -c workload-container

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

訓練工作負載完成後，請將模型轉換回 Hugging Face 格式：

./gcluster job submit \
  --name="gemma4-mt-to-hf" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --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"