GRPO Demo#
This tutorial demonstrates training the Gemma 3 1B-IT model on the GSM8K math reasoning benchmark using Group Relative Policy Optimization (GRPO). GRPO can enhance your model’s problem-solving skills on mathematical word problems, coding problems, etc.
GRPO is an RL algorithm designed to enhance the reasoning abilities of LLMs. It is a variant of Proximal Policy Optimization (PPO) that reduces memory usage by eliminating the need for a separate value function model. GRPO works by generating multiple responses for a given prompt, evaluating these responses using a reward model, and then calculating a relative advantage based on the group’s performance to update the policy.
In this tutorial we use a v6e-1 TPU for Gemma3-1B-it. Let’s get started!
Note that the setup below is for the Gemma3-1B-IT model only. If you want to use another model (say, Qwen2.5), you may need to change the setup (for example, tokenizer, chat template, reward function, etc.).
Install necessary libraries: RESTART AFTER INSTALL FOR COLAB#
import importlib.util
if importlib.util.find_spec('tensorflow') is None:
print("Installing required packages...")
%pip install -q dotenv
%pip install -q kagglehub
%pip install -q ipywidgets
%pip install -q tensorflow
%pip install -q tensorflow_datasets
%pip install -q tensorboardX
%pip install -q transformers
%pip install -q grain
%pip install -q git+https://github.com/jax-ml/jax
%pip install git+https://github.com/google/tunix
%pip install git+https://github.com/google/qwix
%pip uninstall -q flax -y
%pip install git+https://github.com/google/flax
%pip install -q huggingface_hub
%pip install -q datasets
%pip install -q 'numpy>2'
Logging into services#
Also, there is a bug where the code stalls with initializing rl_trainer/grpo_cluster with wandb. Doing this early because it fixes the bug sometimes (but not always - still bugs a lot)
import os
import kagglehub
try:
from google.colab import userdata
USE_COLAB = True
%pip uninstall -q wandb -y # wandb is glitchy with tunix in colab
os.environ["HF_TOKEN"] = userdata.get("HF_TOKEN")
os.environ["KAGGLE_USERNAME"] = userdata.get("KAGGLE_USERNAME")
os.environ["KAGGLE_KEY"] = userdata.get("KAGGLE_KEY")
except:
USE_COLAB = False
from dotenv import load_dotenv
load_dotenv()
print("Using env vars to login")
import nest_asyncio
nest_asyncio.apply()
print("nest_asyncio applied")
# Only using wandb on TPU VM because it has strange bugs on Colab
!pip install -q wandb
import wandb
# Check if WANDB_API_KEY is set before logging in
if "WANDB_API_KEY" in os.environ and os.environ["WANDB_API_KEY"]:
wandb.login(key=os.environ["WANDB_API_KEY"])
else:
print("WANDB_API_KEY not found. Skipping wandb login.")
if "KAGGLE_USERNAME" not in os.environ or "KAGGLE_KEY" not in os.environ:
kagglehub.login()
if "HF_TOKEN" in os.environ and os.environ["HF_TOKEN"]:
hf_token = os.environ["HF_TOKEN"]
!hf auth login --token "$hf_token"
else:
print("HF_TOKEN not found. Skipping Hugging Face login.")
Imports#
import functools
from pprint import pprint
import re
import sys
import csv
import json
import shutil
from flax import nnx
import grain
import humanize
from huggingface_hub import snapshot_download
import jax
import jax.numpy as jnp
import kagglehub
import numpy as np
import optax
from orbax import checkpoint as ocp
from pathlib import Path
import qwix
import tensorflow_datasets as tfds
from tqdm.auto import tqdm
from tunix.generate import sampler as sampler_lib
from tunix.generate import tokenizer_adapter as tokenizer_lib
from tunix.models.gemma3 import model as gemma_lib
from tunix.models.gemma3 import params_safetensors as params_safetensors_lib
from tunix.models.gemma3 import params as gemma_params
from tunix.rl import rl_cluster as rl_cluster_lib
from tunix.rl.grpo.grpo_learner import GRPOConfig, GRPOLearner
from tunix.rl.rollout import base_rollout
from tunix.sft import metrics_logger
Hyperparameters#
Let’s define the configuration we are going to use. Note that this is by no means a “perfect” set of hyperparameters. To get good results, you might have to train the model for longer.
# ====== Model ======
MODEL_ID = "google/gemma-3-1b-it"
GEMMA_TOKENIZER_PATH = "gs://gemma-data/tokenizers/tokenizer_gemma3.model"
# ====== Data ======
TRAIN_DATA_DIR = "./data/train"
TEST_DATA_DIR = "./data/test"
TRAIN_FRACTION = .9
# ====== LoRA ======
RANK = 64
ALPHA = 64.0
# ====== Sharding ======
# Adjust mesh based on your TPU memory and model size.
NUM_TPUS = len(jax.devices())
if NUM_TPUS == 8:
MESH_COUNTS = (1, 4)
elif NUM_TPUS == 1:
MESH_COUNTS = (1, 1)
else:
raise ValueError(f"Unsupported number of TPUs: {NUM_TPUS}")
MESH = [
MESH_COUNTS,
("fsdp", "tp"),
]
# ====== GRPO ======
# === Generation during GRPO training ===
MAX_PROMPT_LENGTH = 256
TOTAL_GENERATION_STEPS = 768
# Important to keep a high-ish temperature for varied, diverse responses during
# training.
TEMPERATURE = 0.9
TOP_P = 1.0
TOP_K = 50
# The number of times the policy generates multiple responses for a given prompt
# within a single training step. This corresponds to `G` in Algorithm 1 in the
# paper. The "group" in GRPO comes from here.
NUM_GENERATIONS = 2
# === other GRPO configs ===
# The number of iterations per batch (𝜇 in GRPO algo 1).
NUM_ITERATIONS = 1
# The coefficient for the KL divergence penalty (đť›˝) in the GRPO loss function.
# Important to keep a high enough value for this, otherwise, the KL divergence
# can increase unchecked.
BETA = 0.08
# Epsilon value for clipping (𝜀 in GRPO loss in paper). Similar to PPO, for
# stable updates.
EPSILON = 0.2
# ====== Training ======
TRAIN_MICRO_BATCH_SIZE = 1
# Increase `NUM_BATCHES` and `MAX_STEPS` for better results.
NUM_BATCHES = 3738
# Keep `NUM_TEST_BATCHES` low so that evaluation runs quickly. It can be
# increased to a max. of 330 (if batch size is 4).
NUM_TEST_BATCHES = 64
EVAL_EVERY_N_STEPS = 64 # this doesn't matter if `TRAIN_FRACTION = 1.0`.
NUM_EPOCHS = 1 # can potentially train for more epochs
# Number of training steps.
MAX_STEPS = int(NUM_BATCHES * NUM_ITERATIONS * TRAIN_FRACTION * NUM_EPOCHS)
# === AdamW, warmup, cosine scheduler ===
LEARNING_RATE = 3e-6
B1 = 0.9
B2 = 0.99
WEIGHT_DECAY = 0.1
# == Cosine decay with warmup scheduler ==
# Linearly increase learning rate from 0. to 5e-6 in the first 10% training
# steps, and then gradually decrease the learning rate to 0 using cosine
# scheduler.
WARMUP_STEPS = 0.1 * MAX_STEPS
# == Grad clipping ==
# Grad clipping to prevent large gradients. Found this
# important to keep KL divergence in check.
MAX_GRAD_NORM = 0.1
# Checkpoint saving
INTERMEDIATE_CKPT_DIR = "/tmp/content/intermediate_ckpt/"
CKPT_DIR = "/tmp/content/ckpts/"
SAVE_INTERVAL_STEPS = 500
MAX_TO_KEEP = 4
# ====== Inference ======
GENERATION_CONFIGS = {
# greedy search
"greedy": {"temperature": None, "top_k": 1, "top_p": None},
# some randomness
"standard": {"temperature": 0.7, "top_k": 50, "top_p": 0.95},
# liberal
"liberal": {"temperature": 0.85, "top_k": 2000, "top_p": 1.0},
}
Utility functions#
def show_hbm_usage():
"""Displays memory usage per device."""
fmt_size = functools.partial(humanize.naturalsize, binary=True)
for d in jax.local_devices():
stats = d.memory_stats()
used = stats["bytes_in_use"]
limit = stats["bytes_limit"]
print(f"Using {fmt_size(used)} / {fmt_size(limit)} ({used/limit:%}) on {d}")
Data preprocessing#
First, let’s define some special tokens. We instruct the model to first reason
between the <reasoning> and </reasoning> tokens. After
reasoning, we expect it to provide the answer between the <answer> and
</answer> tokens.
reasoning_start = "<reasoning>"
reasoning_end = "</reasoning>"
solution_start = "<answer>"
solution_end = "</answer>"
SYSTEM_PROMPT = f"""You are given a problem. First, think about the problem \
and provide your reasoning. Place it between {reasoning_start} and \
{reasoning_end}. Then, provide the final answer (i.e., just one numerical \
value) between {solution_start} and {solution_end}."""
TEMPLATE = """<start_of_turn>user
{system_prompt}
{question}<end_of_turn>
<start_of_turn>model
"""
We use OpenAI’s GSM8K dataset, which comprises grade school math word problems.
def extract_hash_answer(text: str) -> str | None:
if "####" not in text:
return None
return text.split("####")[1].strip()
def _load_from_tfds(data_dir: str, split: str):
import tensorflow_datasets.text.gsm8k
return tfds.data_source(
"gsm8k",
split=split,
data_dir=data_dir,
builder_kwargs={"file_format": tfds.core.FileFormat.ARRAY_RECORD},
download=True,
)
def download_kaggle_dataset(target_dir="./data/gsm8k"):
os.makedirs(target_dir, exist_ok=True)
src = kagglehub.dataset_download("thedevastator/grade-school-math-8k-q-a")
src = Path(src)
dst = Path(target_dir)
for csv_file in src.glob("*.csv"): # match all CSV files
shutil.copy2(csv_file, dst / csv_file.name)
print(f"Copied {csv_file.name} → {dst/csv_file.name}")
return target_dir
def get_dataset(data_dir, split="train", source="tfds") -> grain.MapDataset:
# Download data
if not os.path.exists(data_dir):
os.makedirs(data_dir)
if source == "tfds":
import tensorflow_datasets.text.gsm8k
data = tfds.data_source(
"gsm8k",
split=split,
data_dir=data_dir,
builder_kwargs={"file_format": tfds.core.FileFormat.ARRAY_RECORD},
download=True,
)
elif source == "kaggle":
kaggle_dir = download_kaggle_dataset(data_dir)
file_name = "main_" + split + ".csv"
csv_path = os.path.join(kaggle_dir, file_name) # adjust filename if needed
data = []
with open(csv_path, newline="", encoding="utf-8") as csvfile:
reader = csv.DictReader(csvfile)
for row in reader:
data.append({
"question": row["question"],
"answer": row["answer"],
})
else:
raise ValueError(f"Unknown source: {source}")
def _as_text(v):
return v if isinstance(v, str) else v.decode("utf-8")
dataset = (
grain.MapDataset.source(data)
.shuffle(seed=42)
.map(
lambda x: {
# passed to model forward pass
"prompts": TEMPLATE.format(
system_prompt=SYSTEM_PROMPT,
question=_as_text(x["question"]),
),
# passed to reward functions
"question": _as_text(x["question"]),
# passed to reward functions
"answer": extract_hash_answer(_as_text(x["answer"])),
}
)
)
return dataset
We split the dataset set into train and test sets as usual.
source = input("Choose data source [tfds/kaggle]: ").strip().lower() # SHOULD I LEAVE THIS OR KEEP IT SET TO KAGGLE
if source not in ("tfds", "kaggle"):
print("Invalid choice. Defaulting to 'tfds'.")
source = "tfds"
print(f"Using data source: {source}")
dataset = get_dataset(TRAIN_DATA_DIR, "train", source).batch(TRAIN_MICRO_BATCH_SIZE)[
:NUM_BATCHES
]
if TRAIN_FRACTION == 1.0:
train_dataset = dataset.repeat(NUM_EPOCHS)
val_dataset = None
else:
train_dataset = dataset[: int(len(dataset) * TRAIN_FRACTION)]
train_dataset = train_dataset.repeat(NUM_EPOCHS)
val_dataset = dataset[int(len(dataset) * TRAIN_FRACTION) :].repeat(NUM_EPOCHS)
test_dataset = get_dataset(TEST_DATA_DIR, "test", source).batch(TRAIN_MICRO_BATCH_SIZE)[
:NUM_TEST_BATCHES
]
dataset_lengths = (
len(train_dataset),
len(val_dataset) if val_dataset is not None else 0,
len(test_dataset),
)
print(f"dataset contains {dataset_lengths} of batches")
Let’s see how one batch of the training dataset looks like!
for ele in train_dataset[:1]:
pprint(ele)
Load the policy model and the reference model#
The policy model is the model which is actually trained and whose weights are updated. The reference model is the model with which we compute KL divergence. This is to ensure that the policy updates are not huge and that it does not deviate too much from the reference model.
Typically, the reference model is the base model, and the policy model is the same base model, but with LoRA parameters. Only the LoRA parameters are updated.
Note: We perform full precision (fp32) training. You can, however, leverage Qwix for QAT.
To load the model, you need to be on Kaggle and need to have agreed to the Gemma license here.
ignore_patterns = [
"*.pth", # Ignore PyTorch .pth weight files
]
print(f"Downloading {MODEL_ID} from Hugging Face...")
local_model_path = snapshot_download(
repo_id=MODEL_ID, ignore_patterns=ignore_patterns
)
print(f"Model successfully downloaded to: {local_model_path}")
EOS_TOKENS = []
generation_config_path = os.path.join(local_model_path, "generation_config.json")
if os.path.exists(generation_config_path):
with open(generation_config_path, "r") as f:
generation_configs = json.load(f)
EOS_TOKENS = generation_configs.get("eos_token_id", [])
print(f"Using EOS token IDs: {EOS_TOKENS}")
This code snippet serves as a workaround to re-save the pre-trained model checkpoint from Kaggle into a local format that is compatible with the Flax NNX library. Because the original checkpoint has parameter names and tensor structures that don’t match the target NNX model architecture, it cannot be loaded directly.
We first load the original weights into a temporary model instance, then extract and re-save the model’s state into a new, properly formatted local checkpoint, which can then be successfully loaded by the final sharded NNX model.
MODEL_CP_PATH = local_model_path
model_config = None
if "gemma-3-270m" in MODEL_ID:
model_config = gemma_lib.ModelConfig.gemma3_270m()
elif "gemma-3-1b" in MODEL_ID:
model_config = gemma_lib.ModelConfig.gemma3_1b_it()
else:
raise ValueError(f"Unknown model id: {MODEL_ID}")
mesh = jax.make_mesh(*MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0]))
with mesh:
gemma3 = params_safetensors_lib.create_model_from_safe_tensors(
MODEL_CP_PATH, (model_config), mesh
)
nnx.display(gemma3)
LoRA Application#
The get_lora_model function takes the base model and applies LoRA layers to it. It uses a LoraProvider to select specific layers (like attention and MLP layers) to be adapted. The resulting LoRA-infused model is then sharded and updated to ensure it’s ready for distributed training.
def get_lora_model(base_model, mesh):
lora_provider = qwix.LoraProvider(
module_path=(
".*q_einsum|.*kv_einsum|.*gate_proj|.*down_proj|.*up_proj|"
".*attn_vec_einsum"
),
rank=RANK,
alpha=ALPHA,
)
model_input = base_model.get_model_input()
lora_model = qwix.apply_lora_to_model(
base_model, lora_provider, **model_input
)
with mesh:
state = nnx.state(lora_model)
pspecs = nnx.get_partition_spec(state)
sharded_state = jax.lax.with_sharding_constraint(state, pspecs)
nnx.update(lora_model, sharded_state)
return lora_model
Now we load reference and policy Gemma models using the Flax NNX library and display their structures.
# Policy model
lora_policy = get_lora_model(gemma3, mesh=mesh)
nnx.display(lora_policy)
tokenizer = tokenizer_lib.Tokenizer(tokenizer_path=GEMMA_TOKENIZER_PATH)
if tokenizer.eos_id() not in EOS_TOKENS:
EOS_TOKENS.append(tokenizer.eos_id())
print(f"Using EOS token IDs: {EOS_TOKENS}")
Define reward functions#
We define four reward functions:
reward if the format of the output exactly matches the instruction given in
TEMPLATE;reward if the format of the output approximately matches the instruction given in
TEMPLATE;reward if the answer is correct/partially correct;
Sometimes, the text between
<answer>,</answer>might not be one number. So, we extract the number, and reward the model if the answer is correct.
The reward functions are inspired from here.
First off, let’s define a RegEx for checking whether the format matches.
match_format = re.compile(
rf"^[\s]{{0,}}"
rf"{reasoning_start}.+?{reasoning_end}.*?"
rf"{solution_start}(.+?){solution_end}"
rf"[\s]{{0,}}$",
flags=re.MULTILINE | re.DOTALL,
)
example = f"{reasoning_start}Let me think!{reasoning_end}{solution_start}2{solution_end}"
print("Example:", example)
print("Match:\t", match_format.search(example))
Give the model a reward of 3 points if the format matches exactly.
def match_format_exactly(prompts, completions, **kwargs):
return [
0 if match_format.search(response) is None else 3.0
for response in completions
]
We also reward the model if the format of the output matches partially.
def match_format_approximately(prompts, completions, **kwargs):
scores = []
for completion in completions:
score = 0
response = completion
# Count how many keywords are seen - we penalize if too many!
# If we see 1, then plus some points!
score += 0.5 if response.count(reasoning_start) == 1 else -0.5
score += 0.5 if response.find(reasoning_start) == 0 else -0.5
score += 0.5 if response.count(reasoning_end) == 1 else -0.5
score += 0.5 if response.count(solution_start) == 1 else -0.5
score += 0.5 if response.count(solution_end) == 1 else -0.5
scores.append(score)
return scores
Reward the model if the answer is correct. A reward is also given if the answer does not match exactly, i.e., based on how close the answer is to the correct value.
def check_answer(prompts, completions, answer, **kwargs):
responses = completions
extracted_responses = [
guess.group(1) if r is not None and (guess := match_format.search(r)) is not None else None
for r in responses
]
scores = []
assert len(extracted_responses) == len(
answer
), f"{extracted_responses} and {answer} have mismatching length"
for guess, true_answer in zip(extracted_responses, answer):
score = 0
if guess is None:
scores.append(0)
continue
# Correct answer gets 3 points!
if guess == true_answer:
score += 3.0
# Match if spaces are seen
elif guess.strip() == true_answer.strip():
score += 1.5
else:
# We also reward it if the answer is close via ratios!
# Ie if the answer is within some range, reward it!
try:
ratio = float(guess) / float(true_answer)
if ratio >= 0.9 and ratio <= 1.1:
score += 0.5
elif ratio >= 0.8 and ratio <= 1.2:
score += 0.25
else:
score -= 1.0 # Penalize wrong answers
except:
score -= 0.5 # Penalize
scores.append(score)
return scores
Sometimes, the text between <answer> and </answer> might not be one
number; it can be a sentence. So, we extract the number and compare the answer.
match_numbers = re.compile(
rf"{solution_start}.*?([\d\.]{{1,}})", flags=re.MULTILINE | re.DOTALL
)
match_numbers.findall(f"{solution_start} 0.34 {solution_end}")
def check_numbers(prompts, completions, answer, **kwargs):
question = kwargs["question"]
responses = completions
extracted_responses = [
guess.group(1) if (guess := match_numbers.search(r)) is not None else None
for r in responses
]
scores = []
print("START ============================")
print(f"Question:\t{question[0]}")
print(f"Answer:\t{answer[0]}")
print(f"Response:\t{responses[0]}")
print(f"Extracted:\t{extracted_responses[0]}")
print("END ==============================")
for guess, true_answer in zip(extracted_responses, answer):
if guess is None:
scores.append(0)
continue
# Convert to numbers
try:
true_answer = float(true_answer.strip())
guess = float(guess.strip())
scores.append(1.5 if guess == true_answer else 0.0)
except:
scores.append(0)
continue
return scores
Evaluate#
Before we train the model, let’s evaluate the model on the test set so we can see the improvement post training.
We evaluate it in two ways:
Quantitative
Answer Accuracy: percentage of samples for which the model predicts the correct final numerical answer
Answer (Partial) Accuracy: percentage of samples for which the model predicts a final numerical answer such that the `model answer / answer` ratio lies between 0.9 and 1.1.
Format Accuracy: percentage of samples for which the model outputs the correct format, i.e., reasoning between the reasoning special tokens, and the final answer between the `<start_answer>`, `<end_answer>` tokens.
Qualitative
We’ll also print outputs for a few given questions so that we can compare the generated output later.
We define a helper function to generate an answer, given a prompt.
def generate(
question, sampler, temperature=0.7, top_k=50, top_p=0.95, seed=None
):
"""Given prompt, generates text."""
if isinstance(question, str):
input_batch = [
TEMPLATE.format(
system_prompt=SYSTEM_PROMPT,
question=question,
),
]
else:
input_batch = [
TEMPLATE.format(
system_prompt=SYSTEM_PROMPT,
question=q,
)
for q in question
]
out_data = sampler(
input_strings=input_batch,
max_generation_steps=768,
temperature=temperature,
top_k=top_k,
top_p=top_p,
echo=False,
seed=seed if seed is not None else None,
eos_tokens=EOS_TOKENS,
)
output = out_data.text
if isinstance(question, str):
return output[0]
return output
Another helper function for evaluation.
def evaluate(
dataset,
sampler,
temperature=0.7,
top_k=50,
top_p=0.95,
num_passes=1,
corr_lst=False,
make_lst=False,
):
"""Computes accuracy and percentage of outputs matching the format."""
response_lst = []
corr = 0
partially_corr = 0
corr_format = 0
total = 0
for batch in tqdm(dataset):
answers = batch["answer"]
questions = batch["question"]
multiple_call_responses = [[] for _ in range(len(questions))]
for p in range(num_passes):
responses = generate(
questions, sampler, temperature, top_k, top_p, seed=p
)
for idx, response in enumerate(responses):
multiple_call_responses[idx].append(response)
print(f"Question:\t{questions[idx]}")
print(f"Correct Answer:\t{answers[idx]}")
print(f"Response:\t{response}")
print("-" * 50)
for question, multiple_call_response, answer in zip(
questions, multiple_call_responses, answers
):
# check answer
corr_ctr_per_question = 0
partially_corr_per_question = 0
corr_format_per_question = 0
for response in multiple_call_response:
extracted_response = (
guess.group(1)
if (guess := match_numbers.search(response)) is not None
else "-1000000"
)
try:
if float(extracted_response.strip()) == float(answer.strip()):
corr_ctr_per_question += 1
ratio = float(extracted_response.strip()) / float(answer.strip())
if ratio >= 0.9 and ratio <= 1.1:
partially_corr_per_question += 1
except:
print("SKIPPED")
# check format
if match_format.search(response) is not None:
corr_format_per_question += 1
if (
corr_ctr_per_question > 0
and partially_corr_per_question > 0
and corr_format_per_question > 0
):
break
if corr_ctr_per_question > 0:
corr += 1
if corr_lst and make_lst:
response_lst.append((question, answer, multiple_call_response))
else:
if not corr_lst and make_lst:
response_lst.append((question, answer, multiple_call_response))
if partially_corr_per_question > 0:
partially_corr += 1
if corr_format_per_question > 0:
corr_format += 1
total += 1
if total % 10 == 0:
print(
f"===> {corr=}, {total=}, {corr / total * 100=}, "
f"{partially_corr / total * 100=}, {corr_format / total * 100=}"
)
to_return = (
corr,
total,
corr / total * 100,
partially_corr / total * 100,
corr_format / total * 100,
)
if make_lst:
return to_return, response_lst
return to_return
Now let’s see how the original model does on the test set. You can see the percentages of the mode outputs that are fully correct, partially correct and just correct in format. The following step might take couple of minutes to finish.
# The evaluation might take up to couple of minutes to finish.
sampler = sampler_lib.Sampler(
transformer=lora_policy,
tokenizer=tokenizer,
cache_config=sampler_lib.CacheConfig(
cache_size=MAX_PROMPT_LENGTH + TOTAL_GENERATION_STEPS + 256,
num_layers=model_config.num_layers,
num_kv_heads=model_config.num_kv_heads,
head_dim=model_config.head_dim,
),
)
num_correct, total, accuracy, partial_accuracy, format_accuracy = evaluate(
test_dataset,
sampler,
**GENERATION_CONFIGS["greedy"],
)
print(
f"{num_correct=}, {total=}, {accuracy=}%, {partial_accuracy=}%,"
f" {format_accuracy=}%"
)
Train#
Let’s set up all the configs first - checkpointing, metric logging and training. We then train the model.
# Ckpt saving
checkpointing_options = ocp.CheckpointManagerOptions(
save_interval_steps=SAVE_INTERVAL_STEPS, max_to_keep=MAX_TO_KEEP
)
# Metrics logger
metrics_logging_options = metrics_logger.MetricsLoggerOptions(
log_dir="/tmp/content/tmp/tensorboard/grpo", flush_every_n_steps=20
)
# Logs
%load_ext tensorboard
%tensorboard --logdir /tmp/content/tmp/tensorboard/grpo --port=0
# Optimizer, learning rate scheduler, gradient clipping
optimizer = optax.adamw(
learning_rate=optax.schedules.warmup_cosine_decay_schedule(
init_value=0.0,
peak_value=LEARNING_RATE,
warmup_steps=WARMUP_STEPS,
decay_steps=MAX_STEPS,
end_value=0.0,
),
b1=B1,
b2=B2,
weight_decay=WEIGHT_DECAY,
)
if MAX_GRAD_NORM is not None:
optimizer = optax.chain(
optax.clip_by_global_norm(max_norm=MAX_GRAD_NORM),
optimizer,
)
# Training config
cluster_config = rl_cluster_lib.ClusterConfig(
role_to_mesh={
rl_cluster_lib.Role.ACTOR: mesh,
rl_cluster_lib.Role.REFERENCE: mesh,
rl_cluster_lib.Role.ROLLOUT: mesh,
},
rollout_engine='vanilla',
offload_to_cpu=False,
training_config=rl_cluster_lib.RLTrainingConfig(
actor_optimizer=optimizer,
eval_every_n_steps=EVAL_EVERY_N_STEPS,
max_steps=MAX_STEPS,
mini_batch_size=TRAIN_MICRO_BATCH_SIZE,
train_micro_batch_size=TRAIN_MICRO_BATCH_SIZE,
# metrics logging
metrics_logging_options=metrics_logging_options,
# checkpoint saving
checkpoint_root_directory=CKPT_DIR,
checkpointing_options=checkpointing_options,
),
rollout_config=base_rollout.RolloutConfig(
max_tokens_to_generate=TOTAL_GENERATION_STEPS,
max_prompt_length=MAX_PROMPT_LENGTH,
kv_cache_size=MAX_PROMPT_LENGTH + TOTAL_GENERATION_STEPS + 256,
temperature=TEMPERATURE,
top_p=TOP_P,
top_k=TOP_K,
eos_tokens=EOS_TOKENS,
),
)
grpo_config = GRPOConfig(
num_generations=NUM_GENERATIONS,
num_iterations=NUM_ITERATIONS,
beta=BETA,
epsilon=EPSILON,
)
Setting Up the GRPO Trainer#
Now we initialize our system for training. First, we create an RLCluster instance, which brings together the policy model (actor), a reference model (reference), and a tokenizer. Our actor is a trainable LoRA model, while the reference is a fixed base model that we use to guide the training.
We then create a GRPOLearner, the specialized trainer that uses a list of reward functions to evaluate and optimize the model’s output, completing the RL training setup.
Tunix trainers are integrated with Weights & Biases and Tensorboard to help you visualize the training progress. You can choose how you want to use it:
Colab: Unfortunately, WandB doesn’t work nicely with Colab so we disabled it, instead use the tensorboards rendered above to monitor.
Option 1 (Type 1): If you’re running a quick experiment or just testing things out, choose this. It creates a temporary, private dashboard right in your browser without requiring you to log in or create an account.
Option 2 (Type 2): If you have an existing W&B account and want to save your project’s history to your personal dashboard, choose this. You’ll be prompted to enter your API key or log in.
# RL cluster
rl_cluster = rl_cluster_lib.RLCluster(
actor=lora_policy,
reference=gemma3,
tokenizer=tokenizer,
cluster_config=cluster_config,
)
# GRPO Trainer
grpo_trainer = GRPOLearner(
rl_cluster=rl_cluster,
reward_fns=[
match_format_exactly,
match_format_approximately,
check_answer,
check_numbers,
],
algo_config=grpo_config,
)
Training#
The first couple of training step might take up to 5 minutes to finish. Please be patient. If you experience long training steps, e.g. >10 minutes per step, please open a bug. Really appreciated!
grpo_trainer.train(train_dataset, val_dataset)
if not USE_COLAB: wandb.init()
Final Evaluation#
Let’s evaluate our finetuned model!
# Load checkpoint first.
trained_ckpt_path = os.path.join(
CKPT_DIR, "actor", str(MAX_STEPS), "model_params"
)
abs_params = jax.tree.map(
lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype),
nnx.state(lora_policy, nnx.LoRAParam),
)
checkpointer = ocp.StandardCheckpointer()
trained_lora_params = checkpointer.restore(trained_ckpt_path, target=abs_params)
nnx.update(
lora_policy,
jax.tree.map(
lambda a, b: b,
nnx.state(lora_policy, nnx.LoRAParam),
trained_lora_params,
),
)
# The evaluation might take up to couple of minutes to finish. Please be patient.
sampler = sampler_lib.Sampler(
transformer=lora_policy,
tokenizer=tokenizer,
cache_config=sampler_lib.CacheConfig(
cache_size=MAX_PROMPT_LENGTH + TOTAL_GENERATION_STEPS + 256,
num_layers=model_config.num_layers,
num_kv_heads=model_config.num_kv_heads,
head_dim=model_config.head_dim,
),
)
num_correct, total, accuracy, partial_accuracy, format_accuracy = evaluate(
test_dataset,
sampler,
**GENERATION_CONFIGS["greedy"],
)
print(
f"{num_correct=}, {total=}, {accuracy=}%, {partial_accuracy=}%,"
f" {format_accuracy=}%"
)
With sufficient training, you should see that the percentages of correct model outputs have clearly gone up, which means our training worked.
Export Merged Lora Weights (Huggingface Format)#
output_dir = f"./{MODEL_ID}-lora"
if USE_COLAB:
output_dir = f"/tmp/content/{MODEL_ID}-lora"
if os.path.exists(output_dir):
shutil.rmtree(output_dir)
os.makedirs(output_dir)
print(f"Saving merged LoRA model to {output_dir}")
# Use the save_lora_merged_model_as_safetensors function
gemma_params.save_lora_merged_model_as_safetensors(
local_model_path=local_model_path,
output_dir=output_dir,
lora_model=lora_policy,
rank=RANK,
alpha=ALPHA,
)
print("\n" + "="*60)
print("Model saved successfully!")
print(f"Output directory: {output_dir}")
print("="*60)
print("\nSaved files:")
for f in os.listdir(output_dir):
size = os.path.getsize(os.path.join(output_dir, f)) / (1024 * 1024)
print(f" {f:<30} {size:>10.2f} MB")
# For Colab: Download as zip
if USE_COLAB:
from google.colab import files
shutil.make_archive(output_dir, 'zip', output_dir)
files.download(f"{output_dir}.zip")