TPU v6e 上で Qwen3-30b-a3b のマルチホスト RL トレーニングを実行する

このチュートリアルでは、MaxText と Cluster Toolkit を使用して、Tensor Processing Unit（TPU）v6e-32 クラスタでマルチホスト強化学習（RL）トレーニングを実行する方法について説明します。Cluster Toolkit を使用して、マルチホスト トレーニング ワークロードを実行し、結果を Hugging Face 形式でエクスポートしてサービングします。

目標

• Cluster Toolkit とその依存関係をインストールします。
• Cluster Toolkit クラスタをデプロイします。
• Hugging Face モデルを MaxText 形式に変換します。
• TPU v6e クラスタで RL トレーニング ワークロードを実行する。
• サービング用にファインチューニングされたモデルを 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] ページで、アカウントのアバターをクリックして [アクセス トークン] を選択します。
  2. [アクセス トークン] ページで、[新しいトークンを作成] をクリックします。
  3. [読み取り] トークン タイプを選択し、トークンの名前を入力します。
  4. アクセス トークンが表示されます。トークンは安全な場所に保存してください。

• Hugging Face ウェブサイトで、トレーニングするモデルのライセンス契約に同意します。このチュートリアルでは、モデル qwen3-30b-a3b を使用します。

このチュートリアルを完了するために必要な権限を取得するには、プロジェクトに対する次の IAM ロールを付与するよう管理者に依頼してください。

• TPU 管理者 （roles/tpu.admin）
• Kubernetes Engine 管理者（roles/container.admin）
• Compute 管理者（roles/compute.admin）
• ストレージ管理者 （roles/storage.admin）
• サービス アカウント ユーザー（roles/iam.serviceAccountUser）
• サービス アカウント管理者 （roles/iam.serviceAccountAdmin）
• プロジェクト IAM 管理者 （roles/resourcemanager.projectIamAdmin）
• Service Usage 管理者（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 ロールを付与する必要があります。このロールを付与するには、node_pool_service_account という名前の service-account モジュールを更新して、gke-tpu-v6e-advanced.yaml ファイルを編集します。

   - 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 を使用して Cluster Toolkit クラスタをデプロイし、--vars フラグを使用して必要な変数を渡します。

   ./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 サービス アカウントにストレージ管理者ロール（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 は、デフォルトで 1 時間後に自動的にクリーンアップされます（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 の概要をご覧ください。