Training Gemma-3 for Structured Mathematical Reasoning with Tunix GRPO, LoRA Adapters, and GSM8K Rewards

In this tutorial, we build an end-to-end GRPO training workflow that teaches Gemma-3 to reason through GSM8K math problems using Tunix, JAX, LoRA, and custom reward functions. We start by preparing the environment, authenticating with Hugging Face, loading the Gemma-3 model, and wrapping GSM8K examples into a prompt format that requires both structured reasoning and a final numeric answer. We then define reward functions that assess format adherence and mathematical correctness, attach LoRA adapters to keep training lightweight, evaluate the baseline model, and run GRPO to improve the policy via group-sampled generations. It provides a reinforcement learning tutorial in which we train only the adapter weights while keeping the workflow compact enough for a single-accelerator setup.

Installing Tunix and Configuring GRPO Training

Copy CodeCopiedUse a different Browser

import importlib.util, os, shutil as _sh

if importlib.util.find_spec("tunix") is None:

print("Installing Tunix + JAX ecosystem — this takes ~5-8 min…")

%pip install -q ipywidgets tensorboardX transformers grain nest_asyncio

%pip install -q datasets huggingface_hub "numpy>2"

%pip install -q tensorflow tensorflow_datasets

%pip install -q git+https://github.com/jax-ml/jax

%pip install -q git+https://github.com/google/tunix

%pip install -q git+https://github.com/google/qwix

%pip uninstall -q flax -y

%pip install -q git+https://github.com/google/flax

%pip uninstall -q wandb -y

print("\n\n Install done. The runtime will RESTART now.")

print(" After it restarts, RUN THIS CELL AGAIN to start training.\n")

os.kill(os.getpid(), 9)

import getpass, functools, json, re

import numpy as np

os.environ["WANDB_MODE"] = "disabled"

os.environ["TOKENIZERS_PARALLELISM"] = "false"

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

HF_TOKEN = os.environ.get("HF_TOKEN")

if not HF_TOKEN:

from google.colab import userdata

HF_TOKEN = userdata.get("HF_TOKEN")

except Exception:

HF_TOKEN = getpass.getpass("Hugging Face token (needs Gemma license access): ")

os.environ["HF_TOKEN"] = HF_TOKEN or ""

import nest_asyncio; nest_asyncio.apply()

import tensorflow as tf; tf.config.set_visible_devices([], "GPU")

import jax, jax.numpy as jnp, optax, grain, qwix

from flax import nnx

from orbax import checkpoint as ocp

from huggingface_hub import snapshot_download, login

from datasets import load_dataset

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

if HF_TOKEN:

login(token=HF_TOKEN)

devices = jax.devices()

print("JAX backend :", jax.default_backend(), "|", len(devices), "device(s):", devices)

IS_GPU = _sh.which("nvidia-smi") is not None

if IS_GPU and jax.default_backend() == "cpu":

print(" GPU runtime but JAX sees CPU only. Re-run `pip install -U \"jax[cuda12]\"` "

"and restart, or switch to a TPU runtime.")

MODEL_ID = "google/gemma-3-1b-it"

RANK, ALPHA = 32, 32.0

MAX_PROMPT_LENGTH = 256

TOTAL_GENERATION_STEPS = 512

NUM_GENERATIONS = 2

NUM_ITERATIONS = 1

BETA = 0.08

EPSILON = 0.2

TEMPERATURE, TOP_P, TOP_K = 0.9, 1.0, 50

MAX_STEPS = 100

TRAIN_LIMIT = MAX_STEPS

NUM_TEST = 16

LEARNING_RATE, B1, B2, WEIGHT_DECAY = 3e-6, 0.9, 0.99, 0.1

WARMUP_STEPS, MAX_GRAD_NORM = int(0.1 * MAX_STEPS), 0.1

CKPT_DIR, TB_DIR = "/content/ckpts/", "/content/tb/grpo"

