README.md

---
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](https://www.kaggle.com/competitions/google-tunix-hackathon)
* **Objective:** Enhance step-by-step reasoning and mathematical capabilities of smaller LLMs using JAX-native tools.
* **Author:** Tomasz Słapczyński