Qwen2.5-3B-Instruct GRPO on GSM8K

A LoRA-finetuned variant of Qwen/Qwen2.5-3B-Instruct trained with GRPO (Group Relative Policy Optimization) on the GSM8K dataset (grade-school math word problems with ground-truth numeric answers).

Model Details

Base Model: Qwen/Qwen2.5-3B-Instruct
Training Method: GRPO (Group Relative Policy Optimization) via TRL
Training Data: GSM8K (train split, ~7,473 examples)
LoRA Rank: 32
Quantization: 4-bit (NF4) during training
Hardware: NVIDIA T4 GPU (Colab, ~3 hours)
Training Steps: 120
Per-Step Generations: 4 completions per prompt

Training Details

Reward Functions

The model was optimized to maximize a weighted sum of structured rewards:

Correctness Reward (+2.0 if extracted answer matches ground truth, 0 otherwise)
Integer Reward (+0.5 if model extracts a valid integer from <answer> tags, 0 otherwise)
Format Reward (+0.5 if model follows the XML tag structure <reasoning>...</reasoning><answer>...</answer>, 0 otherwise)

Training Configuration

learning_rate: 5e-6
scheduler: cosine with warmup (10%)
optimizer: AdamW (8-bit)
batch_size: 1 (gradient accumulation: 4, effective: 4)
num_generations: 4
max_prompt_length: 256
max_completion_length: 512
max_steps: 120
save_interval: 30 steps

Prompt Format

Use the model with the following system prompt and format:

system_prompt = """Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""

messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": "Your math question here"},
]

Example Output

Question: "A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?"

Model Output:

<reasoning>
To solve this problem, we need to determine the amount of fiber required for both the blue and white fibers separately, and then sum them up.

Given:
- A robe takes 2 bolts of blue fiber.
- The robe takes half as much white fiber as blue fiber.

First, calculate the amount of white fiber:
- Since the robe takes half as much white fiber as blue fiber, the amount of white fiber is \( \frac{1}{2} \times 2 \) bolts = 1 bolt.

Next, add the amounts of blue and white fibers together to find the total number of bolts:
- Total bolts = Blue fiber + White fiber
- Total bolts = 2 bolts (blue) + 1 bolt (white)
- Total bolts = 3 bolts

So, the robe requires a total of 3 bolts of fiber.
</reasoning>
<answer>3 bolts</answer>

Extracted Answer: 3 bolts ✓ (correct)

Usage

With Transformers + PEFT

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B-Instruct",
    device_map="auto",
    torch_dtype=torch.float16,
)

# Load LoRA weights
model = PeftModel.from_pretrained(
    base_model,
    "sarimahsan101/qwen2.5-3b-grpo-gsm8k",
    device_map="auto",
)

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B-Instruct")

# Generate
prompt_text = tokenizer.apply_chat_template([
    {"role": "system", "content": "Respond in the following format:\n<reasoning>\n...\n</reasoning>\n<answer>\n...\n</answer>"},
    {"role": "user", "content": "Your question here"},
], tokenize=False, add_generation_prompt=True)

inputs = tokenizer(prompt_text, return_tensors="pt").to("cuda")
with torch.no_grad():
    output = model.generate(**inputs, max_new_tokens=512, temperature=0.7, top_p=0.9, do_sample=True)

response = tokenizer.decode(output[0], skip_special_tokens=True)
print(response)

Merge LoRA into Base Model

For inference without PEFT overhead:

from peft import PeftModel
import torch

base_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B-Instruct", device_map="auto")
model = PeftModel.from_pretrained(base_model, "sarimahsan101/qwen2.5-3b-grpo-gsm8k")

# Merge
model = model.merge_and_unload()
model.save_pretrained("qwen2.5-3b-grpo-gsm8k-merged")

Limitations

Trained on GSM8K (grade-school math). Performance on other domains (code, creative writing, general QA) is not evaluated.
LoRA rank 32 is relatively conservative; larger ranks may improve performance at the cost of more parameters.
Reward function design is task-specific; transfer to other domains may require retuning.
Training was limited to 120 steps (~2.5 hrs on T4 GPU); longer training may improve convergence.

Training Environment

Framework: Hugging Face Transformers + TRL + PEFT
Quantization: BitsAndBytes (4-bit NF4)
GPU: NVIDIA T4 (16 GB VRAM, Colab)
Peak Memory: ~14 GB
Training Time: ~3 hours total (install + training + save)

Citation

If you use this model, please cite:

@misc{qwen2.5-3b-grpo-gsm8k,
  title={Qwen2.5-3B-Instruct fine-tuned with GRPO on GSM8K},
  author={sarimahsan101},
  year={2025},
  publisher={Hugging Face},
  howpublished={\url{https://huggingface.co/sarimahsan101/qwen2.5-3b-grpo-gsm8k}},
}