TPU v6e 上で Gemma 4 26B のマルチホスト強化学習トレーニングを実行する

このチュートリアルでは、MaxText と Cluster Toolkit を使用して、Tensor Processing Unit（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 を使用した Group Relative Policy Optimization（GRPO）。このワークロードは、MaxText の組み込み RL トレーナー（maxtext.trainers.post_train.rl.train_rl）と vLLM サンプラー バックエンドを組み合わせて、スケーラブルなトレーニング後のパイプラインを実行します。GRPO は、プロンプトごとに一連の出力を生成し、個別の批評家モデルを必要とせずに相対的な報酬を評価することで、ポリシーを最適化します。
• データセットとテンプレート: OpenMathInstruct-2（openmathinstruct2_rl.json）。大規模な数学指示チューニング データセット。このワークロードは、RL ロールアウト フェーズで構造化されたチャット テンプレートを使用して、マルチステップの数学的推論機能を強化します。
• インフラストラクチャとフレームワーク: Cloud TPU v6e（ct6e-standard-4t、4x8 トポロジの v6e-32 スライス）で、Cluster Toolkit を使用して管理される GKE。フレームワークは、Pathways と JAX バックエンドを備えた MaxText（tpu_post_training:0.2.4）です。

目標

• 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. [Read] トークン タイプを選択し、トークンの名前を入力します。
  4. アクセス トークンが表示されます。トークンを安全な場所に保存します。

• Hugging Face の google/gemma-4-26b-a4b-it モデルカードで、モデルの使用許諾契約に同意します。

このチュートリアルを完了するために必要な権限を取得するには、プロジェクトに対する次の 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 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

注: Cluster Toolkit のインストールのドキュメントでは、最新リリースをダウンロードする方法について説明しています。ただし、ブループリント変数は Cluster Toolkit のバージョンによって異なる可能性があるため、このチュートリアルでは、最新リリースではなく、CLUSTER_TOOLKIT_VERSION（v1.103.0）で指定された正確なバージョンをインストールする必要があります。バンドルを現在の作業ディレクトリに抽出すると、以降の手順で使用する gcluster バイナリと examples/ ブループリントが提供されます。

環境を認証する

デプロイ スクリプトを実行する前に、Google Cloud CLI を認証し、アプリケーションのデフォルト認証情報（ADC）を設定する必要があります。これにより、Cluster Toolkit はリソースをプロビジョニングし、イメージを Artifact Registry に push できます。

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

   送信直後、GKE が Pathways コーディネーター CPU ノードプール（cpu-np）をスケールアップしている間、./gcluster job list コマンドの出力に Suspended が一時的に表示されることがあります。ジョブが開始されると、ステップ 0（Pre RL Training: ... mean_reward=...）で最初の事前トレーニング評価が出力され、最初の XLA コンパイル中にログ出力が数分間一時停止します。その後、50 ステップのトレーニングが行われ、ステップ番号なしで各ステップで Reshard finished とロールアウト生成が出力されます。

   次のステップに進む前に、ログに Post RL Training: ... mean_reward=... が出力され、./gcluster job list コマンドの出力に gemma4-training のステータスが Succeeded と表示されるまで待ちます。

   注: GKE で完了した JobSet は、構成された有効期間（TTL）が切れると自動的にクリーンアップされます。上記のトレーニング ジョブの場合、TTL は 24 時間（--gke-ttl-after-finished="24h"）です。完了したワークロードは、TTL の有効期限が切れると ./gcluster job list コマンドの出力に表示されません。ただし、すべてのチェックポイントとモデルの重みは Cloud Storage バケットに保存されます。

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:?}"

次のステップ

• Cloud TPU の詳細については、Cloud TPU の概要をご覧ください。
• v6e-32 TPU のアーキテクチャと構成の詳細については、TPU v6e をご覧ください。
• Cluster Toolkit の詳細については、Cluster Toolkit の概要をご覧ください。