This is the second step in my applied post-training track. In the first step, Phi-4 on dotnet/runtime, I changed what a small model knows by fitting it to a domain. Here I change how a model thinks. The supervised stages install a reasoning format and content, and the reinforcement stage rewards correct answers under a checker rather than a judge model. The next step, Gemma 4 GlucoLens, takes this same post-training instinct into a domain where the output is a structured rollout and the model must refuse when it is unsure.

Writeup on Kaggle, gemma-3-1b-reasoning.

Training pipeline

Stage 1. SFT on math and science

Builds the reasoning foundation using structured math and science datasets.

Data. GSM8K 70% and ScienceQA 30%
Steps. 1000, batch size 4, lr 2e-5 with warmup and cosine decay
Sequence length. 1024

Stage 2. SFT on diverse tasks

Expands to code, summarization, and creative writing while retaining math performance through weighted sampling.

Data. GSM8K 25%, ScienceQA 15%, MBPP 25%, XSum 20%, WritingPrompts 15%
Steps. 600, batch size 4, lr 1e-5

Stage 3. GRPO with a math reward

Reinforcement learning using Group Relative Policy Optimization. The reward function verifies answer correctness against GSM8K gold labels, so no judge model is needed.

Data. GSM8K prompts
Steps. 50, lr 1e-6
GRPO config. 4 generations, 2 iterations, beta 0.1, epsilon 0.2
Reward. 1.0 for a correct answer, 0.2 for showing reasoning steps, 0.1 for reaching a number, 0.1 for a clean ending, 0.0 for degenerate output

LoRA configuration

Parameter	Value
Rank	32
Alpha	32.0
Target modules	`q_einsum`, `kv_einsum`, `gate_proj`, `down_proj`, `up_proj`, `attn_vec_einsum`

Usage

This adapter was trained with Tunix and Qwix in JAX. To load and use it:

python

from tunix.models.gemma3 import params, model
from tunix.generate import sampler as sampler_lib
import qwix

base = params.create_model_from_checkpoint(
    params.GEMMA3_1B_IT,
    model.ModelConfig.gemma3_1b_it()
)

lora_model = qwix.apply_lora_to_model(
    base,
    qwix.LoraProvider(
        module_path=".*q_einsum|.*kv_einsum|.*gate_proj|.*down_proj|.*up_proj|.*attn_vec_einsum",
        rank=32, alpha=32.0,
    ),
    rngs=nnx.Rngs(0),
    **base.get_model_input(),
)

from safetensors.numpy import load_file
adapter = load_file("adapter_model.safetensors")
# Merge adapter weights into lora_model state

tokenizer = params.create_tokenizer()
sampler = sampler_lib.Sampler(
    transformer=lora_model, tokenizer=tokenizer,
    cache_config=sampler_lib.CacheConfig(
        cache_size=1536,
        num_layers=lora_model.config.num_layers,
        num_kv_heads=lora_model.config.num_kv_heads,
        head_dim=lora_model.config.head_dim,
    ),
)

prompt = "<start_of_turn>user\nWhat is 25 * 13? Think step by step.<end_of_turn>\n<start_of_turn>model\n"
out = sampler(
    input_strings=[prompt],
    max_generation_steps=512,
    temperature=0.7, top_k=50, top_p=0.95,
    echo=False, eos_tokens=[106],
)
print(out.text[0])

Technical details

Framework. JAX with Tunix, Qwix, and Flax NNX
Hardware. TPU v5e on Google Colab
Precision. bfloat16
Optimizer. AdamW with warmup cosine decay for SFT, clipped AdamW for GRPO
GRPO reference model. A frozen copy of the base Gemma 3 1B-IT as the KL anchor

Source

Code https://github.com/kotlarmilos/gemma-3-1b-reasoning
Writeup https://huggingface.co/blog/kotlarmilos/gemma-3-1b-reasoning
Hackathon writeup https://www.kaggle.com/competitions/google-tunix-hackathon/writeups/gemma3-reasoning

Citation

bibtex

@misc{kotlar2025gemma3reasoning,
  title={Gemma 3 1B-IT Reasoning LoRA},
  author={Kotlar, Milos},
  year={2025},
  url={https://github.com/kotlarmilos/gemma-3-1b-reasoning}
}