Mini-GPT
A decoder-only GPT written from scratch in PyTorch, one component at a time, each with its own tests, an experiment and a document that explains the math.
- Python
- PyTorch
Challenge
Understand exactly what happens inside a language model when it receives a sentence, without the shortcuts of nn.Transformer, and justify each design choice with a number measured on the model itself.
Approach
Build the model in eighteen components, from the tokenizer to the BPE tokenizer that replaced it, each with its own test file, an experiment and a document. Pieces written by hand, such as LayerNorm and the cross-entropy loss, have tests that check they match PyTorch's own. Every design choice is recorded with the main alternative it ruled out and the reason, usually a measurement. Training uses AdamW with warmup and gradient clipping, and its checkpoints keep the weights, the optimizer state, the position in the data loader and the random states.
Outcome
A 15-epoch run brings the best validation loss to 2.551, a perplexity of 12.82, and the model was still improving when training stopped. Resuming from a checkpoint reproduces the uninterrupted run bit for bit, which a test checks, and a trained checkpoint ships with the repository, so text can be generated without training first.
Architecture
Decisions and trade-offs
A BPE tokenizer replaced the original character tokenizer
Generation stopped in the middle of words, and BPE covers about 2.1x more characters in the same context length.
A much larger token embedding: the vocabulary went from 112 to 4,096 tokens.
Pre-norm: LayerNorm before attention and the feed-forward
The residual path stays free of normalizations, so the gradient crosses the full depth of the model.
The residual stream is never normalized and its scale grows during training, so a final LayerNorm sits before the LM head.
Weight tying between the token embedding and the LM head
As in GPT-2, it saves V×d parameters and gives every row of the embedding a gradient at every step.
LayerNorm and the cross-entropy loss written by hand
The goal is to learn what they do; tests guarantee they match PyTorch's nn.LayerNorm and F.cross_entropy.
The hand-written cross-entropy is about 0.14 ms slower per step.
Checkpoints save the whole training state, not just the weights
Resuming reproduces the uninterrupted run bit for bit. Without the AdamW state, the first steps after resuming move further than normal and the result is no longer reproducible.
Generation samples with temperature 0.8 and top-p 0.95
Greedy decoding falls into loops, and top-p adapts the number of candidates to how confident the model is.
Key evidence
Validation perplexity of 12.82 after a 15-epoch run