N = jax.device_count()

MESH = [(N, 1), ("fsdp", "tp")]

mesh = jax.make_mesh(*MESH, axis_types=(jax.sharding.AxisType.Auto,) * 2)

We set up the complete Colab environment by installing Tunix, JAX, Flax, Qwix, TensorFlow, datasets, and the supporting notebook utilities needed for GRPO training. We configure authentication, disable unnecessary logging paths, keep TensorFlow away from the accelerator, and verify that JAX can see the available TPU or GPU devices. We also define the core training hyperparameters, LoRA settings, generation limits, checkpoint paths, and device mesh that control how the model trains across the available hardware.

Formatting GSM8K Prompts and Reward Functions

Copy CodeCopiedUse a different Browser

reasoning_start, reasoning_end = "", ""

solution_start, solution_end = "", ""

SYSTEM_PROMPT = (f"You are given a problem. First, think about the problem and provide your "

f"reasoning between {reasoning_start} and {reasoning_end}. Then give the final "

f"answer (just one number) between {solution_start} and {solution_end}.")

TEMPLATE = ("user\n{system_prompt}\n\n{question}\nmodel\n")

def extract_hash_answer(text):

return text.split("####")[1].strip() if "####" in text else None

gsm = load_dataset("openai/gsm8k", "main")

def build_grain(split, limit):

rows = [{"question": r["question"], "answer": r["answer"]}

for r in split.select(range(min(limit, len(split))))]

return (grain.MapDataset.source(rows).shuffle(seed=42).map(

lambda x: {"prompts": TEMPLATE.format(system_prompt=SYSTEM_PROMPT, question=x["question"]),

"question": x["question"],

"answer": extract_hash_answer(x["answer"])}))

train_dataset = build_grain(gsm["train"], TRAIN_LIMIT).batch(1)[:MAX_STEPS]

val_dataset = None

test_rows = [{"question": r["question"], "answer": extract_hash_answer(r["answer"])}

for r in gsm["test"].select(range(NUM_TEST))]

print(f"Train batches: {len(train_dataset)} | Test examples: {len(test_rows)}")

match_format = re.compile(

rf"^[\s]{{0,}}{reasoning_start}.+?{reasoning_end}.*?{solution_start}(.+?){solution_end}[\s]{{0,}}$",

flags=re.MULTILINE | re.DOTALL)

match_numbers = re.compile(rf"{solution_start}.*?([\d\.]{{1,}})", flags=re.MULTILINE | re.DOTALL)

def match_format_exactly(prompts, completions, **kw):

return [3.0 if match_format.search(c) else 0.0 for c in completions]

def match_format_approximately(prompts, completions, **kw):

for c in completions:

s += 0.5 if c.count(reasoning_start) == 1 else -0.5

s += 0.5 if c.count(reasoning_end) == 1 else -0.5

s += 0.5 if c.count(solution_start) == 1 else -0.5

s += 0.5 if c.count(solution_end) == 1 else -0.5

out.append(s)

def check_answer(prompts, completions, answer, **kw):

guesses = [(m.group(1) if (m := match_format.search(c)) else None) for c in completions]

scores = []

for g, a in zip(guesses, answer):

if g is None:

scores.append(0.0); continue

if g == a: scores.append(3.0)

elif g.strip() == a.strip(): scores.append(1.5)

r = float(g) / float(a)

scores.append(0.5 if 0.9 <= r <= 1.1 else 0.25 if 0.8 <= r <= 1.2 else -1.0)

except Exception:

scores.append(-0.5)

return scores

def check_numbers(prompts, completions, answer, **kw):

guesses = [(m.group(1) if (m := match_numbers.search(c)) else None) for c in completions]

scores = []

for g, a in zip(guesses, answer):

scores.append(1.5 if float(g.strip()) == float(a.strip()) else 0.0)

except Exception:

scores.append(0.0)

return scores

