kotlarmilos/gemma-3-1b-reasoning

Gemma 3 1B-IT reasoning LoRA

A LoRA adapter that improves math reasoning in Gemma 3 1B-IT, trained with a three stage pipeline, SFT then SFT then GRPO, using Tunix on TPU v5e.

Where this fits

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:

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

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