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

在多主機 TPU v6e 上執行 Gemma 4 31B 的 SFT

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

工作負載摘要

在 Cloud TPU 和 GKE 上執行這項 SFT 工作負載時，會涉及下列元件：

模型：Gemma 4 31B (google/gemma-4-31b-it)，先進的開放權重模型。在本教學課程中，模型權重會從 Hugging Face 下載，轉換為 MaxText 相容格式，並部署至採用 Cloud TPU v6e (v6e-32) 基礎架構的 GKE 叢集。
訓練方法：使用 MaxText 進行 SFT。工作負載會使用 JAX，在 Cloud TPU v6e 上執行可擴充的微調管道。這個指令碼會使用 MaxText 內建的 SFT 訓練器 (maxtext.trainers.post_train.sft.train_sft) 和指令微調資料集，在 1,000 個步驟中最佳化模型參數。
資料集和範本：UltraChat 200k (HuggingFaceH4/ultrachat_200k)，這是大規模、高品質的多輪對話資料集。工作負載會使用這個資料集微調模型，提升對話和指令遵循能力。
基礎架構和框架：GKE，使用 Cluster Toolkit 管理，在 Cloud TPU v6e (v6e-32，採用 4x8 拓撲) 上執行。框架為 MaxText (tpu_post_training:0.2.4)，後端為 JAX。

目標

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

費用

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

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

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

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

事前準備

您需要 Hugging Face 存取權杖，才能使用本教學課程。您可以前往 Hugging Face 註冊免費帳戶。擁有帳戶後，請產生存取權杖：
1. 在「Welcome to Hugging Face」頁面，點選帳戶顯示圖片，然後選取「存取權杖」。
2. 在「存取權杖」頁面，按一下「建立新權杖」。
3. 選取「Read」權杖類型，並輸入權杖名稱。
4. 系統會顯示存取權杖。將權杖儲存在安全的地方。

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

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

如要完成本教學課程，請按照下列步驟操作：
- TPU 管理員 (roles/tpu.admin)
- 服務帳戶使用者 (roles/iam.serviceAccountUser)
- Compute 管理員 (roles/compute.admin)
- 儲存空間管理員 (roles/storage.admin)
- Kubernetes Engine 管理員 (roles/container.admin)
- 服務帳戶管理員 (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 TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

更改下列內容：

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

根據預設，指令碼會將 MODEL_NAME 設為 gemma4-31b (google/gemma-4-31b-it 的 MaxText 模型名稱)。

安裝 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 能代表您佈建資源，並將映像檔推送至 Artifact Registry，請驗證 Google Cloud CLI 並設定應用程式預設憑證：

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud config set project "${PROJECT:?}"
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 叢集，請完成下列步驟：

建立 Cloud Storage bucket：

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

暫存藍圖設定，並將範例部署藍圖複製到獨立的目前使用的目錄：
```
mkdir -p tmp
cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/
```

更新 tmp/gke-tpu-v6e-advanced.yaml，將 Storage 管理員角色 (roles/storage.admin) 授予叢集節點集區服務帳戶 (gke-np-sa)，以便將檢查點寫入 Cloud Storage bucket。接著，將系統節點集區機型切換為 e2-standard-8，並將系統節點集區固定至目標可用區，避免區域機型缺貨：

# Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" tmp/gke-tpu-v6e-advanced.yaml
# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" tmp/gke-tpu-v6e-advanced.yaml
# 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)]' tmp/gke-tpu-v6e-advanced.yaml

使用 gcluster deploy 指令，透過藍圖 tmp/gke-tpu-v6e-advanced.yaml 部署 Cluster Toolkit 叢集，並使用 --vars 旗標傳遞必要變數：

./gcluster deploy tmp/gke-tpu-v6e-advanced.yaml \
    --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -l IGNORE \
    --auto-approve -w

將模型轉換為 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 hf-to-mt \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
      python3 -m maxtext.checkpoint_conversion.to_maxtext \
        model_name=${MODEL_NAME:?} \
        hf_access_token=${HF_TOKEN:?} \
        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"

指令包含 --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 hf-to-mt --main-only -f

在原始終端機中，確認轉換後的模型檔案位於 Cloud Storage bucket：
```
gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"
```

啟動訓練工作負載

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

./gcluster job submit --name sft \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true JAX_TRACEBACK_FILTERING=off LIBTPU_INIT_ARGS=' --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true ' \
      python3 -m maxtext.trainers.post_train.sft.train_sft \
      run_name=sft \
      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:?} \
      dataset_type=hf \
      hf_path=HuggingFaceH4/ultrachat_200k \
      per_device_batch_size=1 steps=1000 \
      profiler=xplane \
      checkpoint_storage_use_zarr3=0 \
      checkpoint_storage_use_ocdbt=0 \
      skip_jax_distributed_system=False"

指令包含 --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 sft --main-only -f
```
初始 XLA 編譯完成後，協調器會以 completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: ... 格式記錄 1000 步執行作業中每一步的輸出進度，訓練損失從第 5 步的約 1.57 降至第 1000 步的約 0.82。請等待步驟 1000 完成，且 ./gcluster job list 指令的輸出內容顯示 sft 的狀態為 Succeeded，再繼續進行下一節。

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

在原始終端機中，確認 Cloud Storage bucket 中已生成 SFT 訓練檢查點：

gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/sft/checkpoints/1000/model_params/"

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

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

./gcluster job submit --name mt-to-hf \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --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/sft/checkpoints/1000/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"