tslapczynski/gemma-2-2b-it-grpo-gsm8k

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