Skip to content

Commit 735aea3

Browse files
dasotodasoto
committed
Add tutorial for RL GRPO gemma4-e4b on v6e 4x8
1 parent c5b135a commit 735aea3

7 files changed

Lines changed: 497 additions & 1 deletion

File tree

‎docs/tutorials/posttraining/rl_gemma4_e4b.md‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
<!--
2+
Copyright 2023-2026 Google LLC
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
https://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
-->
16+
17+
# Reinforcement Learning with gemma4-e4b on Multi-Host TPUs
18+
19+
This tutorial provides step-by-step instructions for setting up the environment
20+
and training the gemma4-e4b model with GRPO on the [OpenMathInstruct-2 dataset](https://huggingface.co/datasets/nvidia/OpenMathInstruct-2) on a Cloud TPU v6e (Trillium) GKE cluster using a `v6e-32` (4x8) slice.
21+
22+
## Prerequisites
23+
24+
Before starting, ensure you have:
25+
26+
- Access to a Google Cloud Project with TPU quotas.
27+
- 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).
28+
- Permissions for Google Artifact Registry (Artifact Registry Writer role).
29+
- Prerequisites for XPK installed (follow [official documentation](https://github.com/AI-Hypercomputer/xpk/blob/main/docs/installation.md#1-prerequisites)).
30+
- A Pathways-ready GKE cluster (see [create GKE cluster](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster)).
31+
- **Docker** installed and configured for sudoless use. Follow the steps to [configure sudoless Docker](https://docs.docker.com/engine/install/linux-postinstall/).
32+
33+
## Build and Upload MaxText Docker Image
34+
35+
For instructions on building and uploading the MaxText Docker image with post-training dependencies, please refer to the [official documentation](../build_maxtext.md).
36+
37+
## Setup Environment Variables
38+
39+
Set up the following environment variables to configure your training run. Replace
40+
placeholders with your actual values.
41+
42+
```bash
43+
# Your GCP project ID.
44+
# If you've already set it in your local config, you can retrieve it via:
45+
# gcloud config get-value project
46+
export PROJECT_ID=<PROJECT_ID>
47+
48+
# The name of your GKE cluster.
49+
export CLUSTER_NAME=<CLUSTER_NAME>
50+
51+
# The GCP location of your GKE cluster.
52+
export ZONE=<ZONE> # e.g., 'us-central1' or 'us-central1-a'
53+
54+
# Use a GCS bucket you own to store logs and checkpoints.
55+
export BASE_OUTPUT_DIRECTORY=<GCS_BUCKET> # e.g., gs://my-bucket/maxtext-runs
56+
57+
# The Docker image you pushed in the previous step
58+
export CLOUD_IMAGE_NAME=<IMAGE_NAME>
59+
export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}"
60+
```
61+
62+
## Clone MaxText Repository
63+
64+
If you haven't already, clone the MaxText repository to your local machine:
65+
66+
```bash
67+
git clone https://github.com/AI-Hypercomputer/maxtext.git
68+
cd maxtext
69+
```
70+
71+
## Authenticate with Hugging Face
72+
73+
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:
74+
75+
```bash
76+
hf auth login
77+
```
78+
79+
## Get Your MaxText Compatible Model Checkpoint
80+
81+
### Option 1: Using an existing MaxText checkpoint
82+
83+
If you already have a MaxText-compatible model checkpoint, simply set the
84+
following environment variable and move on to the next section.
85+
86+
```bash
87+
export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items
88+
```
89+
90+
### Option 2: Converting from a Hugging Face checkpoint
91+
92+
> **Note:** The conversion downloads the Hugging Face safetensors for the **base** model `google/gemma-4-E4B`. The script passes `--hf_model_path=google/gemma-4-E4B` explicitly, because MaxText otherwise defaults to `HF_IDS[gemma4-e4b]` = `google/gemma-4-E4B-it` (the instruction-tuned model), which we do not want for this base-model RL run. Please verify you have sufficient free disk space before running the conversion script.
93+
>
94+
> **Note:** The `gemma4-e4b` model uses the `gemma4_small` decoder block, which does **not** support scanned layers during conversion. The conversion script therefore runs with `scan_layers=False`, producing an unscanned checkpoint that matches the `scan_layers=False` setting used for the RL run.
95+
96+
```bash
97+
# Optional: If you run out of disk space when downloading Hugging Face safetensors,
98+
# customize your "HF_HOME" to redirect the cache to a larger or mounted disk (e.g., on a TPU VM).
99+
# export HF_HOME="/dev/shm/huggingface_tmp"
100+
101+
# Create and activate a virtual environment
102+
uv venv --python 3.12 --seed tpu_venv
103+
source tpu_venv/bin/activate
104+
uv pip install -e .[tpu] --resolution=lowest
105+
106+
# Run the conversion script to convert the Hugging Face checkpoint to MaxText format
107+
bash scripts/run_gemma4_e4b_hf_to_maxtext.sh
108+
109+
# Deactivate the virtual environment
110+
deactivate
111+
rm -rf tpu_venv
112+
```
113+
114+
## Chat Template Configuration
115+
116+
Unlike an instruction-tuned tokenizer, the `google/gemma-4-E4B` tokenizer does **not** ship with a chat template, so the RL run must supply one explicitly. The [scripts/run_gemma4_e4b_rl.sh](../../../scripts/run_gemma4_e4b_rl.sh) script points at two files that are bundled with the repo (and therefore baked into your Docker image):
117+
118+
- `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.
119+
- `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.
120+
121+
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.
122+
123+
## Run RL Workload
124+
125+
### Submit your workload
126+
127+
```bash
128+
# Create and activate a virtual environment
129+
uv venv --python 3.12 --seed runner_venv
130+
source runner_venv/bin/activate
131+
uv pip install -e .[runner] --resolution=lowest
132+
133+
# Run the RL training script on your cluster
134+
bash scripts/run_gemma4_e4b_rl.sh
135+
136+
# Deactivate the virtual environment
137+
deactivate
138+
rm -rf runner_venv
139+
```
140+
141+
> **Note:** The `scripts/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.
142+
143+
### Monitor your workload
144+
145+
To monitor your job's progress, you can use `kubectl` to check the `Jobset` status and stream logs directly from the pods.
146+
147+
```bash
148+
kubectl get jobset -n default ${WORKLOAD_NAME}
149+
150+
# List pods to find the specific name
151+
kubectl get pods | grep ${WORKLOAD_NAME}
152+
153+
# stream the logs from the running pod (replace <POD_NAME> with the name you found)
154+
kubectl logs -f <POD_NAME>
155+
```
156+
157+
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.
158+
159+
### Monitor RL Metrics
160+
161+
During RL training, you can monitor key metrics to track model convergence, reward trends, and hardware performance.
162+
163+
To enable Tunix-managed metrics measurement, set `enable_tunix_perf_metrics` to `true` in `src/maxtext/configs/post_train/rl.yml`. Note that this flag is already set to `True` by default in the [scripts/run_gemma4_e4b_rl.sh](../../../scripts/run_gemma4_e4b_rl.sh) script for this tutorial workload. When enabled, Tunix automatically collects and uploads these metrics to TensorBoard.
164+
165+
For a complete list of collected metrics, see the [Tunix Metrics Documentation](https://tunix.readthedocs.io/en/latest/metrics.html). Key metrics to monitor include:
166+
167+
- **Model Quality & Reward Metrics:**
168+
- `rewards/mean`: The average reward across the batch (crucial for tracking learning progress).
169+
- `score/mean`: The average raw score from the reward model before applying the KL penalty.
170+
- **Rollout & Generation Metrics:**
171+
- `rollout_time`: How long each rollout step takes.
172+
- `completions/mean_length`: The average token length of generated completions.
173+
- `actor_dequeue_time`: The time spent waiting for data from the rollout workers (relevant when async rollout is enabled).
174+
- **Performance & Efficiency Metrics:**
175+
- `step_time_sec`: The execution time for a single training step.
176+
177+
## Convert Checkpoint to Hugging Face Format
178+
179+
After training, you may want to convert your MaxText checkpoint back to Hugging Face format. Use the following script to perform the conversion:
180+
181+
> **Note:** The script passes `--hf_model_path=google/gemma-4-E4B` so the exported checkpoint bundles the **base** model's tokenizer. Without it, `to_huggingface` sources the tokenizer from `HF_IDS[gemma4-e4b]` = `google/gemma-4-E4B-it` (the instruction-tuned model).
182+
183+
```bash
184+
# Create and activate a virtual environment
185+
uv venv --python 3.12 --seed tpu_venv
186+
source tpu_venv/bin/activate
187+
uv pip install -e .[tpu] --resolution=lowest
188+
189+
# Run the conversion script to convert the MaxText checkpoint back to Hugging Face format
190+
bash scripts/run_gemma4_e4b_maxtext_to_hf.sh
191+
192+
# Deactivate the virtual environment
193+
deactivate
194+
rm -rf tpu_venv
195+
```

‎scripts/run_gemma4_e4b_hf_to_maxtext.sh‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
#!/bin/bash
2+
3+
# This script converts a gemma4-e4b model checkpoint from the Hugging Face
4+
# format to the MaxText (i.e., Orbax) format. It requires the BASE_OUTPUT_DIRECTORY
5+
# environment variable to be set to a GCS path where the converted checkpoint will be stored.
6+
7+
set -e
8+
9+
# --- Environment Setup ---
10+
if ! pip show maxtext &> /dev/null; then
11+
echo "maxtext not found in the environment. Please install it by running:"
12+
echo "uv pip install -e .[tpu] --resolution=lowest"
13+
exit 1
14+
fi
15+
16+
# --- Environment Variables ---
17+
export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # GCS bucket path for outputs (e.g., gs://my-bucket/outputs)
18+
19+
# --- Variable Validation ---
20+
if [ -z "$BASE_OUTPUT_DIRECTORY" ]; then
21+
echo "Error: BASE_OUTPUT_DIRECTORY is not set. Please set it in the script or as an environment variable."
22+
exit 1
23+
fi
24+
25+
# Install torch for conversion
26+
echo "Installing torch for checkpoint conversion..."
27+
python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
28+
29+
# Define the output path for the converted checkpoint
30+
CONVERTED_CKPT_BASE_DIR="${BASE_OUTPUT_DIRECTORY}/checkpoints/gemma4-e4b-converted"
31+
echo "Converted checkpoint will be saved to: ${CONVERTED_CKPT_BASE_DIR}"
32+
33+
# Run the conversion script.
34+
# NOTE: scan_layers must be False for the gemma4_small decoder block (E2B/E4B);
35+
# scanned-layer conversion is not supported for this block.
36+
# NOTE: --hf_model_path pins the conversion to the BASE model (google/gemma-4-E4B).
37+
# Without it, MaxText defaults to HF_IDS[gemma4-e4b] = google/gemma-4-E4B-it (the
38+
# instruction-tuned model), which is not what we want for this base-model RL run.
39+
python3 -m maxtext.checkpoint_conversion.to_maxtext \
40+
src/maxtext/configs/base.yml \
41+
model_name=gemma4-e4b \
42+
base_output_directory="${CONVERTED_CKPT_BASE_DIR}" \
43+
scan_layers=False \
44+
weight_dtype=bfloat16 hardware=cpu skip_jax_distributed_system=True \
45+
checkpoint_storage_use_ocdbt=False checkpoint_storage_use_zarr3=False \
46+
--hf_model_path=google/gemma-4-E4B \
47+
--eager_load_method=safetensors
48+
49+
# Set MAXTEXT_CKPT_PATH to the newly created checkpoint path
50+
export MAXTEXT_CKPT_PATH="${CONVERTED_CKPT_BASE_DIR}/0/items"
51+
echo "Conversion complete. Using checkpoint at: ${MAXTEXT_CKPT_PATH}"

‎scripts/run_gemma4_e4b_maxtext_to_hf.sh‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
#!/bin/bash
2+
3+
# This script converts a gemma4-e4b model checkpoint from the MaxText
4+
# format (orbax) to the Hugging Face format. It requires the MAXTEXT_CKPT_PATH and
5+
# BASE_OUTPUT_DIRECTORY environment variables to be set.
6+
7+
set -e
8+
9+
# --- Environment Setup ---
10+
if ! pip show maxtext &> /dev/null; then
11+
echo "maxtext not found in the environment. Please install it by running:"
12+
echo "uv pip install -e .[tpu] --resolution=lowest"
13+
exit 1
14+
fi
15+
16+
# --- Environment Variables ---
17+
export MAXTEXT_CKPT_PATH="${MAXTEXT_CKPT_PATH:-}" # GCS path to the MaxText checkpoint to convert
18+
export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # GCS bucket path for storing the converted HF checkpoint
19+
20+
# --- Variable Validation ---
21+
if [ -z "$BASE_OUTPUT_DIRECTORY" ]; then
22+
echo "Error: BASE_OUTPUT_DIRECTORY is not set. Please set it in the script or as an environment variable."
23+
exit 1
24+
fi
25+
26+
if [ -z "$MAXTEXT_CKPT_PATH" ]; then
27+
echo "Error: MAXTEXT_CKPT_PATH is not set. Please set it in the script or as an environment variable."
28+
exit 1
29+
fi
30+
31+
# Install torch for conversion
32+
echo "Installing torch for checkpoint conversion..."
33+
python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
34+
35+
# Define the output path for the converted checkpoint
36+
HF_CKPT_OUTPUT_DIR="${BASE_OUTPUT_DIRECTORY}/checkpoints/gemma4-e4b-hf-converted"
37+
echo "Converted Hugging Face checkpoint will be saved to: ${HF_CKPT_OUTPUT_DIR}"
38+
39+
# Run the conversion script.
40+
# NOTE: scan_layers must be False for the gemma4_small decoder block (E2B/E4B);
41+
# scanned-layer conversion is not supported for this block.
42+
# NOTE: --hf_model_path pins the bundled tokenizer to the BASE model
43+
# (google/gemma-4-E4B). Without it, to_huggingface sources the tokenizer from
44+
# HF_IDS[gemma4-e4b] = google/gemma-4-E4B-it (the instruction-tuned model).
45+
python -m maxtext.checkpoint_conversion.to_huggingface \
46+
src/maxtext/configs/base.yml \
47+
model_name=gemma4-e4b \
48+
load_parameters_path="${MAXTEXT_CKPT_PATH}" \
49+
base_output_directory="${HF_CKPT_OUTPUT_DIR}" \
50+
scan_layers=False \
51+
weight_dtype=bfloat16 hardware=cpu skip_jax_distributed_system=True \
52+
--hf_model_path=google/gemma-4-E4B
53+
54+
echo "Conversion to Hugging Face format complete. Checkpoint saved to: ${HF_CKPT_OUTPUT_DIR}"

0 commit comments

Comments
 (0)