Post-training from scratch: a tutorial

This is a walkthrough of a small project: teaching Qwen2.5-0.5B to solve grade-school maths problems, on a MacBook, with every training loop written by hand. The code lives at github.com/towardtype1/grade-school-math.

It is written as a tutorial. By the end you should understand what every piece does and why it exists - including the two training runs that failed, because those taught the most.

The picture

A language model fresh out of pretraining is not an assistant. It has read a large fraction of the internet and learned one skill: given some text, predict the next token. Post-training is the process of turning that raw predictor into something that answers questions. We do it in two stages, and we measure before and after every stage:

 base model ──► [ SFT: imitate worked solutions ] ──► [ RL: reward correct answers ] ──► final
     │                      │                                    │
   eval                   eval                                 eval
  (0.000)                (0.330)                                (?)

The repo mirrors this picture, one file per job:

gsm/
  data.py    - load the dataset, format prompts, extract answers
  eval.py    - measure accuracy of any checkpoint
  sft.py     - supervised fine-tuning loop
  grpo.py    - reinforcement learning loop
  report.py  - HTML report: curves, and sample answers evolving over training

The data, and one very important regex

GSM8K is 8.5k word problems. Every solution ends the same way:

Natalia sold 48/2 = 24 clips in May.
Natalia sold 48 + 24 = 72 clips altogether.
#### 72

That #### 72 line is the whole reason this project works. It means we can check an answer with a regex instead of a human:

ANSWER_RE = re.compile(r"####[\s$*]*([-+]?[\d,]*\.?\d+)")

def extract_answer(text):
    m = ANSWER_RE.search(text)
    return normalize_number(m.group(1)) if m else None

One function, used in exactly two places: scoring evals, and computing RL rewards. This is deliberate. If eval and reward used different parsers, the model could look good on one and be trained on the other, and you would chase ghosts.

Measuring: the eval harness

Evaluation is a loop: show the model a question, let it generate greedily (always pick the most likely next token, so results are repeatable), extract the answer, compare with the gold answer.

prompt = f"Question: {question}\nAnswer:"
completion = generate(model, tokenizer, prompt)      # greedy decode
correct = extract_answer(completion) == gold_answer

We measure the untrained base model two ways, and the two numbers frame everything that follows:

• Zero-shot: 0.000. Just the question. The model actually reasons - it writes things like “The answer is 72” - but it never uses the #### format, and it never stops: it invents a new question and keeps going. A document usually continues; the model continues it.
• 4-shot: 0.305. Paste four solved examples above the question. Now “predict the next token” happens to mean “keep solving in this format”, and the same weights get 30% right.

So the maths is largely already in the model. What is missing is behaviour: answer the question you were asked, in a recognisable format, then stop. That is what post-training installs.

What a model outputs, and what a loss is

One paragraph of theory, because everything else builds on it.

For each position in the text, the model outputs a score for every token in its vocabulary (~152k of them). Softmax turns those scores into probabilities. The training loss for one position is cross-entropy:

loss = -log( p(the correct next token) )

Read it as: if the model gave the right token probability 1, the loss is 0; probability 0.01, the loss is 4.6. Training nudges every weight to make the observed next token more probable - that is all “learning” means here, in both SFT and RL.

Stage 1: SFT - learning by imitation

Supervised fine-tuning continues exactly this training, but only on text we want imitated: real GSM8K questions with their worked solutions. The full loss function from gsm/sft.py:

def loss_fn(model, tokens, mask):
    inputs, targets = tokens[:, :-1], tokens[:, 1:]   # position t predicts token t+1
    logits = model(inputs)
    ce = nn.losses.cross_entropy(logits, targets, reduction="none")
    m = mask[:, 1:]
    return (ce * m).sum() / mx.maximum(m.sum(), 1)

Three things are happening:

The shift. inputs is the text minus its last token, targets is the text minus its first. Lining them up makes position t’s prediction get scored against token t+1 - next-token prediction, mechanically.

The mask. mask is 0 over the question and 1 over the solution. Multiplying the loss by it means the model is never trained to produce questions - the question is context, the solution is the target. Forget this and you train a question generator.

The average. Dividing by the number of unmasked tokens makes batches of different lengths comparable.

The loop around it is: batch examples, compute gradients, AdamW optimizer step, warmup-then-cosine learning-rate schedule, save checkpoints. About 140 lines, no magic.

The sanity check that earns its keep

Before any long run: train on just 4 examples and demand the loss collapse to ~0 in minutes. A loop that cannot memorise 4 examples is broken, and you want to know now.

