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

在 TPU v6e 上執行 Qwen3-30b-a3b 的多主機強化學習訓練

本教學課程說明如何使用 MaxText 和 Cluster Toolkit，在 Tensor 處理單元 (TPU) v6e-32 叢集上執行多主機強化學習 (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 註冊免費帳戶。擁有帳戶後，請產生存取權杖：
1. 在「Welcome to Hugging Face」(歡迎使用 Hugging Face) 頁面中，按一下帳戶顯示圖片，然後選取「Access tokens」(存取權杖)。
2. 在「Access tokens」頁面，按一下「Create new token」。
3. 選取「讀取」權杖類型，然後輸入權杖名稱。
4. 系統會顯示存取權杖。將權杖存放在安全的地方。

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

如要取得完成本教學課程所需的權限，請要求管理員在專案中授予您下列 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 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"

更改下列內容：

YOUR_PROJECT_ID：專案的 ID。 Google Cloud
YOUR_REGION：要部署叢集的區域。
YOUR_ZONE：要部署叢集的可用區。
YOUR_CLUSTER_NAME：Google Kubernetes Engine 叢集的名稱 (最多 20 個字元)。
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 的預先建構套件。

舉例來說，您可以按照下列步驟，將套件下載並解壓縮至目前的工作目錄：

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

將套件解壓縮至目前的工作目錄，即可取得後續步驟中使用的 gcluster二進位檔和 examples/ 藍圖。

建立 Cluster Toolkit 叢集

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

建立 Cloud Storage bucket：

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

將 Cluster Toolkit 藍圖複製到目前的工作目錄：
```
cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml .
```

根據預設，叢集節點集區服務帳戶沒有寫入 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

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

設定 Container Registry 驗證，並將儲存空間管理員角色 (roles/storage.admin) 授予 GKE 服務帳戶：

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

將模型轉換為 MaxText 格式

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

如要簡化後續指令，請使用 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}"

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

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

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

確認轉換後的模型檔案是否位於 Cloud Storage bucket 中：

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

啟動訓練工作負載

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

./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'}}\""

如要查看訓練工作的狀態，請按照下列步驟操作：
```
# 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
```
注意：根據預設，GKE 會在一小時後自動清除已完成的 JobSet (由 gcluster job submit 中的 --gke-ttl-after-finished=1h 設定控管)。如果 ./gcluster job list 的輸出內容不再列出您的工作負載，這是預期行為，並不代表工作失敗。所有檢查點和模型權重都會安全地保留在 Cloud Storage bucket 中。
如要確認訓練檢查點是否已在 Cloud Storage bucket 中產生，請按照下列步驟操作：
```
gcloud storage ls "gs://${GCS_BUCKET}/${MODEL_NAME}/trained/rl/checkpoints/actor/"
```

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

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

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

如要查看轉換工作的狀態，請按照下列步驟操作：

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

如要確認訓練好的 Hugging Face 模型權重和設定檔是否位於 Cloud Storage bucket 中，請按照下列步驟操作：
```
gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/"
```

清除所用資源

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

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

後續步驟

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