Run multi-host reinforcement learning training for Gemma 4 26B on TPU v6e

This tutorial shows you how to run multi-host reinforcement learning (RL) training on a Tensor Processing Unit (TPU) v6e-64 cluster by using MaxText and Cluster Toolkit. You use Cluster Toolkit to execute a multi-host training workload and export the results back to Hugging Face format for serving.

Objectives

• Install Cluster Toolkit and its dependencies.
• Install MaxText and its dependencies.
• Deploy a Cluster Toolkit cluster.
• Convert a Hugging Face model to MaxText format.
• Run an RL training workload on the TPU v6e cluster.
• Convert the fine-tuned model back to Hugging Face format for serving.

Costs

In this document, you use the following billable components of Google Cloud:

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

To generate a cost estimate based on your projected usage, use the pricing calculator.

New Google Cloud users might be eligible for a free trial.

When you finish the tasks that are described in this document, you can avoid continued billing by deleting the resources that you created. For more information, see Clean up.

Important: To complete this tutorial, you must have access to v6e-64 TPUs. This tutorial assumes that you already have a reservation for these TPUs in your Google Cloud project to ensure that you have the capacity available when you need it. To learn more about your options for reserving Cloud TPU capacity for a future date and time, see About Cloud TPU reservations.

Before you begin

You need a Hugging Face access token to use this tutorial. You can sign up for a free account at Hugging Face. After you have an account, generate an access token:

1. On the Welcome to Hugging Face page, click your account avatar and select Access tokens.
2. On the Access tokens page, click Create new token.
3. Select the Read token type and enter a name for your token.
4. Your access token is displayed. Save the token in a safe place.

• On the Hugging Face website, accept the license agreement for the model that you plan to train. This tutorial uses the model gemma4-26b.

To get the permissions that you need to complete this tutorial, ask your administrator to grant you the following IAM roles on your project:

• To complete this tutorial:
  • TPU Admin (roles/tpu.admin)
  • Service Account User (roles/iam.serviceAccountUser)
  • Compute Editor (roles/compute.editor)
• To prepare your MaxText container image:
  • Cloud Build Editor (roles/cloudbuild.builds.editor)
  • Artifact Registry Administrator (roles/artifactregistry.admin)
  • Storage Administrator (roles/storage.admin)
  • Service Usage Administrator (roles/serviceusage.serviceUsageAdmin)

For more information about granting roles, see Manage access to projects, folders, and organizations.

You might also be able to get the required permissions through custom roles or other predefined roles.

Set up your environment variables

Set up your environment variables:

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"

Replace the following:

• YOUR_PROJECT_ID: the ID of your Google Cloud project.
• YOUR_REGION: the region where you want to deploy your cluster.
• YOUR_ZONE: the zone where you want to deploy your cluster.
• YOUR_CLUSTER_NAME: the name of your Google Kubernetes Engine cluster.
• YOUR_REPOSITORY_NAME: the name of the Artifact Registry repository for your MaxText images.
• YOUR_BUCKET_NAME: a globally unique name for a Cloud Storage bucket.
• YOUR_RESERVATION_NAME: the name of your reservation.
• YOUR_HF_TOKEN: your Hugging Face access token.

Install Cluster Toolkit dependencies

To complete this tutorial from a Linux or macOS client or workstation, follow the relevant steps in Install dependencies in the Cluster Toolkit documentation.

If you're using Cloud Shell, then you can skip this section.

Install Cluster Toolkit

Install the prebuilt bundle for Cluster Toolkit by following the instructions at Install Cluster Toolkit.

Prepare your MaxText container image

To prepare your MaxText container image, including installing the required dependencies, complete the following steps:

1. Create a Cloud Storage bucket:

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

2. Create an Artifact Registry repository;

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

3. Create a file in the root directory of your repository with the filename cloudbuild.yaml and the following contents:

   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. Use Cloud Build to build your MaxText Docker image:

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

Create your Cluster Toolkit cluster

To create and deploy a Cluster Toolkit cluster with 64 v6e TPU chips, complete the following steps:

1. Create a custom Identity and Access Management (IAM) role, named gke.gcsfuse.profileUser:

   # 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. Create a Cloud Storage bucket:

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

3. By default, your cluster node pool service account doesn't have the required permissions to write to your Cloud Storage bucket. To allow the node pool service account to write to your Cloud Storage bucket, you must grant it the Storage Admin role. To grant this role, edit the file gke-tpu-v6e-advanced.yaml by updating the service-account module named node_pool_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            # Change from storage.objectViewer
       - artifactregistry.reader

4. Apply custom settings to the gke-tpu-v6e-cluster block that override the default IPv6 and machine type settings:

   - 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. Deploy your Cluster Toolkit cluster by using the blueprint gke-tpu-v6e-advanced.yaml and passing the required variables by using the --vars flag:

   ./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

Convert the model to MaxText format

To train the model in MaxText format, you must convert it from Hugging Face format to MaxText format.

1. After you've finished creating your Cluster Toolkit cluster, configure Docker:

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

2. To simplify subsequent commands, configure your default project, cluster, and location:

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

3. Convert the model from Hugging Face format to MaxText format and store it in your Cloud Storage bucket:

   ./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. Check the status of the conversion job:

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

Start the training workload

After the conversion process has completed, start the RL training workload:

./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'}}\""

Check the status of the training job:

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

Convert the trained model back into Hugging Face format

After the training workload has completed, convert the model back to Hugging Face format:

./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"

Check the status of the conversion job:

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

Clean up

To avoid incurring additional charges, delete the resources created during this tutorial:

./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}

What's next

• For more information about Cloud TPU, see Introduction to Cloud TPU.
• For architecture and configuration details for the v6e-64 TPU, see TPU v6e.
• For more information about Cluster Toolkit, see Cluster Toolkit overview.