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

このチュートリアルでは、MaxText と Cluster Toolkit を使用して、Tensor Processing Unit (TPU) v6e-64クラスタ上でマルチホスト強化学習 (RL) トレーニングを実行する方法を示します。Cluster Toolkitを使用して、マルチホストのトレーニングワークロードを実行し、結果を Hugging Face 形式にエクスポートして提供します。

目標

• Cluster Toolkit とその依存関係をインストールしてください。
• MaxText とその依存関係をインストールしてください。
• Cluster Toolkit クラスターをデプロイします。
• ハグ顔モデルを MaxText 形式に変換します。
• TPU v6e クラスタ上で強化学習のトレーニングワークロードを実行します。
• 微調整したモデルを、提供用にハグフェイス形式に戻します。

費用

このドキュメントでは、課金対象である次の Google Cloudコンポーネントを使用します。

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

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。

新規の Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

このドキュメントに記載されているタスクの完了後、作成したリソースを削除すると、それ以上の請求は発生しません。詳細については、クリーンアップをご覧ください。

重要: このチュートリアルを完了するには、v6e-64 TPU にアクセスできる必要があります。このチュートリアルでは、必要なときに容量が利用できるように、Google Cloud プロジェクトでこれらの TPU の予約が既に行われていることを前提としています。将来の日時を指定してクラウド TPU 容量を予約するためのオプションの詳細については、クラウド TPU の予約について.

を参照してください。

始める前に

このチュートリアルを使用するには、Hugging Face のアクセストークンが必要です。Hugging Faceで無料アカウントに登録できます。アカウントを作成したら、アクセストークンを生成してください。

1. Hugging Face へようこそページで、アカウントのアバターをクリックし、アクセストークンを選択します。
2. オンアクセストークンページをクリック新しいトークンを作成する。
3. Readトークンタイプを選択し、トークンの名前を入力してください。
4. アクセストークンが表示されました。トークンは安全な場所に保管してください。

• Hugging Face のウェブサイトで、トレーニング予定のモデルのライセンス契約に同意してください。このチュートリアルではモデルgemma4-26bを使用します。

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

• このチュートリアルを完了するには：
  • TPU 管理者 (roles/tpu.admin)
  • サービス アカウント ユーザー（roles/iam.serviceAccountUser）
  • コンピュータエディタ (roles/compute.editor)
• MaxText コンテナイメージを準備するには：
  • Cloud Build 編集者 （roles/cloudbuild.builds.editor）
  • Artifact Registry 管理者（roles/artifactregistry.admin）
  • ストレージ管理者 (roles/storage.admin)
  • サービス利用管理者 (roles/serviceusage.serviceUsageAdmin)

ロールの付与については、プロジェクト、フォルダ、組織へのアクセス権の管理をご覧ください。

必要な権限は、カスタムロールや他の事前定義ロールから取得することもできます。

環境変数を設定する

環境変数を設定します。

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-64"
export TOPOLOGY="8x8"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export HF_TOKEN="YOUR_HF_TOKEN"

次のように置き換えます。

• YOUR_PROJECT_ID: 実際の Google Cloud プロジェクト ID。
• YOUR_REGION: クラスターをデプロイするリージョン。
• YOUR_ZONE: クラスターをデプロイするゾーン。
• YOUR_CLUSTER_NAME: Google Kubernetes Engine クラスタの名前。
• YOUR_REPOSITORY_NAME: MaxText イメージのアーティファクトレジストリリポジトリの名前。
• YOUR_BUCKET_NAME: クラウドストレージバケットのグローバルに一意な名前。
• YOUR_RESERVATION_NAME: 予約の名前。
• YOUR_HF_TOKEN: Hugging Face アクセス トークン。

Cluster Toolkit の依存関係をインストールする

Linux または macOS クライアントまたはワークステーションからこのチュートリアルを完了するには、Cluster Toolkit ドキュメントの 依存関係のインストール の関連手順に従ってください。

Cloud Shell を使用している場合は、このセクションはスキップしてください。

Cluster Toolkit をインストールする

Cluster Toolkit をインストールするの手順に沿って、Cluster Toolkit 用のビルド済みバンドルをインストールします。

MaxText コンテナイメージを準備する

MaxText コンテナイメージを準備するには、必要な依存関係のインストールを含め、以下の手順を実行してください。

1. Cloud Storage バケットを作成します。

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

2. アーティファクトレジストリリポジトリを作成する。

   gcloud artifacts repositories create ${REPOSITORY_NAME} \
       --repository-format=docker \
       --location=$REGION \
       --project=$PROJECT \
       --description="Docker repository for MaxText images in $REGION" || true

3. リポジトリのルートディレクトリに、ファイル名が cloudbuild.yaml で、以下の内容を記述したファイルを作成してください。

   steps:
     - name: 'gcr.io/cloud-builders/docker'
       entrypoint: 'bash'
       args:
         - '-c'
         - |
           set -euo pipefail
   
           # 0. Install prerequisites (if needed)
           apt-get update && apt-get install -y curl || apk add curl || true
   
           # 1. Install uv
           curl -LsSf https://astral.sh/uv/install.sh | sh
           source $$HOME/.local/bin/env
   
           # 2. Setup Python environment and install MaxText runner
           uv venv --python 3.12 --seed maxtext_venv
           source maxtext_venv/bin/activate
           uv pip install maxtext[runner]==0.2.4 --resolution=lowest
   
           # 3. Build the Docker image (Cloud Build has Docker pre-configured)
           build_maxtext_docker_image WORKFLOW=post-training
   
           # 4. Tag the image properly
           docker tag maxtext_base_image ${_CLOUD_IMAGE_NAME}
   
   # Cloud Build automatically pushes images listed here
   images:
     - '${_CLOUD_IMAGE_NAME}'
   
   options:
     machineType: 'E2_HIGHCPU_32'

