Frame 05Open sourceFoundationsAug 2025
The GPT-2 training recipe in one readable file
A decoder-only transformer and its training loop, written from scratch in PyTorch after Andrej Karpathy's nanoGPT and "Let's reproduce GPT-2" recipe. I built it to see the whole modern stack in one place. It is a study, not a trained model.
- Stack
- Python, PyTorch, tiktoken (GPT-2 BPE), CUDA
- Scope
- Sole author, following the nanoGPT recipe
- One PyTorch file holds the model and the loop: 12 pre-norm blocks at 768 dimensions with 12 heads and a 1,024-token context, tied embeddings, fused causal attention, bf16 autocast, torch.compile, gradient accumulation, warmup then cosine.
- Every setting is readable in
GPT.py. The 524,288-token step, the clip at 1.0 and the 6e-4 peak learning rate are the nanoGPT numbers, not mine. - The committed run is a 10-step smoke test on one novel. There are no tests, logs or checkpoints, and the README claims figures the repo cannot back.
What could go wrong, and how we would know
A training script can be wrong and still print a falling loss. A mask that leaks one future token, a schedule that never decays, an accumulated gradient that is never scaled: all three look like learning for a while. This file's safety net is small, and I want to be exact about how small.
Causal masking is delegated to scaled_dot_product_attention(q, k, v, is_causal=True) instead of a hand-built triangle, and the hand-built version is left commented out beside it so the two can be compared. get_lr() asserts the decay ratio stays inside [0, 1]. Each optimizer step prints loss, milliseconds and tokens per second only after torch.cuda.synchronize(), so the timing is real. That is the whole net. Nothing in the repo can catch a number the code never produced, and the README shows it: its parameter count, speedup and loss curve match no artifact in the tree.
What I built
The model is GPT(vocab_size=50304, model_dim=768, context_length=1024, num_transformer=12, num_heads=12): token and position embeddings, twelve blocks that apply LayerNorm before attention and again before the MLP with a residual around each, a final LayerNorm, and an output projection whose weight is the token embedding's weight. Linear and embedding weights start from a normal with std 0.02, biases at zero. The vocabulary is padded from tiktoken's 50,257 GPT-2 tokens to 50,304, a multiple of 64, so the output matmul lands on tensor-core-friendly shapes.
Attention is twelve single-head modules, each with its own bias-free q, k and v projections, concatenated and passed through one output linear. Each head calls the fused SDPA kernel with is_causal=True. Above that call sits the version I wrote first: scores, a scale, tril mask, masked_fill with negative infinity, softmax. I kept it as a comment on purpose, as the spec the fused call has to match. It is not an exact spec: that draft divides by the square root of the score matrix's last dimension, which is the sequence length, where SDPA divides by the square root of the head size.
The loop compiles the model, sets float32_matmul_precision('high'), and runs the forward under bf16 autocast on CUDA. The micro-batch is 4 by 1,024 tokens. The loop accumulates total_batch_size // (B*T) of them, 128 passes, into one 524,288-token optimizer step, dividing each loss by the accumulation count before backward(). Gradients are clipped to a norm of 1.0, get_lr(step) is written into every AdamW param group, then optimizer.step().
The schedule is meant to warm up linearly to max_lr = 6e-4 and follow a half cosine down to min_lr = max_lr * 0.1, which is 6e-5. As written, the warmup line multiplies min_lr rather than max_lr, so the ramp climbs to 6e-5 and the rate jumps to 6e-4 on the first cosine step. The demo below keeps that as committed. warm_up = 1 and max_steps = 10, because the committed run is ten steps. A loader walks one text file in B*T + 1 windows and wraps at the end; Generate() samples a multinomial over the last position and decodes with tiktoken.
| Chose | Over | Because | Cost |
|---|---|---|---|
Fused SDPA with is_causal=True | The hand-rolled masked softmax | Scale, mask and softmax in one kernel call, and it is what I would rely on in real code | The mask logic is no longer visible in the forward; the old block stays as a comment |
| Twelve separate heads in a Python loop | One fused QKV projection | One head at a time was easier to read and check against the commented-out math | 36 small matmuls per block instead of one; slower than the standard layout |
| Accumulate to 524,288 tokens | The largest micro-batch the GPU allows | The recipe's batch size is a training-dynamics choice, not a memory one | 128 forward and backward passes per optimizer step |
| bf16 autocast, no GradScaler | fp16 with loss scaling | bf16 keeps fp32's exponent range, so scaling is unnecessary | The README still mentions GradScaler; the README is wrong |
Try it
Three panels: the repo's learning-rate function, the mask its attention call applies, and the BPE merge algorithm run on whatever you type. None of them fakes a loss curve or a throughput number.
get_lr() from GPT.py ported line for line, with warm_up, max_steps and max_lr as sliders and min_lr fixed at a tenth of max_lr. The mask grid is the triangle is_causal=True applies; hover a row to see which keys that query may read. The BPE panel is the algorithm only, minbpe-style merges trained on your text; the repo uses tiktoken's pretrained GPT-2 encoding, so the merges shown are illustrative, not the shipped tokenizer.What I would do differently
Give the MLP the 4x expansion GPT-2 uses; as written it is 768 to 768 to 768, which is why the model is smaller than the configuration the README names, and why I print no parameter count here until I have measured it by instantiation. Multiply the warmup ramp by max_lr, not min_lr. Fuse the twelve heads into one QKV projection. Wrap Generate() in torch.no_grad() and model.eval(), since the 0.2 dropout is active while sampling. Write the one test that matters, that the output at position t does not change when tokens after t change, and run it against both attention implementations. Then either do a real run and commit the log, or delete the README's numbers.
- Single Colab export: a hard-coded Drive path, the model built and moved to CUDA at import time, notebook cells left in as comments. Six commits, all on one afternoon.
- The committed run is a 10-step smoke test on one Project Gutenberg novel (A Room with a View), followed by 100 sampled tokens. No checkpoint, log or plot is saved.
- No tests, no CI, no requirements file, no license.
- Two slips I only saw while writing this page: the warmup ramps to
min_lrinstead ofmax_lr, and the commented-out attention scales by the sequence length, not the head size. Neither is caught by anything in the repo. - The README's parameter count, speedup, tokens per second, loss curve and perplexity are not backed by anything in the repo, and the count does not match the code. I do not repeat them here.
- The recipe is Andrej Karpathy's nanoGPT and "Let's reproduce GPT-2". The code is my own idiom, not a copy, and not novel.
Sources GPT.py README Scope: I (sole author) Verified 3 Sep 2026