在 TPU v6e 上运行 Gemma 4 26B 的多主机强化学习训练

本教程介绍了如何使用 MaxText 和 Cluster Toolkit 在张量处理单元 (TPU) v6e-32 集群上运行多主机强化学习 (RL) 训练。您可以使用 Cluster Toolkit 执行多主机训练工作负载，并将结果导出回 Hugging Face 格式以进行传送。

工作负载摘要

在 Cloud TPU 和 GKE 上运行此 RL 后训练工作负载涉及以下组件：

  • 模型：Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it)，这是一款先进的开放权重指令调优模型。在本教程中，模型权重从 Hugging Face 下载，转换为 MaxText 兼容格式，并部署在具有 Cloud TPU v6e 基础设施的 GKE 集群上。
  • 强化学习 (RL) 算法：使用 MaxText 和 Pathways 的群组相对策略优化 (GRPO)。该工作负载使用 MaxText 的内置 RL 训练器 (maxtext.trainers.post_train.rl.train_rl) 搭配 vLLM 采样器后端，执行可伸缩的训练后流水线。GRPO 通过为每个提示生成一组输出并评估相对奖励来优化策略，而无需单独的批评者模型。
  • 数据集和模板：OpenMathInstruct-2 (openmathinstruct2_rl.json)，一个大规模数学指令调优数据集。该工作负载在强化学习 (RL) 部署阶段使用结构化聊天模板，以增强多步数学推理能力。
  • 基础架构和框架：GKE，使用 Cluster Toolkit 进行管理，在 Cloud TPU v6e（ct6e-standard-4t，v6e-32 切片，采用 4x8 拓扑）上运行。框架为 MaxText (tpu_post_training:0.2.4)，采用 Pathways 和 JAX 后端。

目标

  • 安装 Cluster Toolkit 及其依赖项。
  • 部署 Cluster Toolkit 集群。
  • 将 Hugging Face 模型转换为 MaxText 格式。
  • 在 TPU v6e 集群上运行强化学习训练工作负载。
  • 将经过微调的模型转换回 Hugging Face 格式，以便进行应用。

费用

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

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

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

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

准备工作

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

    1. 在欢迎使用 Hugging Face 页面上，点击您的账号头像，然后选择访问令牌。
    2. 在访问令牌页面上，点击创建新令牌。
    3. 选择读取令牌类型，然后输入令牌名称。
    4. 系统会显示您的访问令牌。将令牌保存在安全的位置。

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

如需详细了解如何授予角色，请参阅管理对项目、文件夹和组织的访问权限。

您也可以通过自定义角色或其他预定义角色来获取所需的权限。

设置环境变量

设置环境变量：

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 COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-32"
export TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"
export HF_TOKEN="YOUR_HF_TOKEN"

替换以下内容：

  • YOUR_PROJECT_ID：您的 Google Cloud 项目的 ID。
  • YOUR_REGION：TPU 预留所在的区域以及您部署集群的区域。
  • YOUR_ZONE：TPU 预留所在的可用区以及您部署集群的可用区。
  • YOUR_CLUSTER_NAME：GKE 集群的名称。名称必须以小写字母开头，只能包含小写字母、数字和连字符，且长度不超过 20 个字符。
  • YOUR_BUCKET_NAME：Cloud Storage 存储桶的全球唯一名称。
  • YOUR_RESERVATION_NAME：TPU v6e 预留的名称。
  • YOUR_HF_TOKEN：您的 Hugging Face 访问令牌。

安装 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

对您的环境进行身份验证

在运行部署脚本之前，您必须对 Google Cloud CLI 进行身份验证并设置应用默认凭据 (ADC)。这可确保 Cluster Toolkit 能够代表您预配资源并将映像推送到 Artifact Registry。

gcloud auth login
gcloud auth application-default login --no-launch-browser
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 集群，请完成以下步骤：

  1. 创建 Cloud Storage 存储桶，请运行以下命令：

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

  2. 使用 gcluster deploy 命令，通过使用蓝图 gke-tpu-v6e-advanced.yaml 并使用 --vars 标志传递所需变量来部署 Cluster Toolkit 集群：

    TMP_DIR=$(mktemp -d)
BLUEPRINT_FILE="${TMP_DIR}/gke-tpu-v6e-advanced.yaml"

cp ./examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml "${BLUEPRINT_FILE}"

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

./gcluster deploy "${BLUEPRINT_FILE}" \
    --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=${CLUSTER_NODEPOOL_COUNT:?},tpu_topology=${TOPOLOGY:?},authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    -l IGNORE --auto-approve -w

将模型转换为 MaxText 格式

如需以 MaxText 格式训练模型，您必须将其从 Hugging Face 格式转换为 MaxText 格式。

  1. 为简化后续命令，请使用 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:?}"

  2. 使用 gcluster job submit 命令将模型从 Hugging Face 格式转换为 MaxText 格式，并将其存储在 Cloud Storage 存储桶中：

    ./gcluster job submit \
  --name="gemma4-hf-to-mt" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --await-job-completion \
  --command="[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
    python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name=${MODEL_NAME:?} \
    hf_access_token=${HF_TOKEN:?} \
    --hf_model_path='google/gemma-4-26b-a4b-it' \
    base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME}/max-text-format/ \
    scan_layers=False \
    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 标志，这意味着它会在当前终端中等待，直到转换完成。

  3. 可选：如需在作业运行时流式传输转换日志，请在提取 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 gemma4-hf-to-mt --main-only -f

  4. 在原始终端中，验证转换后的模型文件是否在 Cloud Storage 存储桶中可用：

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"

