Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

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 Cloud ユーザーは無料トライアルをご利用いただける場合があります。

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

を参照してください。

始める前に

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

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

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

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

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

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

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

リポジトリのルートディレクトリに、ファイル名が 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'

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

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

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

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

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"

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

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

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

デフォルトの 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

ブループリント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 形式に変換する必要があります。

Cluster Toolkit クラスターの作成が完了したら、Docker を設定します。

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

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

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

モデルを 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"

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

# 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"