85 lines
3.0 KiB
Markdown
85 lines
3.0 KiB
Markdown
---
|
|
license: apache-2.0
|
|
base_model: Qwen/Qwen2.5-0.5B
|
|
datasets:
|
|
- openai/gsm8k
|
|
language:
|
|
- en
|
|
pipeline_tag: text-generation
|
|
tags:
|
|
- math
|
|
- gsm8k
|
|
- grpo
|
|
- lora
|
|
- mlx
|
|
model-index:
|
|
- name: qwen2.5-0.5b-gsm8k
|
|
results:
|
|
- task:
|
|
type: text-generation
|
|
dataset:
|
|
name: GSM8K
|
|
type: openai/gsm8k
|
|
metrics:
|
|
- type: accuracy
|
|
value: 0.416
|
|
name: GSM8K test, zero-shot, strict `#### answer` match
|
|
---
|
|
|
|
# qwen2.5-0.5b-gsm8k
|
|
|
|
[Qwen/Qwen2.5-0.5B](https://huggingface.co/Qwen/Qwen2.5-0.5B) (base) post-trained for grade-school math on a MacBook Pro, with every training loop written by hand in [MLX](https://github.com/ml-explore/mlx) as a learning exercise.
|
|
|
|
Two stages, both via rank-16 LoRA on the attention q/v projections (1,081,344 trainable parameters, 0.2% of the model), adapters merged into the weights afterwards:
|
|
|
|
1. **SFT** on GSM8K's worked solutions (900 steps, lr 1e-4).
|
|
2. **GRPO** with verifiable rewards - 1.0 if the extracted `#### answer` matches gold, else 0.0 - with a k3 KL penalty against the frozen SFT reference (100 steps, groups of 8 at temperature 0.8, lr 1e-5, kl 0.1, grad clip 1.0).
|
|
|
|
## Results (GSM8K test, zero-shot, greedy, strict `#### <number>` match)
|
|
|
|
| Model | Accuracy |
|
|
|---|---|
|
|
| Qwen2.5-0.5B base | 0.000 (0.305 with 4-shot prompting) |
|
|
| + LoRA SFT | 0.345 |
|
|
| **+ LoRA GRPO (this model)** | **0.416** |
|
|
| full-fine-tune SFT, for comparison | 0.325 |
|
|
| full-fine-tune GRPO, for comparison | 0.303 |
|
|
|
|
Full test set (1,319 questions) for all rows.
|
|
The interesting finding: the same GRPO recipe that stalled on full fine-tuning gained 7 points under the low-rank constraint - the small trainable space acts as a structural leash on noisy small-batch policy gradients.
|
|
|
|
## Usage
|
|
|
|
Prompt format (the model is a narrow GSM8K-format specialist, not a chat model):
|
|
|
|
```
|
|
Question: <your word problem>
|
|
Answer:
|
|
```
|
|
|
|
It emits step-by-step reasoning ending in `#### <number>`.
|
|
|
|
With mlx-lm:
|
|
|
|
```python
|
|
from mlx_lm import load, generate
|
|
model, tokenizer = load("towardtype1/qwen2.5-0.5b-gsm8k")
|
|
print(generate(model, tokenizer, prompt="Question: Tom has 3 boxes of 12 pencils. He gives away 5. How many are left?\nAnswer:", max_tokens=200))
|
|
```
|
|
|
|
With transformers (weights are standard Qwen2 layout):
|
|
|
|
```python
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
model = AutoModelForCausalLM.from_pretrained("towardtype1/qwen2.5-0.5b-gsm8k")
|
|
```
|
|
|
|
## Limitations
|
|
|
|
- Single-task specialist: trained and evaluated only on GSM8K-style word problems in the exact format above. It is not aligned, not a chat assistant, and inherits any base-model limitations and biases.
|
|
- 0.5B parameters: arithmetic slips are common; verify answers.
|
|
|
|
## Training write-up
|
|
|
|
The full story - including two instructive GRPO failure modes (format collapse without a KL penalty, gradient-noise drift without clipping) and why the LoRA constraint fixed them - is in the blog post [Post-training from scratch](https://towardtype1.github.io/essays/post-training-from-scratch/) and the code repo [towardtype1/grade-school-math](https://github.com/towardtype1/grade-school-math).
|