Reinforcement Learning with gemma4-e4b on Multi-Host TPUs

Reinforcement Learning with gemma4-e4b on Multi-Host TPUs#

This tutorial provides step-by-step instructions for setting up the environment and training the gemma4-e4b model with GRPO on the OpenMathInstruct-2 dataset on a Cloud TPU v6e (Trillium) GKE cluster using a v6e-32 (4x8) slice with Cluster Toolkit.

Prerequisites#

Before starting, ensure you have:

  • Access to a Google Cloud Project with TPU quotas.

  • A Hugging Face account with an access token for downloading models (the google/gemma-4-E4B and google/gemma-4-E4B-it repositories are gated; request access before proceeding).

  • Permissions for Google Artifact Registry (Artifact Registry Writer role).

  • Cluster Toolkit installed and configured. Follow Running MaxText with Cluster Toolkit for gcluster setup.

  • A Pathways-ready GKE cluster configured for Cluster Toolkit, including healthy Kueue and JobSet components (see create a GKE cluster with Pathways and Cluster Toolkit documentation).

  • Docker installed and configured for sudoless use. Follow the steps to configure sudoless Docker.

Setup Environment Variables#

Set up the following environment variables to configure your training run. Replace placeholders with your actual values.

# Your GCP project ID.
# If you've already set it in your local config, you can retrieve it via:
# gcloud config get-value project
export PROJECT_ID=<PROJECT_ID>

# The name of your GKE cluster.
export CLUSTER_NAME=<CLUSTER_NAME>

# The GCP location of your GKE cluster.
export ZONE=<ZONE> # e.g., 'us-central1' or 'us-central1-a'

# Use a GCS bucket you own to store logs and checkpoints.
export BASE_OUTPUT_DIRECTORY=<GCS_BUCKET> # e.g., gs://my-bucket/maxtext-runs

# An arbitrary string to identify this specific run.
export RUN_NAME=<RUN_NAME>

Authenticate with Hugging Face#

To download the gemma4-e4b model checkpoint from Hugging Face, you need to authenticate using your Hugging Face account credentials. Run the following command and follow the prompts to log in:

hf auth login

Get Your MaxText Compatible Model Checkpoint#

Option 1: Using an existing MaxText checkpoint#

If you already have a MaxText-compatible model checkpoint, simply set the following environment variable and move on to the next section.

export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items

Option 2: Converting from a Hugging Face checkpoint#

Refer to Hugging Face to MaxText to convert a Hugging Face checkpoint to MaxText format. After conversion finishes, set MAXTEXT_CKPT_PATH to the converted MaxText checkpoint path.

export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items

Note (Gemma 4 E4B specifics):

  • For the gemma4-e4b model, you must run the conversion with scan_layers=False (the gemma4_small decoder block is incompatible with nn.scan); the resulting unscanned checkpoint matches the scan_layers=False setting used for the RL run below.

  • This RL recipe fine-tunes the base model (google/gemma-4-E4B), not the instruction-tuned default. Pass --hf_model_path=google/gemma-4-E4B explicitly when running to_maxtext — otherwise MaxText defaults to HF_IDS[gemma4-e4b] = google/gemma-4-E4B-it.

Chat Template Configuration#

Unlike an instruction-tuned tokenizer, the google/gemma-4-E4B base tokenizer does not ship with a chat template, so the RL run must supply one explicitly. The run_gemma4_e4b_rl.sh script points at two files bundled with the repo (and therefore baked into your Docker image):

  • data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json — a stripped-down data template that injects only the system prompt and question (no turn markers). This avoids the doubled <start_of_turn> delimiters that would occur if the default gsm8k_rl.json (which bakes literal Gemma turn markers into the message content) were combined with the tokenizer’s Jinja chat template.

  • chat_template_path=maxtext/examples/chat_templates/gemma-3-27b-chat_template.json — the Gemma 3 chat template, wrapped in a JSON object with a chat_template key. It is applied by the tokenizer as the single source of truth for turn-marker formatting.

Both files are already included under src/maxtext/examples/chat_templates/, so no additional setup is required. If you want to customize the prompt formatting, edit these files (or point the config at your own) before building the Docker image.

Run RL Workload#

Build and Upload MaxText Docker Image#

For instructions on building and uploading the MaxText Docker image with post-training dependencies, please refer to the official documentation.

Submit your workload#

# The Docker image you pushed in the previous step
export DOCKER_IMAGE=<IMAGE_NAME>

# Run the RL training script on your cluster
run_tutorial maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh

Note: The run_gemma4_e4b_rl.sh script pins the Pathways component images to specific versions via the gcluster --pathways-server-image, --pathways-worker-image, and --pathways-proxy-server-image flags (set through the PATHWAYS_SERVER_IMAGE and PATHWAYS_PROXY_SERVER_IMAGE variables at the top of the script). PATHWAYS_SERVER_IMAGE is used for both the Pathways resource-manager server and the workers (the reference config uses the same image for both). Update these variables if you need a different Pathways release.

Monitor your workload#

To monitor your job’s progress, you can use gcluster or kubectl to check the JobSet status and stream logs directly:

# Check job status with Cluster Toolkit
gcluster job list

# Stream logs with Cluster Toolkit (specify --main-only=false for Pathways workloads)
gcluster job logs ${RUN_NAME?} --main-only=false

# Alternatively, check JobSet status with kubectl
kubectl get jobset -l gcluster.google.com/workload=${RUN_NAME?}

# List pods (use jobset-name to select both head and worker pods in Pathways)
kubectl get pods -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?}

# Stream logs with kubectl
kubectl logs -f -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} --all-containers=true --max-log-requests=64

Alternatively, gcluster job submit provides a link to the Google Cloud Console to view your workload logs. Follow the link to view logs and monitor your workload’s progress in the Cloud Console.

Monitor RL Metrics#

During RL training, you can monitor key metrics to track model convergence, reward trends, and hardware performance.

To enable Tunix-managed metrics measurement, set enable_tunix_perf_metrics to true in RL configurations. Note that this flag is already set to True by default for this tutorial workload. When enabled, Tunix automatically collects and uploads these metrics to TensorBoard.

For a complete list of collected metrics, see the Tunix Metrics Documentation. Key metrics to monitor include:

  • Model Quality & Reward Metrics:

    • rewards/mean: The average reward across the batch (crucial for tracking learning progress).

    • score/mean: The average raw score from the reward model before applying the KL penalty.

  • Rollout & Generation Metrics:

    • rollout_time: How long each rollout step takes.

    • completions/mean_length: The average token length of generated completions.

    • actor_dequeue_time: The time spent waiting for data from the rollout workers (relevant when async rollout is enabled).

  • Performance & Efficiency Metrics:

    • step_time_sec: The execution time for a single training step.

Convert Checkpoint to Hugging Face Format#

Refer to MaxText to Hugging Face to convert a MaxText checkpoint back to Hugging Face format.

Note (Gemma 4 E4B specifics): Because this recipe fine-tunes the base model, pass --hf_model_path=google/gemma-4-E4B to to_huggingface so the exported checkpoint bundles the base tokenizer. Without it, to_huggingface sources the tokenizer from HF_IDS[gemma4-e4b] = google/gemma-4-E4B-it (the instruction-tuned model). Also keep scan_layers=False, since the gemma4_small decoder block is not compatible with scanned layers.