Lecture 10: Reproducing GPT-2 (124M) from Scratch

What this lecture covers

This final lecture in the Zero to Hero series answers a single question: what does it actually take to reproduce GPT-2's smallest model, 124 million parameters, from nothing? Andrej Karpathy first loads OpenAI's released weights through Hugging Face to establish a target, then writes the GPT-2 architecture from scratch in PyTorch, matching Hugging Face's parameter names so the pretrained weights can be copied in directly. It builds on the tokenization and Transformer internals covered earlier in the playlist ("Let's build GPT from scratch" and the tokenizer video), and assumes you're comfortable with attention, multi-head attention and basic PyTorch modules.

After the model is correct, the lecture shifts to making training fast: precision tricks, torch.compile, Flash Attention, nice power-of-two numbers, and multi-GPU training. It then follows the GPT-3 paper for optimization hyperparameters, builds a proper data pipeline on the FineWeb-Edu dataset, and finally trains the model overnight on 8 GPUs. By the end you can write your own decoder-only Transformer, load pretrained weights into it, and run a real pretraining loop with evaluation against a benchmark, HellaSwag, that you can also implement yourself.

Key ideas

• Weight tying: GPT-2 shares the same matrix between the token embedding table and the final classifier layer, which saves about 30% of parameters and is motivated by wanting similar tokens to have similar embeddings and similar output probabilities.
• Pre-normalization residual stream: GPT-2 moves the layer norms before the attention and MLP blocks (rather than after), keeping the residual pathway free of normalization so gradients flow through it unchanged.
• Attention as reduce, MLP as map: attention exchanges information between tokens (a pooling operation), while the MLP processes each token's representation independently, so a Transformer block is a repeated map-reduce.
• Mixed precision and TF32/BF16: dropping numeric precision (torch.set_float32_matmul_precision('high') for TF32, torch.autocast for BF16) speeds up matrix multiplies on tensor cores because many workloads are memory-bandwidth bound, not compute bound.
• torch.compile and kernel fusion: compiling the model removes Python interpreter overhead and fuses element-wise operations so intermediate results stay on the GPU chip instead of round-tripping to memory.
• Flash Attention: a kernel-fusion rewrite of attention that avoids ever materializing the full attention matrix in memory, making it faster despite doing more floating-point operations.
• Gradient accumulation: simulates a large batch size (GPT-3 uses about 0.5 million tokens per step) on limited GPU memory by summing gradients over several smaller forward-backward passes before one optimizer step, with the loss scaled down to keep the math equivalent to one big batch.
• Distributed Data Parallel (DDP): runs one copy of the model per GPU, each on a different data shard, and averages gradients across GPUs after the backward pass so every process ends up with the same updated weights.

Walkthrough

Loading OpenAI's GPT-2 and writing the model from scratch (0:00)

The lecture opens by loading the released GPT-2 124M weights through Hugging Face Transformers and sampling from them to establish a working target, since the original OpenAI code was written in TensorFlow. It inspects the state dict, visualizes the learned position embeddings, and notes that GPT-2 differs from the original Transformer paper by being decoder-only (no cross-attention) and by moving layer norms before each sub-block, with an extra layer norm before the final classifier.

Implementing the GPT-2 module (13:47)

Karpathy builds the model class to mirror Hugging Face's naming (wte, wpe, h, ln_f, lm_head) so pretrained weights load directly, then implements the block (attention plus MLP, GELU with the tanh approximation), the multi-head self-attention as a single fused module, and the forward pass that returns logits and, when targets are given, a cross-entropy loss. He then adds weight tying, GPT-2-style weight initialization (0.02 standard deviation, zeroed biases) with an extra scaling factor on residual projection weights to control variance growth across layers, and confirms the model can load OpenAI's checkpoint and generate coherent text.

Building a data loader and overfitting a batch (45:50)

Using the tiny Shakespeare dataset for debugging, the lecture shows how to reshape a flat token stream into batch-by-time tensors with shifted targets, checks that the initial loss is close to the theoretical -log(1/50257), and writes a simple optimization loop with AdamW that first overfits a single batch to sanity-check the pipeline, then a minimal DataLoaderLite that walks through the dataset in chunks.

Speeding up training: precision, compile, Flash Attention, nice numbers (1:22:18)

This is the core optimization section. Switching from float32 to TF32 gives roughly a 3x speedup for one line of code; switching activations to bfloat16 with torch.autocast adds a smaller further gain. torch.compile removes Python overhead and fuses kernels for another large jump, and replacing manual attention with F.scaled_dot_product_attention (Flash Attention) cuts time again by avoiding materializing the full attention matrix. Finally, padding the vocabulary size from the "ugly" 50257 to the power-of-two-friendly 50304 speeds up CUDA kernels further. Combined, these changes take the code from roughly 1000ms to under 100ms per step, about an 11x improvement.

Following GPT-3's hyperparameters and scaling to multiple GPUs (2:14:55)

Because the GPT-2 paper is vague on optimization details, the lecture turns to GPT-3's paper for AdamW betas, gradient clipping, and a cosine-decay learning rate schedule with linear warmup. It implements gradient accumulation to simulate GPT-3's roughly 0.5-million-token batch size on a single GPU, carefully scaling the loss so accumulated gradients match a true large-batch gradient, and separates parameters into weight-decayed (2D matrices) and non-decayed (biases, layer norm) groups.

Distributed training and dataset upgrade (2:46:52)

The lecture introduces PyTorch's DistributedDataParallel and torchrun to run eight processes across eight GPUs, each handling a different data shard, with gradients averaged (all-reduce) after the backward pass only on the final gradient-accumulation micro-step. It then swaps tiny Shakespeare for a 10-billion-token sample of the FineWeb-Edu dataset, pre-tokenized into shards, to run a real pretraining job.

Evaluation with HellaSwag and the overnight results (3:28:23)

HellaSwag is introduced as a sentence-completion benchmark scored by comparing average token probabilities across candidate completions rather than multiple-choice labels, since small models can't reason about labels directly. Wired into the training loop alongside periodic validation loss and sample generation, HellaSwag tracks progress over a roughly two-hour, 10-billion-token run, which surpasses GPT-2 124M's validation loss and HellaSwag accuracy, and an eight-hour, four-epoch overnight run that approaches GPT-3's 124M accuracy. The lecture closes by checkpointing the model, pointing at the companion llm.c CUDA implementation, and noting open issues (loss periodicity from unshuffled data, torch.compile breaking generation) for viewers to pursue.

Before you watch

• Watch "Let's build GPT: from scratch, in code, spelled out" first, since this lecture assumes you already understand attention, multi-head attention, and the Transformer forward pass.
• Review the tokenization lecture in this series; GPT-2's byte-pair encoding vocabulary and special tokens are used without re-explanation here.
• Have a basic familiarity with PyTorch training loops (optimizers, backward(), state_dict) since the lecture moves quickly through standard boilerplate.
• Access to a CUDA GPU makes the second half meaningfully easier to follow, though the lecture notes CPU and Apple Silicon (MPS) also work for the earlier parts.

Check your understanding

1. Why does weight tying between the token embedding and the final classifier make sense, and how much does it reduce the model's parameter count?
2. Why does gradient accumulation require dividing the loss by the number of accumulation steps before calling backward()?
3. Explain why lowering numeric precision (TF32, bfloat16) can speed up training even though the model performs the same number of floating-point operations.
4. Why does Flash Attention run faster than the naive four-line attention implementation despite doing more total FLOPs?
5. How is HellaSwag scored for a small language model, and why can't the model just be asked to pick "A", "B", "C" or "D"?