Google 会使用 AI 技术将内容翻译成您偏好的语言。AI 翻译可能包含错误。

在 TPU v6e 上对 Gemma 4 运行监督式微调

本教程介绍了如何使用 MaxText 和 Cluster Toolkit 在张量处理单元 (TPU) v6e 集群上运行监督式微调 (SFT)。您可以使用 Cluster Toolkit 执行多主机训练工作负载，并将结果导出为 Hugging Face 格式以供服务。

目标

安装 Cluster Toolkit 及其依赖项。
安装 MaxText 及其依赖项。
部署 Cluster Toolkit 集群。
将 Hugging Face 模型转换为 MaxText 格式。
在 TPU 上运行 SFT 训练工作负载。
将微调后的模型转换回 Hugging Face 格式以用于提供服务。

费用

在本文档中，您将使用 Google Cloud的以下收费组件：

如需根据您的预计使用情况来估算费用，请使用价格计算器。

新 Google Cloud 用户可能有资格申请免费试用。

完成本文档中描述的任务后，您可以通过删除所创建的资源来避免继续计费。如需了解详情，请参阅清理。

准备工作

您需要拥有 Hugging Face 访问令牌才能使用本教程。您可以在 Hugging Face 上注册免费账号。拥有账号后，生成访问令牌：

在 Welcome to Hugging Face 页面上，点击您的账号头像，然后选择 Access tokens。
在访问令牌页面上，点击创建新令牌。
选择读取令牌类型，然后输入令牌的名称。
系统会显示您的访问令牌。将令牌保存在安全的位置。

在 Hugging Face 网站上，接受您计划训练的模型的许可协议。本教程使用模型 gemma4-31b。

如需获得完成本教程所需的权限，请让您的管理员为您授予项目的以下 IAM 角色：

如需完成本教程，请执行以下操作：
- TPU Admin (roles/tpu.admin)
- Service Account User (roles/iam.serviceAccountUser)
- Compute Editor (roles/compute.editor)
如需准备 MaxText 容器映像，请执行以下操作：
- Cloud Build Editor (roles/cloudbuild.builds.editor)
- Artifact Registry Administrator (roles/artifactregistry.admin)
- Storage Administrator (roles/storage.admin)
- Service Usage Administrator (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：您的 Google Cloud 项目的 ID。
YOUR_REGION：您要在其中部署集群的区域。
YOUR_ZONE：您要部署集群的可用区。
YOUR_REPOSITORY_NAME：MaxText 映像的 Artifact Registry 代码库的名称。
YOUR_RESERVATION_NAME：预留的名称。
YOUR_HF_TOKEN：您的 Hugging Face 访问令牌。
YOUR_BUCKET_NAME：Cloud Storage 存储桶的全球唯一名称。

安装 Cluster Toolkit 依赖项

如需从 Linux 或 macOS 客户端或工作站完成本教程，请按照 Cluster Toolkit 文档中的安装依赖项中的相关步骤操作。

如果您使用的是 Cloud Shell，则可以跳过此部分。

安装 Cluster Toolkit

按照安装 Cluster Toolkit 中的说明安装 Cluster Toolkit 的预构建软件包。

准备 MaxText 容器映像

如需准备 MaxText 容器映像（包括安装所需依赖项），请完成以下步骤：

创建 Cloud Storage 存储分区，请运行以下命令：

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 存储分区，请运行以下命令：

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

默认情况下，集群节点池服务账号没有写入 Cloud Storage 存储桶所需的权限。如需允许节点池服务账号写入您的 Cloud Storage 存储桶，您必须为其授予 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 存储桶中，请运行以下脚本：

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