|

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

ByRicardo
✅

In this tutorial, we construct an end-to-end GRPO coaching workflow that teaches Gemma-3 to purpose via GSM8K math issues utilizing Tunix, JAX, LoRA, and customized reward capabilities. We begin by making ready the surroundings, authenticating with Hugging Face, loading the Gemma-3 mannequin, and wrapping GSM8K examples right into a immediate format that requires each structured reasoning and a last numeric reply. We then outline reward capabilities that assess format adherence and mathematical correctness, connect LoRA adapters to maintain coaching light-weight, consider the baseline mannequin, and run GRPO to enhance the coverage through group-sampled generations. It offers a reinforcement studying tutorial during which we prepare solely the adapter weights whereas preserving the workflow compact sufficient for a single-accelerator setup.

Installing Tunix and Configuring GRPO Training

Copy Code
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 set up -q ipywidgets tensorboardX transformers grain nest_asyncio
   %pip set up -q datasets huggingface_hub "numpy>2"
   %pip set up -q tensorflow tensorflow_datasets
   %pip set up -q git+https://github.com/jax-ml/jax
   %pip set up -q git+https://github.com/google/tunix
   %pip set up -q git+https://github.com/google/qwix
   %pip uninstall -q flax -y
   %pip set up -q git+https://github.com/google/flax
   %pip uninstall -q wandb -y
   print("nn✅ Install executed. The runtime will RESTART now.")
   print("👉  After it restarts, RUN THIS CELL AGAIN to begin coaching.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:
   attempt:
       from google.colab import userdata
       HF_TOKEN = userdata.get("HF_TOKEN")
   besides Exception:
       HF_TOKEN = getpass.getpass("Hugging Face token (wants Gemma license entry): ")
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.fashions.gemma3 import mannequin as gemma_lib
from tunix.fashions.gemma3 import params_safetensors as params_safetensors_lib
from tunix.fashions.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)
units = jax.units()
print("JAX backend :", jax.default_backend(), "|", len(units), "system(s):", units)
IS_GPU = _sh.which("nvidia-smi") will not be None
if IS_GPU and jax.default_backend() == "cpu":
   print("⚠  GPU runtime however JAX sees CPU solely. Re-run `pip set up -U "jax[cuda12]"` "
         "and restart, or change 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 material/ckpts/", "/content material/tb/grpo"
N = jax.device_count()
MESH = [(N, 1), ("fsdp", "tp")]
mesh = jax.make_mesh(*MESH, axis_types=(jax.sharding.AxisType.Auto,) * 2)

We arrange the entire Colab surroundings by putting in Tunix, JAX, Flax, Qwix, TensorFlow, datasets, and the supporting pocket book utilities wanted for GRPO coaching. We configure authentication, disable pointless logging paths, hold TensorFlow away from the accelerator, and confirm that JAX can see the obtainable TPU or GPU units. We additionally outline the core coaching hyperparameters, LoRA settings, era limits, checkpoint paths, and system mesh that management how the mannequin trains throughout the obtainable {hardware}.

Formatting GSM8K Prompts and Reward Functions

Copy Code
reasoning_start, reasoning_end = "<reasoning>", "</reasoning>"
solution_start,  solution_end  = "<reply>",   "</reply>"
SYSTEM_PROMPT = (f"You are given an issue. First, take into consideration the issue and present your "
                f"reasoning between {reasoning_start} and {reasoning_end}. Then give the ultimate "
                f"reply (only one quantity) between {solution_start} and {solution_end}.")
TEMPLATE = ("<start_of_turn>usern{system_prompt}nn{query}<end_of_turn>n<start_of_turn>modeln")
def extract_hash_answer(textual content):
   return textual content.break up("####")[1].strip() if "####" in textual content else None
gsm = load_dataset("openai/gsm8k", "important")
def build_grain(break up, restrict):
   rows = [{"question": r["question"], "reply": r["answer"]}
           for r in break up.choose(vary(min(restrict, len(break up))))]
   return (grain.MapDataset.supply(rows).shuffle(seed=42).map(
       lambda x: {"prompts":  TEMPLATE.format(system_prompt=SYSTEM_PROMPT, query=x["question"]),
                  "query": x["question"],
                  "reply":   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"], "reply": extract_hash_answer(r["answer"])}
                for r in gsm["test"].choose(vary(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):
   out = []
   for c in completions:
       s = 0.0
       s += 0.5 if c.rely(reasoning_start) == 1 else -0.5
       s += 0.5 if c.rely(reasoning_end)   == 1 else -0.5
       s += 0.5 if c.rely(solution_start)  == 1 else -0.5
       s += 0.5 if c.rely(solution_end)    == 1 else -0.5
       out.append(s)
   return out
def check_answer(prompts, completions, reply, **kw):
   guesses = [(m.group(1) if (m := match_format.search(c)) else None) for c in completions]
   scores = []
   for g, a in zip(guesses, reply):
       if g is None:
           scores.append(0.0); proceed
       if g == a:               scores.append(3.0)
       elif g.strip() == a.strip(): scores.append(1.5)
       else:
           attempt:
               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)
           besides Exception:
               scores.append(-0.5)
   return scores
def check_numbers(prompts, completions, reply, **kw):
   guesses = [(m.group(1) if (m := match_numbers.search(c)) else None) for c in completions]
   scores = []
   for g, a in zip(guesses, reply):
       attempt:
           scores.append(1.5 if float(g.strip()) == float(a.strip()) else 0.0)
       besides Exception:
           scores.append(0.0)
   return scores
REWARD_FNS = [match_format_exactly, match_format_approximately, check_answer, check_numbers]

We outline the structured reasoning format that asks the mannequin to position its reasoning inside reasoning tags and its last numeric reply inside reply tags. We load GSM8K from Hugging Face, extract the ground-truth solutions, and convert every math drawback into the immediate format anticipated by the GRPO rollout pipeline. We then create reward capabilities that rating precise format matching, approximate tag utilization, reply correctness, and fallback numeric extraction so the mannequin receives helpful suggestions from a number of alerts.

Loading Gemma-3 and Attaching LoRA Adapters

Copy Code
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())
with mesh:
   base_model = params_safetensors_lib.create_model_from_safe_tensors(
       local_model_path, model_config, mesh)
TOK = "/content material/tokenizer_gemma3.mannequin"
if not os.path.exists(TOK):
   cand = os.path.be part of(local_model_path, "tokenizer.mannequin")
   if os.path.exists(cand):
       shutil.copy(cand, TOK)
   else:
       import urllib.request
       urllib.request.urlretrieve(
           "https://storage.googleapis.com/gemma-data/tokenizers/tokenizer_gemma3.mannequin", TOK)
tokenizer = tokenizer_lib.Tokenizer(tokenizer_path=TOK)
EOS_TOKENS = []
gcfg = os.path.be part of(local_model_path, "generation_config.json")
if os.path.exists(gcfg):
   EOS_TOKENS = checklist(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):
   supplier = 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, supplier, **base.get_model_input())
   with mesh:
       st = nnx.state(m)
       st = jax.lax.with_sharding_constraint(st, nnx.get_partition_spec(st))
       nnx.replace(m, st)
   return m
lora_policy = apply_lora(base_model, mesh)
def make_sampler(mannequin):
   return sampler_lib.Sampler(
       transformer=mannequin, 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 consider(rows, mannequin, tag, n_show=2):
   sampler = make_sampler(mannequin)
   right = fmt_ok = 0
   for i in vary(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).textual content
       for j, (o, r) in enumerate(zip(outs, chunk)):
           m = match_numbers.search(o); g = m.group(1) if m else None
           attempt:
               right += int(float(g) == float(r["answer"]))
           besides Exception:
               cross
           fmt_ok += int(match_format.search(o) will not be 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*right/len(rows):.1f}%   format = {100*fmt_ok/len(rows):.1f}%n")
print("n════════ BASELINE (earlier than GRPO) ════════")
consider(test_rows, lora_policy, "baseline")

We obtain the chosen Gemma-3 checkpoint, create the bottom mannequin from safetensors, and put together the tokenizer and EOS token checklist for era. We connect LoRA adapters to the eye and MLP projection modules, enabling us to coach a light-weight coverage with out updating the total mannequin weights. We additionally construct a sampler-based analysis perform and run a baseline take a look at earlier than GRPO coaching to measure the mannequin’s preliminary accuracy and adherence to the format.

Configuring the Tunix RL Cluster and Training

Copy Code
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.prepare(train_dataset, val_dataset)

We create the learning-rate schedule, optimizer, gradient clipping, and AdamW configuration that information the LoRA coverage updates throughout coaching. We outline the Tunix RL cluster with actor, reference, and rollout roles, then join the rollout configuration to the tokenizer, era limits, sampling temperature, and EOS tokens. We initialize the GRPO learner, launch TensorBoard for reside metrics, and begin the GRPO coaching loop on the ready GSM8K batches.

Evaluating and Exporting the Trained Model

Copy Code
print("n════════ AFTER GRPO ════════")
consider(test_rows, lora_policy, "after-grpo")
attempt:
   out_dir = "/content material/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 mannequin saved to {out_dir}")
besides Exception as e:
   print(f"n(Skipped merged export: {e})")
print("nDone. Checkpoints are in", CKPT_DIR)

We consider the skilled LoRA coverage after GRPO to check its math accuracy and response format towards the baseline end result. We then attempt to export the LoRA-merged Gemma-3 mannequin right into a Hugging Face safetensors listing so the skilled checkpoint may be reused outdoors the pocket book. We end by reporting the output checkpoint path, giving us an entire workflow from coaching setup to post-training analysis and non-obligatory mannequin export.

Conclusion

In conclusion, we accomplished a full GRPO fine-tuning loop for mathematical reasoning with Gemma-3, from dataset preparation and reward design to LoRA-based coverage coaching and post-training analysis. We noticed how Tunix organizes the actor, reference mannequin, rollout engine, optimizer, checkpoints, and metrics right into a reusable reinforcement studying pipeline. We additionally in contrast the mannequin earlier than and after GRPO to watch whether or not its reply format and numeric accuracy enhance throughout coaching. Finally, we optionally exported the merged LoRA mannequin, offering an entire path from uncooked GSM8K examples to a skilled, reasoning-oriented Gemma-3 checkpoint.

Check out the FULL CODES HERE. Also, be at liberty to comply with us on Twitter and don’t overlook to affix 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 companion with us for selling your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar and so on.? Connect with us

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

Similar Posts