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

本教學向您展示如何使用 MaxText 和 Cluster Toolkit 在 Tensor Processing Unit (TPU) v6e-64 叢集上執行多主機強化學習 (RL) 訓練。您可以利用 Cluster Toolkit 執行多主機訓練工作負載，並將結果匯出為 Hugging Face 格式以進行服務。

目標

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

費用

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

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

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

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

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

重要事項： 如要完成本教學課程，您必須有權存取 v6e-64 TPU。本教學課程假設您已在Google Cloud 專案中預留這些 TPU，確保需要時有可用容量。如要進一步瞭解如何預訂未來日期和時間的 Cloud TPU 容量，請參閱「關於 Cloud TPU 預訂」一文。

事前準備

如要使用本教學課程，您需要 Hugging Face 存取權杖。您可以在 Hugging Face 註冊一個免費帳戶。註冊帳號後，產生訪問令牌：

1. 在「Welcome to Hugging Face」(歡迎使用 Hugging Face) 頁面中，按一下帳戶顯示圖片，然後選取「Access tokens」(存取權杖)。
2. 在「存取權杖」頁面，按一下「建立新權杖」。
3. 選取「讀取」權杖類型，然後輸入權杖名稱。
4. 您的訪問令牌已顯示。將權杖儲存在安全的地方。

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

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

• 如要完成本教學課程：
  • TPU 管理員 (roles/tpu.admin)
  • 服務帳戶使用者 (roles/iam.serviceAccountUser)
  • 運算編輯者 (roles/compute.editor)
• 準備 MaxText 容器映像檔：
  • Cloud Build 編輯者 (roles/cloudbuild.builds.editor)
  • Artifact Registry 管理員 (roles/artifactregistry.admin)
  • 儲存管理員 (roles/storage.admin)
  • 服務使用情形管理員 (roles/serviceusage.serviceUsageAdmin)

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

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

設定環境變數

設定環境變數：

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
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_REPOSITORY_NAME：MaxText 圖片的 Artifact Registry 存放區名稱。
• 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 的預先建構套件。

準備 MaxText 容器映像檔

如要準備 MaxText 容器映像檔 (包括安裝必要依附元件)，請完成下列步驟：

1. 建立 Cloud Storage bucket：

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

2. 建立 Artifact Registry 存放區；

   gcloud artifacts repositories create ${REPOSITORY_NAME} \
       --repository-format=docker \
       --location=$REGION \
       --project=$PROJECT \
       --description="Docker repository for MaxText images in $REGION" || true

3. 在存放區的根目錄中建立名為 cloudbuild.yaml 的檔案，並加入以下內容：

   steps:
     - name: 'gcr.io/cloud-builders/docker'
       entrypoint: 'bash'
       args:
         - '-c'
         - |
           set -euo pipefail
   
           # 0. Install prerequisites (if needed)
           apt-get update && apt-get install -y curl || apk add curl || true
   
           # 1. Install uv
           curl -LsSf https://astral.sh/uv/install.sh | sh
           source $$HOME/.local/bin/env
   
           # 2. Setup Python environment and install MaxText runner
           uv venv --python 3.12 --seed maxtext_venv
           source maxtext_venv/bin/activate
           uv pip install maxtext[runner]==0.2.4 --resolution=lowest
   
           # 3. Build the Docker image (Cloud Build has Docker pre-configured)
           build_maxtext_docker_image WORKFLOW=post-training
   
           # 4. Tag the image properly
           docker tag maxtext_base_image ${_CLOUD_IMAGE_NAME}
   
   # Cloud Build automatically pushes images listed here
   images:
     - '${_CLOUD_IMAGE_NAME}'
   
   options:
     machineType: 'E2_HIGHCPU_32'

4. 使用 Cloud Build 建構 MaxText Docker 映像檔：

   gcloud builds submit . \
       --project=$PROJECT \
       --region=$REGION \
       --substitutions=_CLOUD_IMAGE_NAME="${CLOUD_IMAGE_NAME}"

建立 Cluster Toolkit 叢集

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

1. 建立名為 gke.gcsfuse.profileUser 的自訂 Identity and Access Management (IAM) 角色：

   # The GKE TPU v6e blueprint uses GCS Fuse CSI Storage Profiles which requires a custom IAM role.
   # If this role is not already created in your project, you must create it before deploying.
   gcloud iam roles create gke.gcsfuse.profileUser \
     --project=${PROJECT} \
     --title="GKE GCSFuse Profile User" \
     --description="Allows scanning GCS buckets for objects, retrieving bucket metadata, and creating Anywhere Caches." \
     --permissions="storage.objects.list,storage.buckets.get,storage.anywhereCaches.create,storage.anywhereCaches.get,storage.anywhereCaches.list,storage.anywhereCaches.update"

2. 建立 Cloud Storage bucket：

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

3. 依預設，叢集節點集區服務帳戶不具備寫入 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            # Change from storage.objectViewer
       - artifactregistry.reader

4. 將自訂設定套用至 gke-tpu-v6e-cluster 區塊，覆寫預設的 IPv6 和機器類型設定：

   - id: gke-tpu-v6e-cluster
     source: modules/scheduler/gke-cluster
     use: [gke-tpu-v6e-net-0, workload_service_account]
     settings:
       enable_private_ipv6_google_access: false
       system_node_pool_disk_size_gb: $(vars.system_node_pool_disk_size_gb)
       system_node_pool_taints: []
       enable_private_endpoint: false # Allows access from authorized public IPs
       enable_pathways_for_tpus: $(vars.enable_pathways_for_tpus)
       enable_dataplane_v2: true
       configure_workload_identity_sa: true

5. 使用藍圖 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 格式。

1. 建立 Cluster Toolkit 叢集後，請設定 Docker：

   # Configure docker for pulling images
   gcloud auth configure-docker gcr.io --quiet
   gcloud auth configure-docker ${REGION}-docker.pkg.dev --quiet

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

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

3. 將模型從 Hugging Face 格式轉換為 MaxText 格式，並將其儲存在您的雲端儲存桶中：

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

4. 檢查轉換作業的狀態：

   # 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}

開始訓練工作負載

轉換過程完成後，開始強化學習訓練工作負荷：

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

# 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-training --main-only -f \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

將訓練好的模型轉換回 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=True \
      use_multimodal=False \
      weight_dtype=bfloat16 \
      --override_model_architecture"

檢查轉換作業的狀態：

# 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-mt-to-hf --main-only -f \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

清除所用資源

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

./gcluster destroy ${CLUSTER_NAME} --robust
gcloud storage rm -r gs://${GCS_BUCKET}
gcloud artifacts repositories delete ${REPOSITORY_NAME} --location=${REGION} --project=${PROJECT} --quiet

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

後續步驟

• 有關 Cloud TPU 的更多信息，請參閱 Cloud TPU 簡介。
• 如要瞭解 v6e-64 TPU 的架構和設定詳情，請參閱 TPU v6e。
• 有關集群工具包的更多信息，請參閱集群工具包概述。