A tiny GPT from scratch, and why it needs warmup

What I built

This is the third module of llm-lab. It puts the autograd engine and the transformer block from the first two modules together into a small GPT, and trains it on The Adventures of Sherlock Holmes, one character at a time.

The model has 4 layers, 4 heads, 128-dimensional vectors and a context of 128 characters, 0.82M parameters in all. The token embedding doubles as the output layer. AdamW, the learning-rate schedule (linear warmup, then cosine decay), gradient clipping and sampling with temperature, top-k and top-p are all written by hand.

The only new op the engine needed was the embedding lookup, and it had the quietest bug of the module. The obvious backward, W.grad[idx] += g, adds only one gradient per token when a token shows up more than once in the batch, which in text is always. NumPy doesn't complain, and the model still trains, just worse. The fix is one line:

def _backward():
    np.add.at(weight.grad, idx, out.grad)

Proving it's right

There are 13 tests: a gradient check of the whole model, a test that no position can see the future, overfitting a single batch, the properties of AdamW and clipping, the sampling filters, and saving and loading a checkpoint. The whole-model gradient check is what caught the embedding bug, with an error of 0.39.

That same check gives 2e-7 on a correct model, which looked too high next to the 1e-10 I'd come to expect. It isn't a bug. The embeddings start at 0.02, so the gradients reaching the blocks are around 1e-3, and the finite-difference noise of about 1e-9 is a bigger fraction of something that small. With larger embeddings the same check gives 7e-9.

Two engineering problems came before any of that. The first full training run used all 16 GB of memory within a few steps. Every node's backward function points back at the node, so each step's graph is a reference cycle, and Python only frees those when its garbage collector runs. The collector is triggered by how many objects you create, not how many bytes they hold, and one step is about 2,000 objects holding about 3 GB. A gc.collect() after each step fixed it. Then a step took 10.3 s, and it dropped to 2.4 s after I gave the (B, T, C) @ (C, D) matmul its own path in the engine.

The model

The full run took about 2.7 hours on a laptop CPU. Validation loss bottomed out at 1.42 between steps 1,500 and 3,000, about 18 passes over the book. After that it memorizes: training loss keeps falling to 0.79 while validation climbs back to 1.67 by step 5,000. With 0.82M parameters, half a million characters and no dropout, it runs out of new things to learn early. Training keeps the best checkpoint, so that's the one used for sampling.

At temperature 0.8, prompted with "Holmes":

Holmes gammed dried
the fearsal wound from the fire. She him brought of him in the writing powerson.”

“What will be an insigation don?”

“Never. Jone was the old, the must rooms from who done upon the Eysander on
presence, which was of his lens? Where you guide’s so.”

It has the look of the book right: dialogue, nested quotes, line breaks, mostly real words. It has none of the meaning. At 600 steps it was mostly made-up words.

Why does it need warmup?

Warmup starts the learning rate near zero and raises it over the first steps. Most GPT training recipes use it, and the usual explanation is that it prevents early loss spikes. I set out to see those spikes.

I didn't. I trained 8 short runs with and without warmup and clipping, at learning rates up to 10x the one I train with, and none of them spiked. Something else happened instead. Every run without warmup dropped to a loss of 3.1 and sat there. 3.1 is exactly the unigram loss: the model was predicting the overall character frequencies and nothing else. At a learning rate of 3e-2 without clipping, it was still there after 200 steps. With warmup it never got stuck.

So I measured what the model was doing. At 3e-2 without warmup, two steps in, it makes the same prediction at every position, whatever the text says. Block 0 is where it happens. After it, the vectors at different positions are almost identical: about 1% of their size is left that differs between positions, against about 100% at the start.

The cause is Adam's first steps. At the start, every position's gradient says the same thing: predict the common characters more. So the gradients of most weight matrices point mostly in one direction, 80–99% of their energy for the value, output, MLP and embedding matrices. Adam's first update moves every weight by about the learning rate, whatever the size of its gradient. At 3e-2 that one shared direction ends up 0.7x to 4.5x as strong as the strongest direction already in the matrix, and 6.7x in the embedding. So block 0 adds the same large vector to every position, and the differences between positions get buried.

Testing it

If that's right, it's the direction of the update that matters, not its size. So I edited every update during training, at 3e-2 with no warmup, for the first 30 steps:

Position spread is how different the positions still are after block 0, where 1 is roughly where they start. Context use is how much the predicted distribution changes between positions. A unigram model scores 0 on it, and the bigram model scores 0.64.

Shrinking the update to the same size changes nothing. Removing that one direction stops the collapse, and that run is ahead of the warmup run after 30 steps. So warmup works by keeping that shared direction small while all the gradients still agree. Even the learning rate I actually train with, 3e-3, leaves the model blind to the context at step 10 without warmup, although it recovers by step 30.

What it wasn't

I had two other suspects. Attention collapsing to one token at high learning rates is a known failure (Wortsman et al., 2023), but here the model is already blind to the context at step 2, when its largest attention score is still about 18. And the run with warmup ends up with much larger scores than the stuck one, about 2,600 against 400 at step 200, and trains fine.

The biases, which add the same amount to every position, aren't it either. With them frozen at 0, or with layer norm on the queries and keys, the model still collapses at step 2. Both do let it escape after 90–120 steps, which the normal run doesn't do within 200.

This was one model size, one seed, one book and 200-step runs. Removing the top direction is a diagnostic, not something I'd train with.

The √d check from module 2

The attention module found that with plain SGD, dividing by √d lets learning rates from 0.1 to 3 work at every head size. I said I'd check it with Adam here, so I reran the same retrieval task with Adam.

Both scalings still work over a wider range than no scaling, and ÷d has the widest. But every range moves down as the heads grow, and the top falls about 10x from d = 16 to 256 for all three. With ÷√d only 0.003–0.01 works at every size, against 0.1–3 with SGD, so a learning rate picked from the middle of the d = 16 range fails at d = 256.

The reason is the same one as above: Adam moves every weight by about the learning rate whatever its gradient, so the scores move further per step as d grows, however they're scaled. So module 2's "tune a small model, then make the heads bigger without retuning" still holds with Adam, but only inside that narrow band. This is what μP (Yang et al., Tensor Programs V) is designed for: it shrinks Adam's learning rate as the model gets wider. I only measured the edges of each range here, not the best learning rate.