AI

By the end of this build, a from-scratch GPT-2 124M — implemented in under 100 lines, trained on 8xA100 GPUs — matches OpenAI's original GPT-2 124M on validation loss and HellaSwag accuracy after roughly two hours on 10 billion tokens of FineWeb-Edu, and comes close to GPT-3's 124M checkpoint after an overnight, four-epoch, 40-billion-token run. Along the way, a chain of hardware- and algorithm-level optimizations cuts per-step time on an A100 from about 1,000ms to 93ms, an 11x speedup, at an estimated cloud cost of about $10 for the shorter run.

Loading the target model. Rather than starting from the original TensorFlow GPT-2 code released by OpenAI, the build starts from Hugging Face's PyTorch port, loading its state dict to inspect the real shapes: a 50,257×768 token embedding table (one row per BPE token, from the tokenizer covered in Karpathy's earlier tokenization series) and a 1,024×768 position embedding table. Plotting the position embeddings shows a jagged but broadly sinusoidal structure — evidence the vectors were learned from scratch rather than fixed as in the original 2017 Transformer, and a sign, from their raggedness, that this checkpoint isn't fully converged. Sampling five completions of "Hello, I'm a language model," from the Hugging Face pipeline gives a target: coherent, on-topic text to match once a custom implementation exists.

Rebuilding the architecture from scratch. GPT-2 departs from the original Transformer in two ways: it drops the encoder and cross-attention entirely (decoder-only), and it moves layer normalization to before each attention and MLP block instead of after, adding one more final layer norm before the classifier. The custom implementation mirrors Hugging Face's module names exactly (wte, wpe, h.0...h.11, ln_f) so weights can be copied over directly, and comes in at under 100 lines versus Hugging Face's roughly 2,000. Attention is described as a "reduce" — tokens exchange information by aggregating weighted values from earlier positions — while the MLP is a "map" applied identically to every token, making the Transformer a repeated map-reduce. The MLP uses GELU with the tanh approximation (GPT-2's original choice, made because the exact erf-based GELU was slow in TensorFlow at the time), preferred over ReLU because a smooth nonlinearity avoids the fully dead, zero-gradient neurons a flat ReLU tail can produce. Multi-head attention is implemented as one fused module using batched tensor reshaping — algorithmically identical to separate per-head modules, just faster — with a causal mask stopping any position attending to the future and a softmax normalizing attention weights to sum to one.

Loading real weights and validating generation. After copying OpenAI's released weights into the custom class (accounting for the fact that a few TensorFlow-derived weight matrices are transposed relative to PyTorch's convention), the custom model reproduces coherent completions, confirming the from-scratch architecture is correct — though its exact sampled tokens don't match Hugging Face's pipeline output bit-for-bit, a discrepancy Karpathy attributes to undocumented pipeline defaults (top-k 50 sampling among them) rather than a bug in the reimplementation. Switching to random initialization instead produces gibberish, as expected, and confirms the model is ready to train rather than merely replay OpenAI's weights.

A minimal training loop and sanity checks. Tokenizing the tiny Shakespeare dataset (about 1 million characters, roughly 338,000 GPT-2 tokens) and reshaping the token stream into (B, T) input batches, with targets simply the same sequence shifted by one position, gives the first working forward-backward-update loop using AdamW at a 3e-4 learning rate. Two sanity checks matter here: at initialization, the loss should be close to -ln(1/50257) ≈ 10.82, since an untrained model should assign roughly uniform probability across the vocabulary — the observed ~11 confirms a healthy random start; and the model should be able to overfit a single small batch to near-zero loss, which it does, confirming the optimization mechanics work before scaling up to real data.