REWARD_FNS = [match_format_exactly, match_format_approximately, check_answer, check_numbers]

We define the structured reasoning format that asks the model to place its reasoning inside reasoning tags and its final numeric answer inside answer tags. We load GSM8K from Hugging Face, extract the ground-truth answers, and convert each math problem into the prompt format expected by the GRPO rollout pipeline. We then create reward functions that score exact format matching, approximate tag usage, answer correctness, and fallback numeric extraction so the model receives useful feedback from multiple signals.

Loading Gemma-3 and Attaching LoRA Adapters

Copy CodeCopiedUse a different Browser

print(f"Downloading {MODEL_ID} …")

local_model_path = snapshot_download(repo_id=MODEL_ID, ignore_patterns=["*.pth"])

model_config = (gemma_lib.ModelConfig.gemma3_270m() if "270m" in MODEL_ID

else gemma_lib.ModelConfig.gemma3_1b_it())

base_model = params_safetensors_lib.create_model_from_safe_tensors(

local_model_path, model_config, mesh)

TOK = "/content/tokenizer_gemma3.model"

if not os.path.exists(TOK):

cand = os.path.join(local_model_path, "tokenizer.model")

if os.path.exists(cand):

shutil.copy(cand, TOK)

import urllib.request

urllib.request.urlretrieve(

"https://storage.googleapis.com/gemma-data/tokenizers/tokenizer_gemma3.model", TOK)

tokenizer = tokenizer_lib.Tokenizer(tokenizer_path=TOK)

EOS_TOKENS = []

gcfg = os.path.join(local_model_path, "generation_config.json")

if os.path.exists(gcfg):

EOS_TOKENS = list(json.load(open(gcfg)).get("eos_token_id", []) or [])

if tokenizer.eos_id() not in EOS_TOKENS:

EOS_TOKENS.append(tokenizer.eos_id())

print("EOS tokens:", EOS_TOKENS)

def apply_lora(base, mesh):

provider = qwix.LoraProvider(

module_path=".*q_einsum|.*kv_einsum|.*gate_proj|.*down_proj|.*up_proj|.*attn_vec_einsum",

rank=RANK, alpha=ALPHA)

m = qwix.apply_lora_to_model(base, provider, **base.get_model_input())

st = nnx.state(m)

st = jax.lax.with_sharding_constraint(st, nnx.get_partition_spec(st))

nnx.update(m, st)

lora_policy = apply_lora(base_model, mesh)

def make_sampler(model):

return sampler_lib.Sampler(

transformer=model, 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))

def evaluate(rows, model, tag, n_show=2):

sampler = make_sampler(model)

correct = fmt_ok = 0

for i in range(0, len(rows), 4):

chunk = rows[i:i + 4]

qs = [r["question"] for r in chunk]

prompts = [TEMPLATE.format(system_prompt=SYSTEM_PROMPT, question=q) for q in qs]

outs = sampler(input_strings=prompts, max_generation_steps=TOTAL_GENERATION_STEPS,

temperature=None, top_k=1, top_p=None, echo=False, eos_tokens=EOS_TOKENS).text

for j, (o, r) in enumerate(zip(outs, chunk)):

m = match_numbers.search(o); g = m.group(1) if m else None

correct += int(float(g) == float(r["answer"]))

except Exception:

fmt_ok += int(match_format.search(o) is not None)

if i + j < n_show:

print(f" Q: {qs[j][:70]}…\n gold={r['answer']} pred={g}\n out: {o[:150]}…\n")

print(f"[{tag}] accuracy = {100*correct/len(rows):.1f}% format = {100*fmt_ok/len(rows):.1f}%\n")

print("\n════════ BASELINE (before GRPO) ════════")

evaluate(test_rows, lora_policy, "baseline")

