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

在 TPU v6e 上對 Gemma 4 執行監督式微調

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

目標

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

費用

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

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

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

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

事前準備

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

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

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

如要取得完成本教學課程所需的權限，請要求管理員在專案中授予您下列 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="gke-tpu-v6e"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-31b"
export HF_TOKEN="YOUR_HF_TOKEN"
export GCS_BUCKET="YOUR_BUCKET_NAME"

更改下列內容：

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

安裝 Cluster Toolkit 依附元件

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

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

安裝 Cluster Toolkit

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

準備 MaxText 容器映像檔

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

建立 Cloud Storage bucket：

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

建立 Artifact Registry 存放區：

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

在存放區的根目錄中建立名為 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.3 --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'

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

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

建立 Cluster Toolkit 叢集

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

建立名為 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"

建立 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 模組：

- 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

使用藍圖 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=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -w

將模型轉換為 MaxText 格式

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

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

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

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

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

如要檢查轉換工作的狀態，請執行下列指令：

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

開始訓練工作負載

轉換程序完成後，您可以執行下列指令來啟動 SFT 工作負載：

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

如要檢查訓練工作的狀態，請執行下列指令：

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

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

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

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