Weight tying and initialization details. Inspecting the loaded state dict reveals that OpenAI's token embedding matrix and the final classifier's weight matrix are literally the same tensor — a weight-tying scheme from the original 2017 Transformer paper, motivated by the idea that semantically similar tokens should sit near each other in both the input embedding space and the output probability space. Pointing both layers at the same tensor saves about 40 million of the model's 124 million parameters, roughly 30%, and is expected to help a small model train more efficiently with fewer independent parameters to learn. For initialization, GPT-2's own source uses standard deviation 0.02 for weights and zero for biases (differing from PyTorch's default uniform bias init), values roughly consistent with the Xavier-style 1/√(fan-in) rule at typical GPT-2 channel counts. A second, less obvious detail scales every residual-path weight's standard deviation by an additional 1/√(2·n_layer): because each block adds its own contribution into the residual stream, repeatedly summing unscaled, unit-variance contributions inflates its variance — a toy example of 100 additions of standard-normal noise grows the running standard deviation to about 10 — so this correction keeps activation growth bounded regardless of depth.

Precision and compilation: from 1,000ms to 130ms per step. With batch size 16 and full GPT-2 sequence length 1,024, a plain float32 forward-backward step on an A100 costs about 1,000ms. Enabling TensorFloat-32 (TF32) — which truncates float32's mantissa from 23 to 10 bits inside the tensor-core matrix multiply only, transparently to the rest of the code — cuts this to about 333ms, a 3x gain rather than TF32's theoretical 8x, because the workload remains memory-bandwidth-bound: activations are still shuttled around in full float32 even though the multiply itself is faster. Switching to automatic mixed precision with bfloat16 (via torch.autocast, chosen over float16 because bfloat16 keeps float32's exponent range and needs no gradient scaler) brings this to about 300ms, converting activations like the logits to bf16 while parameters and operations like layer norm stay in float32. The largest single jump comes from torch.compile, which analyzes the whole model graph ahead of time rather than dispatching operations one by one from the Python interpreter, fusing elementwise sequences (GELU is the worked example) into single kernels that round-trip GPU high-bandwidth memory once instead of several times — dropping the step to about 130ms, a 2.3x gain from one line of code.

FlashAttention and vocabulary padding. torch.compile cannot find every fusion opportunity; the attention operation (torch.nn.functional.scaled_dot_product_attention, backed by FlashAttention) requires an algorithmic rewrite it won't discover on its own. FlashAttention performs more floating-point operations than the naive four-line attention implementation but is dramatically faster — the paper cites up to 7.6x — because it never materializes the full T×T attention matrix in high-bandwidth memory at all, instead computing softmax incrementally in an online, streaming fashion (a technique that traces back to a 2018 NVIDIA paper on online softmax normalization, four years before FlashAttention applied it). Swapping in FlashAttention drops the step to about 96ms. A final, smaller optimization pads GPT-2's awkward 50,257-token vocabulary up to 50,304 — a number divisible by 8, 16, 32, and 64 — because CUDA kernels operate in fixed-size block tiles and an input that doesn't divide evenly forces slow boundary-handling code paths; despite the added compute from the extra fake tokens (which the model simply learns to assign near-zero probability), this padding brings the step time down to about 93ms, an roughly 11x total speedup from the 1,000ms starting point.

Matching GPT-3's optimization hyperparameters. Because GPT-2's own paper is vague on optimizer details, hyperparameters are pulled from GPT-3's paper instead, given the architectures are otherwise very similar. This means AdamW with betas (0.9, 0.95) and epsilon 1e-8 in place of PyTorch's defaults, global gradient-norm clipping at 1.0 to guard against destabilizing loss spikes, and a cosine learning-rate schedule with linear warmup — rising to a peak of 6e-4 for a 124M-parameter model, then decaying to 10% of that peak, implemented as under a dozen lines rather than PyTorch's built-in scheduler so every step of the math stays visible. Weight decay of 0.1 is applied only to two-dimensional parameters (matrix weights and embeddings), leaving one-dimensional biases and layer-norm parameters undecayed, and a fused AdamW kernel (available in newer PyTorch, detected via inspect.signature) shaves the step time slightly further by updating all parameter tensors in one kernel launch instead of iterating over them individually.

Gradient accumulation, and the bug in getting it right. GPT-3's stated 0.5-million-token batch size doesn't fit in one GPU's memory, so gradient accumulation simulates it: multiple forward-backward passes accumulate gradients before a single optimizer step. A toy example with a 4-example batch versus four size-1 gradient-accumulation steps demonstrates the trap: because the loss functions used (MSE, and later cross-entropy) default to averaging over the batch, naively accumulating per-example losses via repeated .backward() calls sums rather than averages the contributions, producing gradients too large by the accumulation factor; the fix is dividing each micro-batch's loss by the number of accumulation steps before calling backward, recovering an identical result to true large-batch training.

Scaling to 8 GPUs with DistributedDataParallel. Launched via torchrun, eight copies of the training script run in parallel, each pinned to one GPU and identified by a DDP rank used to stride each process through a distinct slice of the data loader so no two GPUs see the same batch. DDP keeps the forward pass unchanged but, after the backward pass, averages gradients across all eight ranks via an all-reduce, so every rank ends the step with identical, batch-averaged gradients; to avoid wasteful synchronization on every micro-step during gradient accumulation, Karpathy directly toggles DDP's internal require_backward_grad_sync flag so the all-reduce fires only on each macro-step's final micro-step, rather than using PyTorch's documented (but, in his view, uglier) no_sync context-manager pattern. The accumulated loss value itself also needs an explicit dist.all_reduce average across ranks before the master process logs it, since otherwise each process would only report the loss it locally observed. With 8 A100s, gradient accumulation drops to 4 steps and throughput reaches about 1.5 million tokens per second.

Data: FineWeb-Edu, and evaluation with HellaSwag. GPT-2 and GPT-3 both trained on undisclosed, Reddit- and Common-Crawl-derived datasets, so the build substitutes a 10-billion-token sample of Hugging Face's FineWeb-Edu, a Common Crawl subset filtered for educational content (using Llama 3 70B as an automatic quality judge), tokenized into 100 shards of 100 million tokens apiece. Alongside a held-out validation shard, the model is periodically scored on HellaSwag, a 2019 multiple-choice sentence-completion benchmark where wrong answers are adversarially generated to fool language models while remaining easy for humans (95% human accuracy versus 48% for the best models of that year); since a 124M model can't natively perform multiple-choice reasoning, each candidate ending is instead scored by its average per-token loss under the model, with the lowest-loss completion taken as the prediction, and GPT-2 124M is measured at 29.5% (versus 25% random chance) as the baseline to beat.

Results and open caveats. Training for one epoch (about 19,073 steps, 10 billion tokens, roughly two hours) produces a validation loss below OpenAI's GPT-2 124M checkpoint's own score on this validation set, and a HellaSwag accuracy surpassing GPT-2 124M's 29.5% baseline — with 10x fewer training tokens than GPT-2's original 100-billion-token budget, tentatively credited to FineWeb-Edu's narrower, English-only, higher-quality domain and to broader post-2019 gains in data curation and filtering. An overnight extension to four epochs (about 40 billion tokens) pushes HellaSwag to 33.24%, nearly reaching GPT-3 124M's own reported score despite its 300-billion-token budget. Two rough edges stay unresolved: a periodic, sawtooth pattern in the loss curve, suspected to come from FineWeb-Edu shards not being shuffled or permuted across epochs, and an unexplained error that forces torch.compile off during generation and HellaSwag evaluation. The lecture closes by pointing to llm.c, a parallel CUDA/C reimplementation that reproduces identical losses and norms step-for-step while running faster than this PyTorch version, and by noting that everything built here is pretraining only — a conversational model, as in ChatGPT, would need a further supervised fine-tuning pass on conversational data, not covered here.