Ours was broken. Loss spiked from 1.18 to 16.3. The code was fine - the problem was numerical. MLX loads models in bfloat16, a 16-bit format with only 8 bits of precision, and the optimiser was keeping its internal statistics in the same format. AdamW divides by the square root of a running average of squared gradients; at 8 bits of precision that division is noise. Fix: convert the model to 32-bit floats for training, save checkpoints back as compact bfloat16 copies.

After 900 real steps (about an hour):

The model now beats its own few-shot baseline with no examples in the prompt. And SFT was cheap: one epoch, one laptop-hour.

But notice the ceiling built into the method. SFT grades the model on matching the reference text, token by token. A wrong answer written in perfect style scores the same as a right one. To optimise being right, the training signal has to know what “right” means.

Stage 2: RL - learning from consequences

Reinforcement learning flips the direction of learning. Instead of showing the model what to write, we let it write, check whether it was right, and adjust.

The algorithm is GRPO. One training step, in words:

1. Take a question. Sample a group of 8 answers from the model at temperature 0.8 (temperature adds randomness, so the 8 answers differ).
2. Reward each answer: 1.0 if extract_answer matches gold, else 0.0. Same parser as eval.
3. Convert rewards to advantages by normalising within the group.
4. Nudge the model: make high-advantage answers more probable, low-advantage answers less probable.

Step 3 is GRPO’s clever bit, so let’s do it by hand. Say 1 of the 8 answers was correct: rewards [1, 0, 0, 0, 0, 0, 0, 0].

mean = 0.125
std  = 0.33
advantage(correct) = (1 - 0.125) / 0.33 = +2.65
advantage(wrong)   = (0 - 0.125) / 0.33 = -0.38

The one success gets a strong push up; the failures get a mild push down. “Better or worse than the group” replaces the separate value network that classic RL (PPO) trains just to estimate expectations - the group IS the expectation. Note a built-in edge case: if all 8 answers are right, or all wrong, every advantage is exactly 0 and the step teaches nothing. Remember that; it comes back to bite.

Step 4 is the policy-gradient loss, from gsm/grpo.py:

def grpo_loss(model, tokens, mask, advantages):
    logp = -nn.losses.cross_entropy(logits, targets, reduction="none")
    per_seq = (logp * m).sum(axis=1) / m.sum(axis=1)   # avg log-prob of each answer
    return -(advantages * per_seq).mean()

Read the last line slowly. per_seq is how probable the model finds each of its own sampled answers. Multiplying by the advantage and minimising the negative means: raise the probability of answers with positive advantage, lower it for negative ones. That is the entire mechanism of RL on language models. Notice there is no reference answer anywhere - the model learns from its own attempts.

Failure 1: the reward collapses

We ran exactly the loop above. Within 16 steps the mean reward went from 0.18 to 0.00 and never recovered.

The logged samples showed why. The answers stayed coherent - they just grew longer each step, until they ran past the token limit without ever writing ####. No ####, no parse, no reward.

The mechanism is worth internalising. With 1-in-8 correct, each step pushes down seven answers for every one it pushes up. But all eight answers share their skeleton - the format, the reasoning style, the “finish and answer” move. Pushing down seven wrong-but-well-formatted answers erodes that shared skeleton seven times as fast as the one success rebuilds it. The model drifts toward never finishing: a perfect strategy for avoiding being wrong, and for never being right.

Then the edge case from earlier arrives: all 8 answers score 0, the group is uniform, all advantages are exactly 0, and the gradient is gone. A policy that never succeeds gets no signal to recover with. The run was not diverging; it was dead.

The fix: a KL leash

Every production RL recipe includes a term that now makes visceral sense: a KL divergence penalty tying the trained policy to a frozen copy of its starting point (the SFT checkpoint).

The intuition: the reward only measures one thing (the final answer), but the policy carries a thousand behaviours the reward never mentions - formatting, stopping, staying on topic. Unconstrained optimisation will happily sacrifice unmentioned behaviours to move the mentioned number. The KL term prices that sacrifice: drift from the reference and pay a penalty, everywhere, not just where the reward looks.

The estimator, term by term:

diff = ref_logp - logp                # how differently the frozen reference
                                      # and current policy rate these tokens
kl = exp(diff) - diff - 1             # ≥ 0; equals 0 only when diff = 0
loss = pg_loss + 0.1 * kl.mean()      # coefficient = the price of drift

Check it does what the intuition says: if the policy hasn’t moved, diff = 0 and the penalty is exactly 0 - staying put is free. The further apart the two models rate the same tokens, the faster the penalty grows.

Failure 2: noise, and the fix you can measure

The rerun with a small KL penalty kept the format alive (the leash worked) but the reward still decayed - this time the arithmetic degraded, with samples like $2,400 / 3 = $8,000 / 3 = $240. And a tell: reward never improved, not even at step 2. That pattern says noise, not learning.

