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

在多主机 TPU v6e 上针对 Gemma 4 31B 运行 SFT

本教程介绍了如何使用 MaxText 和 Cluster Toolkit 在多主机张量处理单元 (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的以下收费组件：

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

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

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

准备工作

您需要拥有 Hugging Face 访问令牌才能使用本教程。您可以在 Hugging Face 注册免费账号。拥有账号后，生成访问令牌：
1. 在欢迎使用 Hugging Face 页面上，点击您的账号头像，然后选择访问令牌。
2. 在访问令牌页面上，点击创建新令牌。
3. 选择读取令牌类型，然后输入令牌名称。
4. 系统会显示您的访问令牌。将令牌保存在安全的位置。

在 Hugging Face 上的 google/gemma-4-31b-it 模型卡片中，接受模型的许可协议。

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

如需完成本教程，您需要：
- TPU Admin (roles/tpu.admin)
- Service Account User (roles/iam.serviceAccountUser)
- Compute Admin (roles/compute.admin)
- Storage Admin (roles/storage.admin)
- Kubernetes Engine Admin (roles/container.admin)
- Service Account Admin (roles/iam.serviceAccountAdmin)
- Project IAM Admin (roles/resourcemanager.projectIamAdmin)
- Service Usage Admin (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 存储桶的全球唯一名称。
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 存储桶，请运行以下命令：

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，向集群节点池服务账号 (gke-np-sa) 授予 Storage Admin 角色 (roles/storage.admin)，以便将检查点写入 Cloud Storage 存储桶；将系统节点池机器类型切换为 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 并使用 --vars 标志传递所需变量来部署 Cluster Toolkit 集群：

./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 存储桶中：

./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 存储桶中可用：
```
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)（默认为 1 小时 [--gke-ttl-after-finished=1h]）到期后自动清理。已完成的工作负载在 TTL 到期后不会列在 ./gcluster job list 命令的输出中；但是，所有检查点和模型权重都会存储在 Cloud Storage 存储桶中。

在原始终端中，验证 SFT 训练检查点是否已在 Cloud Storage 存储桶中生成：

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"