Skip to content

rl_gemma4_e4b.md

Latest commit

 

History

History
152 lines (102 loc) · 8.55 KB

rl_gemma4_e4b.md

File metadata and controls

152 lines (102 loc) · 8.55 KB
Raw

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.

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).
  • Prerequisites for XPK installed (follow official documentation).
  • A Pathways-ready GKE cluster (see create GKE cluster).
  • 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

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 CLOUD_IMAGE_NAME=<IMAGE_NAME>
export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_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 xpk --server-image and --proxy-server-image flags (set through the PATHWAYS_SERVER_IMAGE and PATHWAYS_PROXY_SERVER_IMAGE variables at the top of the script). The --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 kubectl to check the Jobset status and stream logs directly from the pods.

kubectl get jobset -n default ${WORKLOAD_NAME}

# List pods to find the specific name
kubectl get pods | grep ${WORKLOAD_NAME}

# stream the logs from the running pod (replace <POD_NAME> with the name you found)
kubectl logs -f <POD_NAME>

Alternatively, after running the bash script, you will also get 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.