base_model: google/gemma-2-2b-it
library_name: jax
license: gemma
language:
- en
pipeline_tag: text-generation
tags:
- math
- reasoning
- grpo
- reinforcement-learning
- gsm8k
- gemma-2
- jax
- lora
- tpu
- google-tunix-hackathon
- kaggle
Gemma 2 2B - GRPO Math Solver 🧮🧠
This model is a fine-tuned version of google/gemma-2-2b-it, specifically optimized for solving mathematical word problems. It was trained on the GSM8K benchmark using Group Relative Policy Optimization (GRPO) via the JAX-native Tunix library.
The model was developed and submitted as part of the Google Tunix Hackathon hosted on Kaggle.
🚀 Model Details
- Base Architecture: Gemma 2 (2B Parameters, Instruction Tuned)
- Task: Mathematical Reasoning & High-Precision Problem Solving
- Dataset: GSM8K (Grade School Math 8K)
- Training Algorithm: GRPO (Group Relative Policy Optimization)
- Hardware: Trained on a Google v5e-8 TPU.
- Format: JAX/Tunix checkpoints.
⚙️ Training Hyperparameters (LoRA & GRPO)
To reduce hallucinations and encourage logical step-by-step thinking ("Show your work"), the model was fine-tuned using Low-Rank Adaptation (LoRA) integrated with a low-temperature GRPO setup.
LoRA Config:
- Rank: 64
- Alpha: 64.0
GRPO & Generation Config:
- Max Prompt Length: 256
- Total Generation Steps: 512
- Temperature: 0.6 (Optimized for math: less chaotic, highly logical, yet diverse enough for GRPO)
- Top-P: 0.95 (Cuts off the 5% least likely tokens to prevent hallucinations)
- Top-K: 50
- Num Generations: 4
- Beta / Epsilon: 0.08 / 0.2
Training Details:
- Learning Rate: 3e-6 (AdamW, Cosine Scheduler)
- Weight Decay: 0.1
- Max Grad Norm: 0.1
- Training Steps: 1000 batches (Micro-batch size: 2)
🧠 Why GRPO?
Instead of standard Supervised Fine-Tuning (SFT), this model utilizes GRPO, a Reinforcement Learning (RL) algorithm that enhances reasoning abilities. By generating multiple responses for a given prompt and evaluating them via a reward model, it calculates a relative advantage to update the policy without the heavy memory footprint of a separate value model.
🏆 Hackathon Context
- Competition: Google Tunix Hackathon - Train a model to show its work
- Objective: Enhance step-by-step reasoning and mathematical capabilities of smaller LLMs using JAX-native tools.
- Author: Tomasz Słapczyński