启动训练工作负载

  1. 转换过程完成后，启动 RL 训练工作负载：

    ./gcluster job submit \
  --name="gemma4-training" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET:?}/pathways/" \
  --gke-ttl-after-finished="24h" \
  --restarts=0 \
  --env="GRPC_DNS_RESOLVER=native" \
  --env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, re; p = pathlib.Path('/deps/src/maxtext/trainers/post_train/rl/utils_rl.py'); p.write_text(re.sub('optax[.]adamw[(][^)]+[)]', 'optax.adafactor(learning_rate=learning_rate)', p.read_text()))\" && \
    JAX_PLATFORMS=proxy,cpu ENABLE_PATHWAYS_PERSISTENCE=1 \
    python3 -m maxtext.trainers.post_train.rl.train_rl \
    run_name=rl \
    base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME:?}/trained/ \
    model_name=${MODEL_NAME} \
    scan_layers=False \
    load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
    hf_access_token=${HF_TOKEN:?} \
    data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json \
    num_batches=50 \
    batch_size=4 \
    train_micro_batch_size=2 \
    max_target_length=640 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=4 \
    rollout_expert_parallelism=1 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='google/gemma-4-26b-a4b-it' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.7 \
    remat_policy=full \
    async_scheduling=False \
    allow_split_physical_axes=true \
    ragged_gather_reduce_fallback=True \
    enable_dp_attention=False \
    decode_sampling_temperature=0.8 \
    decode_sampling_top_k=50 \
    decode_sampling_nucleus_p=0.95 \
    learning_rate=2e-5 \
    learning_rate_schedule_steps=100 \
    rl.num_generations=4 \
    rl.reshard_chunk_size=4 \
    debug=True \
    vllm_hf_overrides='{\"architectures\": [\"MaxTextForCausalLM\"]}' \
    vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'model_call_mode': 'inference', 'enable_dp_attention': false, 'allow_split_physical_axes': true, 'use_ragged_sort': false, 'ragged_gather_reduce_fallback': true, 'prefuse_moe_weights': true, 'weight_dtype': 'bfloat16'}}\""

  2. 检查状态，等待协调器 Pod 准备就绪，然后流式传输训练作业的日志：

    # Use the list command to check status (initially shows "Suspended" briefly while the CPU node pool scales up)
./gcluster job list

# Ensure kubectl is configured
gcloud container clusters get-credentials "${CLUSTER_NAME:?}" \
  --location="${REGION:?}" \
  --project="${PROJECT:?}"

# Wait for the Pathways head pod to be ready before streaming logs
kubectl wait --for=condition=Ready pod \
    -l job-name=gemma4-training-pathways-head-0 \
    --timeout=600s

# Stream training logs (runs for 50 steps; wait for "Post RL Training: ... mean_reward=..." and job status Succeeded before converting to Hugging Face)
kubectl logs -f \
    -l job-name=gemma4-training-pathways-head-0 \
    -c workload-container

    提交后，./gcluster job list 命令的输出可能会在 GKE 扩容 Pathways 协调器 CPU 节点池 (cpu-np) 时短暂显示 Suspended。作业开始后，它会在第 0 步 (Pre RL Training: ... mean_reward=...) 输出初始预训练评估，在初始 XLA 编译期间暂停日志输出几分钟，然后训练 50 步，在每一步打印 Reshard finished 和发布生成，不显示步数。

    在继续执行下一步之前，请等待日志输出 Post RL Training: ... mean_reward=...，并且 ./gcluster job list 命令的输出显示 gemma4-training 的状态为 Succeeded。

  3. 验证 Cloud Storage 存储桶中是否生成了训练检查点：

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/rl/checkpoints/actor/"

将经过训练的模型转换回 Hugging Face 格式

  1. 训练工作负载完成后，将模型转换回 Hugging Face 格式：

    ./gcluster job submit \
  --name="gemma4-mt-to-hf" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --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/rl/checkpoints/actor/50/model_params/ \
      base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ \
      skip_jax_distributed_system=true \
      hardware=cpu \
      scan_layers=False \
      use_multimodal=False \
      weight_dtype=bfloat16 \
      --override_model_architecture"

  2. 可选：如需在第一个终端中等待上述 --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 gemma4-mt-to-hf --main-only -f

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

  3. 在原始终端中，验证经过训练的 Hugging Face 模型权重和配置文件是否存在于 Cloud Storage 存储桶中：

    gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/hf-trained/"

清理

为避免产生额外费用，请使用 gcluster destroy 命令删除本教程中创建的资源：

./gcluster destroy "${CLUSTER_NAME:?}" --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

# To delete the local deployment folder
rm -rf .ghpc "${CLUSTER_NAME:?}"

后续步骤