GPT-2 from scratch, 124M parameters
Writing a transformer that produces correct shapes takes an afternoon. Writing one that trains at a speed you can afford is a different exercise, and it is the one that teaches you what the libraries have been doing on your behalf. I built the architecture and the training pipeline end to end, then made it roughly ten times faster than the version that merely worked.
- Type
- Personal project, open source
- Stack
- PyTorch · CUDA · tiktoken
- Model
- 124M parameters, GPT-2 small configuration
- Precision
- bfloat16 mixed precision
Why build it by hand
Most of my work sits above the model: pipelines, evaluation, the systems that decide whether an output can be trusted. That is a comfortable place to stop understanding things. Building the model itself is what turns attention, tokenization, and the training loop from concepts I could describe into mechanics I have actually debugged.
The practical payoff shows up constantly. Knowing why a fused attention kernel matters, or what mixed precision is actually trading away, is the difference between guessing at an inference bill and reasoning about it.
What the speedup came from
The gain is not one trick. It is a stack of them, each removing a different bottleneck, and the order matters, because the first fix makes the next one measurable.
- bfloat16 mixed precision: halves memory traffic and lets the tensor cores do the work they were built for, without the loss-scaling fragility of float16.
- torch.compile: fuses the pointwise operations that otherwise round-trip to memory between every step of the block.
- Fused scaled-dot-product attention: replaces the materialized attention matrix with a kernel that never writes it out, which is where both the time and the memory were going.
- Gradient accumulation: decouples the batch size the optimizer sees from the batch size the GPU can hold, so the schedule stays correct on smaller hardware.
- Cosine schedule with warmup: warmup keeps the early steps from destabilizing the weights; the decay is what gets the final loss down rather than merely converging.
- Weight tying: shares the token embedding with the output projection, removing a large parameter block that gains nothing by being independent.
What I would do differently
The pipeline has no evaluation harness beyond training loss. Loss curves tell you the optimizer is working, not that the model is good. Wiring in a held-out benchmark would make the training run answer a question worth asking.
I would also add distributed training. Single-device is the right scope for learning the mechanics, but data-parallel training is where the next set of real problems lives, and they are problems you cannot see from one GPU.