qwen2.5-0.5b-gsm8k

qwen2.5-0.5b-gsm8k

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 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)

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:

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):

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 and the code repo towardtype1/grade-school-math.