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:

1. Correctness Reward (+2.0 if extracted answer matches ground truth, 0 otherwise)
2. Integer Reward (+0.5 if model extracts a valid integer from <answer> tags, 0 otherwise)
3. 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}},
}

References

• GRPO Paper: Group Relative Policy Optimization
• GSM8K Dataset: Cobbe et al., 2021
• Qwen2.5: Qwen/Qwen2.5 Models
• TRL: Hugging Face Transformers Reinforcement Learning

License

Apache License 2.0 (inherited from Qwen2.5-3B-Instruct)

---

Model Card Date: 2025-08-02
Last Updated: 2025-08-02