Transformer from Scratch

I’d used transformers plenty, but I’d never written one where I could point at every line and say what it does. So I wrote one, from the embedding layer up, and then trained it long enough to actually translate.

TLDR: an encoder–decoder transformer in plain PyTorch, the shape of the original “Attention Is All You Need” model. 6 encoder blocks, 6 decoder blocks, d_model 512, 8 heads, a 32,000-word vocabulary on each side. 93.4M parameters. I trained it for 200,000 steps on a million English–French sentence pairs on my laptop’s GPU (Apple MPS), about 21 hours, and added a KV cache that makes decoding 2.2–2.5x faster.

Building it piece by piece

The git log is basically the table of contents. Embeddings and multi-head attention first, then the encoder, then the decoder, then cross-attention, then the full model. I kept the earlier versions of attention in the file (single-head, then one weight matrix per head) because each one was a step I needed to see work before the next.

Git log listing 16 commits from August 28 to September 5, from initial embedding and attention through decoder, training and KV cache — Sixteen commits over nine days. Reading it bottom to top is reading the model from the inside out.

The attention module is where most of the thinking went. It handles self-attention, masked self-attention and cross-attention with the same code, plus the cache.

Python source for MultiHeadAttention.forward, including the KV cache path and the causal mask — MultiHeadAttention.forward. When a cache is passed in, only the new token gets projected; the old keys and values are reused.

Data

The corpus is a 22.5M-row English–French CSV. My first loader just read the whole thing into pandas and tokenised in __getitem__. That’s fine on the 5,000-row sample and hopeless on the real file.

The version that trained the model does two passes over the CSV, once to count words and build the vocabularies and once to encode every row, and writes the result to a binary cache. After that the dataset is memory-mapped, so start-up is instant and fetching a sentence is an array slice. Batches are length-bucketed: sentences of similar length get batched together so you’re not padding a 6-word sentence out to 120.

Python source for LengthBucketBatchSampler, which sorts pools of examples by length before batching — The length-bucketed sampler. Shuffle, cut into pools, sort each pool by length, then batch. Much less padding.

Tokenisation is the dumb kind: lowercase and split on spaces. Punctuation stays stuck to words, so “paris.” and “paris” are different tokens, and anything outside the top 32,000 becomes <unk>. That shows up in every output below.

Training on a laptop

I trained on the first million rows, 200,000 steps at batch 32, which is about six passes over that subset. It ran for 21 hours at around 3,200 tokens a second. Memory was the constant fight: a 4,096-token batch peaks near 7.8 GB, and somewhere around 32,000 tokens per batch the MPS backend just dies. The training loop skips a batch on out-of-memory instead of crashing, and drops the allocator cache every so often because length bucketing keeps changing tensor shapes.

Three plots of training loss, perplexity and learning rate over 200,000 steps — Loss went from 9.40 to 2.31 (perplexity about 10). The little bumps every 31,250 steps are the running average resetting at each epoch, not the model getting worse.

The training script writes the model’s output for a few fixed sentences every 1,000 steps, and that log turned out to be my favourite record of the model learning. It starts with random French-ish words, picks up common phrases, and eventually gets whole sentences right.

Table of the same four input sentences translated at steps 0, 1000, 5000, 20000, 50000 and 200000 — Same inputs at six points in training. At step 1,000 everything is 'il y a'. By 200,000 most of it is right, apart from the words it never learned.

KV cache

Without a cache, generating each new word re-runs attention over every word so far. With one, the encoder runs once and each step only projects the newest token. I pass the cache in and out of the model instead of storing it on the modules, so two translations can run against the same model and nothing leaks into the next training step.

Measured on CPU, cached decoding is 2.15–2.49x faster, and the outputs match the full forward pass to within 2.3e-6.

Terminal output of cached versus uncached decode timings, speedups between 2.15x and 2.49x with identical tokens — Cached against uncached decode of 96 tokens, on a fresh model and on the trained checkpoint. Same tokens every run.

What it gets right and wrong

Eight English sentences and the model's French translations, some perfect and some with unknown-word tokens — Real greedy outputs from the final checkpoint. 'ce rapport a été publié en 2005' is perfect. 'where is the library?' becomes 'où est le <unk>' because 'library?' with the question mark was never in the vocabulary.

The attention maps convinced me it had learned real structure. In the later decoder blocks, cross-attention lines up English and French words almost along a diagonal, bending where French word order differs.

What I’d change

A real subword tokenizer (BPE) first, since half the mistakes above come from the vocabulary. Then a validation set, because right now I only know training loss. There’s also a decoder-only version with grouped-query attention in the repo that I haven’t trained yet.