Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)
The author builds and trains a GPT-2 small model from scratch in JAX, starting from a basic bigram-style model and incrementally adding components like LayerNorm and Transformer blocks. Achieved a final loss of 3.418, beating their PyTorch version (3.538) and original GPT-2 small (3.499) on the same test dataset.
Background
- This is part 34b of a long-running blog series where the author builds a GPT-2-style LLM from scratch, documenting each component step by step. The series started in late 2024 as a companion to Sebastian Raschka's book "Build a Large Language Model (from Scratch)".
- The author trained a GPT-2 Small model (124M parameters) using JAX—a high-performance numerical computing library from Google—instead of the more common PyTorch, to ensure they were truly building everything from scratch and not just copying existing code. Training took ~37 hours on an RTX 3090 GPU.
- Key results: their JAX model achieved better loss (3.42) than their own PyTorch model (3.54) and even slightly outperformed the original OpenAI GPT-2 Small (3.50) on the same test data, though OpenAI's weights still won on instruction-following tasks.
- This post walks through the incremental build: starting from a trivial "A-to-A" model that just reproduces its input, then adding components (layer normalisation, position embeddings, multi-head attention, transformer blocks) one at a time, showing that each addition measurably improved the loss.