We download the selected Gemma-3 checkpoint, create the base model from safetensors, and prepare the tokenizer and EOS token list for generation. We attach LoRA adapters to the attention and MLP projection modules, enabling us to train a lightweight policy without updating the full model weights. We also build a sampler-based evaluation function and run a baseline test before GRPO training to measure the model’s initial accuracy and adherence to the format.

Configuring the Tunix RL Cluster and Training

Copy CodeCopiedUse a different Browser

schedule = 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)

optimizer = optax.chain(

optax.clip_by_global_norm(MAX_GRAD_NORM),

optax.adamw(learning_rate=schedule, b1=B1, b2=B2, weight_decay=WEIGHT_DECAY))

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=10**9,

max_steps=MAX_STEPS,

mini_batch_size=1, train_micro_batch_size=1,

metrics_logging_options=metrics_logger.MetricsLoggerOptions(

log_dir=TB_DIR, flush_every_n_steps=10),

checkpoint_root_directory=CKPT_DIR,

checkpointing_options=ocp.CheckpointManagerOptions(save_interval_steps=50, max_to_keep=2)),

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))

rl_cluster = rl_cluster_lib.RLCluster(

actor=lora_policy, reference=base_model, tokenizer=tokenizer, cluster_config=cluster_config)

grpo_trainer = GRPOLearner(

rl_cluster=rl_cluster, reward_fns=REWARD_FNS,

algo_config=GRPOConfig(num_generations=NUM_GENERATIONS, num_iterations=NUM_ITERATIONS,

beta=BETA, epsilon=EPSILON))

%load_ext tensorboard

%tensorboard --logdir $TB_DIR --port=0

print("\n════════ TRAINING (GRPO) ════════")

grpo_trainer.train(train_dataset, val_dataset)

We create the learning-rate schedule, optimizer, gradient clipping, and AdamW configuration that guide the LoRA policy updates during training. We define the Tunix RL cluster with actor, reference, and rollout roles, then connect the rollout configuration to the tokenizer, generation limits, sampling temperature, and EOS tokens. We initialize the GRPO learner, launch TensorBoard for live metrics, and start the GRPO training loop on the prepared GSM8K batches.

Evaluating and Exporting the Trained Model

Copy CodeCopiedUse a different Browser

print("\n════════ AFTER GRPO ════════")

evaluate(test_rows, lora_policy, "after-grpo")

out_dir = "/content/gemma3-grpo-merged"

if os.path.exists(out_dir): shutil.rmtree(out_dir)

gemma_params.save_lora_merged_model_as_safetensors(

local_model_path=local_model_path, output_dir=out_dir,

lora_model=lora_policy, rank=RANK, alpha=ALPHA)

print(f"\n Merged model saved to {out_dir}")

except Exception as e:

print(f"\n(Skipped merged export: {e})")

print("\nDone. Checkpoints are in", CKPT_DIR)

We evaluate the trained LoRA policy after GRPO to compare its math accuracy and response format against the baseline result. We then try to export the LoRA-merged Gemma-3 model into a Hugging Face safetensors directory so the trained checkpoint can be reused outside the notebook. We finish by reporting the output checkpoint path, giving us a complete workflow from training setup to post-training evaluation and optional model export.

In conclusion, we completed a full GRPO fine-tuning loop for mathematical reasoning with Gemma-3, from dataset preparation and reward design to LoRA-based policy training and post-training evaluation. We saw how Tunix organizes the actor, reference model, rollout engine, optimizer, checkpoints, and metrics into a reusable reinforcement learning pipeline. We also compared the model before and after GRPO to observe whether its answer format and numeric accuracy improve during training. Finally, we optionally exported the merged LoRA model, providing a complete path from raw GSM8K examples to a trained, reasoning-oriented Gemma-3 checkpoint.

Check out the FULL CODES HERE. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.

Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us

The post Training Gemma-3 for Structured Mathematical Reasoning with Tunix GRPO, LoRA Adapters, and GSM8K Rewards appeared first on MarkTechPost.