Each of our updates was built from only 2-4 usable groups. Production GRPO uses hundreds of questions per update; averaging over many gradients cancels their noise. Ours barely averaged anything, and once we logged it, we could measure the problem: our gradients had norms around 6, swinging the weights wildly every step.

Three standard fixes, all cheap:

grads, grad_norm = optim.clip_grad_norm(grads, max_norm=1.0)  # cap update size

plus a lower learning rate, and a stronger KL coefficient.

And one borrowed habit: before betting hours on a full run, we ran an RL overfit check - the same idea as SFT’s overfit-one-batch, but for RL. Train on a frozen pool of 8 questions for 20 steps; if the loop can learn at all, reward on that pool must trend up. It did: from ~0.06 to ~0.44 on the hard questions and ~0.5 to ~0.9 on the easy ones, with bounded KL. Only then did the full run get launched.

Where it landed

The hardened run went the distance: 100 stable steps, KL flat at ~0.05, no collapse. Final numbers, on the full 1,319-question test set:

GRPO did not beat SFT. But the logs show something more interesting than a null result. Over the run, the model’s sampled success rate - how often an attempt at temperature 0.8 is correct, which is exactly what the reward measures - climbed from 0.23 to roughly 0.31. Its greedy accuracy - the single most confident answer, which is what benchmarks measure - never moved (the in-run eval scored the same 12 of 50 questions at steps 20, 40 and 60).

Read those two facts together: RL genuinely optimised the objective it was given, and the objective was not the benchmark. We trained at temperature and graded greedily, with 48 samples per update where production runs use thousands. At this scale, GRPO reshuffles probability mass among behaviours the model already has - it selects, rather than teaches - and its selection pressure never reached the argmax path.

That looked like the end of the story. It wasn’t.

The LoRA twist

Everything above trained all 494 million weights. There is another way: freeze the whole model and inject tiny trainable matrices into each attention layer. That is LoRA, and the entire mechanism is ~20 lines:

class LoRALinear(nn.Module):          # wraps a frozen linear layer
    def __init__(self, linear, rank, alpha):
        self.linear = linear                        # W, frozen, bf16
        self.lora_a = mx.random.normal((rank, in_dims), scale=1/math.sqrt(in_dims))
        self.lora_b = mx.zeros((out_dims, rank))    # zero: delta starts at 0
        self.scale = alpha / rank

    def __call__(self, x):
        y = self.linear(x)
        return y + (self.scale * ((x @ self.lora_a.T) @ self.lora_b.T)).astype(y.dtype)

Every update is now confined to a rank-16 subspace: 1,081,344 trainable parameters, 0.2% of the model. Intuition says constraining the update space this hard must cost quality. Same data, same steps, same eval:

The SFT row: no cost at all (0.345 vs 0.325 is at the edge of noise - call it a match), in a third of the training time, with no fp32 upcast needed - the bf16-AdamW instability lives in the trained parameters’ optimiser state, and LoRA’s are small enough to keep in fp32 for free.

The GRPO row is the result of the whole project. The same RL recipe that collapsed once, drifted once, and then merely-held-steady on full fine-tuning gained seven points on LoRA - the only run where greedy accuracy genuinely improved (in-run evals: 0.26, 0.28, 0.42, 0.44 against a 0.34 baseline). And the drift telemetry says why: KL against the reference averaged 0.008 all run, roughly six times below the full-FT stable run’s 0.048. The low-rank constraint is itself a leash - there is simply not enough room in a rank-16 subspace to strip-mine distant behaviours, so the noisy 48-samples-per-update gradient could only spend its budget near the reward.

So the earlier conclusion needs an honest amendment. RL at laptop scale is not doomed to reshuffle - it worked, decisively, once the update space was constrained tightly enough for its noisy signal. Full fine-tuning gave the noise 494 million directions to do damage in; LoRA gave it a million, mostly pointing somewhere useful.

What the six runs demonstrated

1. Post-training is behaviour installation. The maths knowledge was in the base model all along (0.305 few-shot, zero training); SFT installed “answer and stop”; RL - properly constrained - taught the model to lead with its best attempt (0.416).
2. RL optimises the reward you wrote, not the behaviour you meant. Destructively when unconstrained (two dead runs), inertly when merely stabilised (the flat full-FT run), and productively only when the trainable space itself was small enough to aim (the LoRA run).
3. Constraints are not compromises. LoRA looked like the budget option and beat the full-parameter version on both stages - fewer knobs meant less noise damage, not less learning.
4. Sanity checks scale down. Overfit-one-batch, sample logging, gradient-norm logging, RL overfit gates, and same-slice baselines caught every failure in this project early, on hardware that costs less than a conference ticket.