在 TPU v6e 上运行 Qwen3-30b-a3b 的多主机强化学习训练

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

目标

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

费用

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Compute Engine

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

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

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

重要提示： 如需完成本教程，您必须有权访问 v6e-32 TPU。本教程假设您已在Google Cloud 项目中预留了这些 TPU，以确保您在需要时有可用容量。如需详细了解如何预留 Cloud TPU 容量以供日后使用，请参阅 Cloud TPU 预留简介。

准备工作

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

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

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

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

• TPU Admin (roles/tpu.admin)
• Kubernetes Engine Admin (roles/container.admin)
• Compute Admin (roles/compute.admin)
• Storage Admin (roles/storage.admin)
• Service Account User (roles/iam.serviceAccountUser)
• 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 COMPUTE_TYPE="ct6e-standard-4t"
export TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-30b-a3b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"
export HF_TOKEN="YOUR_HF_TOKEN"

替换以下内容：

• YOUR_PROJECT_ID：您的 Google Cloud 项目的 ID。
• YOUR_REGION：您要在其中部署集群的区域。
• YOUR_ZONE：您要在其中部署集群的可用区。
• YOUR_CLUSTER_NAME：Google Kubernetes Engine 集群的名称（最多 20 个字符）。
• YOUR_BUCKET_NAME：Cloud Storage 存储桶的全球唯一名称。
• YOUR_RESERVATION_NAME：预留的名称。
• YOUR_HF_TOKEN：您的 Hugging Face 访问令牌。

安装 Cluster Toolkit 依赖项

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

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

安装 Cluster Toolkit

按照安装 Cluster Toolkit 中的说明，在当前工作目录中安装 Cluster Toolkit 的预构建软件包。

例如，您可以按如下方式下载并将软件包提取到当前工作目录中：

wget -qO- "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:-v1.103.0}/gcluster_bundle_linux_amd64.tgz" | tar -xz

将软件包提取到当前工作目录中，可获得 gcluster 二进制文件和 examples/ 蓝图，您将在后续步骤中使用这些文件。

创建 Cluster Toolkit 集群

如需创建并部署包含 32 个 v6e TPU 芯片的 Cluster Toolkit 集群，请完成以下步骤：

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

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

2. 将 Cluster Toolkit 蓝图复制到当前工作目录：

   cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml .

3. 默认情况下，集群节点池服务账号没有写入 Cloud Storage 存储桶所需的权限。如需允许节点池服务账号写入您的 Cloud Storage 存储桶，您必须向其授予 Storage Admin 角色。如需授予此角色，请修改文件 gke-tpu-v6e-advanced.yaml，方法是更新名为 node_pool_service_account 的 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
       - artifactregistry.reader

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

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

5. 配置 Container Registry 身份验证，并向您的 GKE 服务账号授予 Storage Admin 角色 (roles/storage.admin)：

   gcloud auth configure-docker gcr.io --quiet
   gcloud auth configure-docker "${REGION}-docker.pkg.dev" --quiet
   gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet
   gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-np-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet

将模型转换为 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 qwen-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='Qwen/Qwen3-30B-A3B-Instruct-2507' \
     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"

3. 使用 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 qwen-hf-to-mt --main-only -f

4. 验证转换后的模型文件是否已在您的 Cloud Storage 存储桶中提供：

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

启动训练工作负载

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

./gcluster job submit \
  --name="qwen-rl" \
  --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="(echo 190G > /sys/fs/cgroup/memory.max || echo 190G > /sys/fs/cgroup/memory/memory.limit_in_bytes) 2>/dev/null || true && \
    export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, tunix.generate.vllm_sampler as vs; p = pathlib.Path(vs.__file__); p.write_text(p.read_text().replace('reshard_chunk_size: Optional[int] = None', 'reshard_chunk_size: Optional[int] = 4').replace('reshard_chunk_size=self.config.reshard_chunk_size', 'reshard_chunk_size=4'))\" && \
    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} \
    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 \
    num_test_batches=0 \
    batch_size=8 \
    train_micro_batch_size=8 \
    max_target_length=512 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=1 \
    rollout_expert_parallelism=4 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='Qwen/Qwen3-30B-A3B-Instruct-2507' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.55 \
    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'}}\""

• 如需查看训练作业的状态，请执行以下操作：

  # Use the list command to check status
  ./gcluster job list
  
  # Ensure kubectl credentials are configured
  gcloud container clusters get-credentials "${CLUSTER_NAME}" \
      --location="${REGION}" \
      --project="${PROJECT}"
  
  # Check progress of the job
  kubectl logs -f \
      -l jobset.sigs.k8s.io/replicatedjob-name=pathways-head \
      -c workload-container

  注意：默认情况下，GKE 上已完成的 JobSet 会在一小时后自动清理（由 gcluster job submit 中的 --gke-ttl-after-finished=1h 设置控制）。如果您的工作负载不再列在 ./gcluster job list 的输出中，这是预期行为，并不意味着作业失败。所有检查点和模型权重都会安全地保留在您的 Cloud Storage 存储桶中。

• 如需验证训练检查点是否已在 Cloud Storage 存储桶中生成，请执行以下操作：

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

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

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

./gcluster job submit \
  --name="qwen-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=True \
  use_multimodal=False \
  weight_dtype=bfloat16 \
  --override_model_architecture"

• 如需检查转换作业的状态，请执行以下操作：

  # 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 qwen-mt-to-hf --main-only -f
  # The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

• 如需验证训练后的 Hugging Face 模型权重和配置文件是否已存在于您的 Cloud Storage 存储桶中，请执行以下操作：

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

清理

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

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

# To delete the local deployment folder and copied blueprint
rm -rf .ghpc "${CLUSTER_NAME}" gke-tpu-v6e-advanced.yaml

后续步骤

• 如需详细了解 Cloud TPU，请参阅 Cloud TPU 简介。
• 如需详细了解 v6e-32 TPU 的架构和配置，请参阅 TPU v6e。
• 如需详细了解 Cluster Toolkit，请参阅 Cluster Toolkit 概览。