4. Cloud Build を使用して MaxText の Docker イメージをビルドします。

   gcloud builds submit . \
       --project=$PROJECT \
       --region=$REGION \
       --substitutions=_CLOUD_IMAGE_NAME="${CLOUD_IMAGE_NAME}"

Cluster Toolkit クラスターを作成します

Cluster Toolkit クラスターを 64 で作成およびデプロイするにはv6eTPU チップを使用する場合は、以下の手順を実行してください。

1. gke.gcsfuse.profileUser という名前のカスタム ID およびアクセス管理 (IAM) ロールを作成します。

   # The GKE TPU v6e blueprint uses GCS Fuse CSI Storage Profiles which requires a custom IAM role.
   # If this role is not already created in your project, you must create it before deploying.
   gcloud iam roles create gke.gcsfuse.profileUser \
     --project=${PROJECT} \
     --title="GKE GCSFuse Profile User" \
     --description="Allows scanning GCS buckets for objects, retrieving bucket metadata, and creating Anywhere Caches." \
     --permissions="storage.objects.list,storage.buckets.get,storage.anywhereCaches.create,storage.anywhereCaches.get,storage.anywhereCaches.list,storage.anywhereCaches.update"

2. Cloud Storage バケットを作成します。

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

3. デフォルトでは、クラスタノードプールサービスアカウントには、クラウドストレージバケットへの書き込みに必要な権限がありません。ノードプールサービスアカウントがクラウドストレージバケットに書き込めるようにするには、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            # Change from storage.objectViewer
       - artifactregistry.reader

4. デフォルトの IPv6 およびマシンタイプの設定を上書きするカスタム設定を gke-tpu-v6e-cluster ブロックに適用します。

   - id: gke-tpu-v6e-cluster
     source: modules/scheduler/gke-cluster
     use: [gke-tpu-v6e-net-0, workload_service_account]
     settings:
       enable_private_ipv6_google_access: false
       system_node_pool_disk_size_gb: $(vars.system_node_pool_disk_size_gb)
       system_node_pool_taints: []
       enable_private_endpoint: false # Allows access from authorized public IPs
       enable_pathways_for_tpus: $(vars.enable_pathways_for_tpus)
       enable_dataplane_v2: true
       configure_workload_identity_sa: true

5. ブループリントgke-tpu-v6e-advanced.yamlを使用し、--varsフラグを使用して必要な変数を渡すことで、Cluster Toolkit クラスターをデプロイします。

   ./gcluster deploy examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml \
       --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. Cluster Toolkit クラスターの作成が完了したら、Docker を設定します。

   # Configure docker for pulling images
   gcloud auth configure-docker gcr.io --quiet
   gcloud auth configure-docker ${REGION}-docker.pkg.dev --quiet

2. 後続のコマンドを簡素化するには、デフォルトのプロジェクト、クラスタ、ロケーションを構成します。

   # Configure gcluster Defaults
   ./gcluster job config set project ${PROJECT}
   ./gcluster job config set cluster ${CLUSTER_NAME}
   ./gcluster job config set location ${REGION}

3. モデルを Hugging Face 形式から MaxText 形式に変換し、Cloud Storage バケットに保存します。

   ./gcluster job submit \
     --name="gemma4-hf-to-mt" \
     --cluster="${CLUSTER_NAME}" \
     --project="${PROJECT}" \
     --location="${REGION}" \
     --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=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"

4. 変換ジョブのステータスを確認します。

   # Use the list command to check status
   ./gcluster job list \
       --cluster ${CLUSTER_NAME} \
       --project ${PROJECT} \
       --location ${REGION}
   
   # 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 \
       --cluster ${CLUSTER_NAME} \
       --project ${PROJECT} \
       --location ${REGION}

トレーニング ワークロードを開始する

変換プロセスが完了したら、RL トレーニング ワークロードを開始します。

./gcluster job submit \
  --name="gemma4-training" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET}/pathways/" \
  --env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}'); \
      python3 -c \"import pathlib, tpu_inference.layers.common.fused_moe_gmm as f; p = pathlib.Path(f.__file__); p.write_text(p.read_text().replace('onehot_moe_permute_threshold: int = 0,', 'onehot_moe_permute_threshold: int = 100000,'))\"; \
      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} \
      num_batches=50 \
      batch_size=8 \
      rollout_tensor_parallelism=2 \
      rollout_expert_parallelism=4 \
      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 \
      hbm_utilization_vllm=0.55 \
      remat_policy=full \
      async_scheduling=False \
      allow_split_physical_axes=true \
      ragged_gather_reduce_fallback=True \
      vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \
      vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', '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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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-training --main-only -f \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

トレーニング済みモデルを Hugging Face 形式に変換する

トレーニング ワークロードが完了したら、モデルを Hugging Face 形式に戻します。

./gcluster job submit \
  --name="gemma4-mt-to-hf" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

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

クリーンアップ

追加料金が発生しないようにするには、このチュートリアルで作成したリソースを削除します。

./gcluster destroy ${CLUSTER_NAME} --robust
gcloud storage rm -r gs://${GCS_BUCKET}
gcloud artifacts repositories delete ${REPOSITORY_NAME} --location=${REGION} --project=${PROJECT} --quiet

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

次のステップ

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