使用 MaxText 在多部主機上對 Qwen3-14b 模型執行監督式微調

本教學課程提供逐步指南，說明如何使用 Cloud TPU 上的 MaxText，對 Qwen3-14b 模型執行監督式微調 (SFT)。您將瞭解如何建構專用容器映像檔、使用 Accelerated Processing Kit (XPK)，透過 Pathways 佈建 Google Kubernetes Engine (GKE) 叢集，以及執行多主機訓練工作負載。

目標

• 瞭解如何建構自訂 MaxText 容器映像檔，並針對訓練後作業進行最佳化。
• 使用 XPK 佈建 GKE 叢集，並啟用 Pathways。
• 將 Qwen3 14b 模型從 Hugging Face 格式轉換為 MaxText 格式。
• 在 Cloud TPU 上執行多主機 SFT 訓練工作負載。
• 將微調模型轉換回 Hugging Face 格式，以供使用。

費用

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine, for the temporary build virtual machine (VM)

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

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

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

事前準備

• 確認使用者帳戶或服務帳戶具備下列角色：
  • roles/compute.admin，建立建構 VM
  • roles/artifactregistry.admin，管理 Docker 存放區
  • roles/storage.admin，管理資料 bucket
  • roles/container.admin，建立及管理 Google Kubernetes Engine 叢集
  • roles/iam.serviceAccountAdmin，建立工作負載服務帳戶
  • roles/resourcemanager.projectIamAdmin，設定身分與存取權管理 (IAM) 政策
  • roles/iam.serviceAccountUser，以服務帳戶身分執行動作
• 安裝並初始化 Google Cloud CLI。
• 確認工作站已安裝 Python 3.12 以上版本。

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

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

設定環境

執行下列指令碼，設定環境變數：

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export CLOUD_IMAGE_NAME="$REGION-docker.pkg.dev/$PROJECT/maxtext-images/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-14b"
export HF_TOKEN="YOUR_HF_TOKEN"

更改下列內容：

• YOUR_PROJECT_ID：您的 Google Cloud 專案 ID
• YOUR_REGION：要使用的區域
• YOUR_ZONE：要使用的可用區
• YOUR_CLUSTER_NAME：Google Kubernetes Engine 叢集的名稱
• YOUR_GCS_BUCKET：Cloud Storage bucket 的專屬名稱
• YOUR_RESERVATION_NAME：運算資源預留項目
• YOUR_HF_TOKEN：您的 Hugging Face 存取權杖

準備 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 maxtext-images \
       --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.1 --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:
     # We use a high-CPU machine to match the n4-standard-16 from the VM tutorial
     machineType: 'E2_HIGHCPU_32'

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

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

建立 Google Kubernetes Engine 叢集

如要在 Qwen3 14b 模型上執行 SFT 訓練，您需要配備 TPU 晶片的 Google Kubernetes Engine 叢集。安裝 Accelerated Processing Kit (XPK)，並建立支援 Pathways 的 GKE 叢集。

# Start with creating a new virtual environment to install XPK in.
VENV_DIR=venvp3
python3 -m venv $VENV_DIR
source $VENV_DIR/bin/activate
pip install xpk==1.14.0

xpk cluster create-pathways \
  --num-slices=${CLUSTER_NODEPOOL_COUNT} \
  --tpu-type=${TPU_TYPE} \
  --pathways-gce-machine-type=${PW_CPU_MACHINE_TYPE} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --cluster=${CLUSTER_NAME} \
  --custom-cluster-arguments="--enable-ip-alias" \
  --reservation=$RESERVATION \
  --default-pool-cpu-machine-type=n4-standard-16

gcloud container clusters get-credentials $CLUSTER_NAME \
  --location=$REGION \
  --project $PROJECT

準備訓練模型

使用 CPU 型工作負載，將基礎模型轉換為 MaxText 格式。請勿在多部機器上並行執行這項工作。下列指令包含檢查，確保轉換只在一個 TPU 節點上執行。

xpk workload create \
  --workload "qwen-hf-to-mt" \
  --docker-image $CLOUD_IMAGE_NAME \
  --cluster ${CLUSTER_NAME} \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --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}/qwen-3-14b/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  hardware=cpu \
  --lazy_load_tensors=True"

追蹤模型轉換進度

如要追蹤轉換進度，請按照下列步驟操作：

1. 如要列出 GKE 叢集上排定的 Pod，請執行 kubectl get pod 指令。
2. 找出名為 qwen-hf-to-mt-slice-job-0-0-HASH 的 Pod。
3. 如要即時檢查 Pod 的輸出內容，請執行 kubectl logs -f POD_NAME 指令。

開始訓練工作負載

轉換程序完成後，您可以使用 XPK 啟動 SFT 微調工作負載。

xpk workload create-pathways \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-training" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \
  python3 -m maxtext.trainers.post_train.sft.train_sft \
  run_name=sft \
  base_output_directory=gs://${GCS_BUCKET}/qwen-3-14b/trained/ \
  model_name=${MODEL_NAME} \
  load_parameters_path=gs://${GCS_BUCKET}/qwen-3-14b/max-text-format/0/items/ \
  hf_access_token=${HF_TOKEN} \
  per_device_batch_size=1 \
  steps=1000 \
  profiler=xplane \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  enable_single_controller=True"

監控訓練工作負載

使用 XPK 指令列介面 (CLI) 監控工作負載狀態。

xpk workload list --cluster ${CLUSTER_NAME} --project ${PROJECT} --zone ${ZONE}

如要查看記錄和 TPU 使用率，請使用 Google Cloud 控制台。您也可以執行下列指令來查看記錄：

kubectl logs -f qwen-training-pathways-head-0-0-HASH

將 HASH 替換為 Pod 名稱中的數字雜湊。如要驗證這個雜湊值，請執行 kubectl get pod 指令，並檢查傳回的 Pod 清單。

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

訓練工作負載完成後，請將檢查點轉換回 Hugging Face 格式。

xpk workload create \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-mt-to-hf" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --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}/qwen-3-14b/trained/sft/checkpoints/1000/model_params/ \
  base_output_directory=gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ \
  skip_jax_distributed_system=true \
  hardware=cpu \
  scan_layers=True \
  use_multimodal=False \
  weight_dtype=bfloat16"

如要追蹤轉換進度，請執行 kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH 指令，並將 HASH 替換為 Pod 名稱中的數字雜湊。

轉換完成後，儲存在 gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ 中的調整後模型即可使用。

清除所用資源

為避免產生額外費用，請刪除本教學課程中建立的資源，包括 Google Kubernetes Engine 叢集、Cloud Storage 值區和 Artifact Registry 存放區。

注意： 如果您打算繼續使用在本教學課程中訓練的模型，請先儲存可存取的版本，再刪除包含該模型的 Cloud Storage 值區。

如要刪除您為本教學課程建立的資源，請執行下列指令：

xpk cluster delete --cluster $CLUSTER_NAME --project $PROJECT --zone $ZONE --force

gcloud storage rm --recursive gs://$GCS_BUCKET

gcloud artifacts repositories delete maxtext-images --location=$REGION --project=$PROJECT --quiet

後續步驟

• 如要進一步瞭解 Cloud TPU，請參閱「Cloud TPU 簡介」。
• 如要瞭解 v6e-32 TPU 的架構和設定詳情，請參閱 TPU v6e。