Let’s reproduce GPT-2 (124M) Karpathy · Jun 2024 YouTube Transcript

A lecture notebook

Zero to Hero · continuation

Let’s reproduce GPT-2 (124M)

A knowledge book for Andrej Karpathy’s lecture of 9 June 2024.

4h 1m 26s · YouTube l8pRSuU81PU · karpathy/build-nanogpt

Hover a timestamp for the explainer. Click it to play that moment. Search the bar to filter chapters, or open the annotated transcript at the same time.

Opening

The brief, and the checkpoint

What 124M actually is — twelve layers, 768 channels, not the 1.5B model — and why the GPT-3 paper is the concrete recipe. Then he opens the released OpenAI weights through Hugging Face and looks at wte, wpe, and a first sample.

Let's reproduce GPT-2 (124M)

Zero-to-Hero continuation. The target is GPT-2 124M — twelve layers, 768 channels — not the 1.5B model people usually mean. OpenAI released weights in 2019 but not the recipe; the paper's parameter table is wrong. He will borrow AdamW and the cosine + warmup schedule from the GPT-3 paper, because the architecture barely changed.

Split screen: train_gpt2.py in VS Code beside the GPT-2 paper.
VS Code + the GPT-2 paper (Language Models are Unsupervised Multitask Learners).

Exploring the GPT-2 (124M) OpenAI checkpoint

Original OpenAI code is TensorFlow; he loads GPT-2 124M through Hugging Face instead (HF name gpt2; gpt2-xl is the 1.5B). The state_dict shows wte at 50257×768 and wpe at 1024×768. Learned positions look sinusoidal but jagged — he reads that as undertrained. He samples the HF pipeline with the "Hello, I'm a language model," prefix; later his own loop will use top-k 50.

Grayscale imshow of GPT-2 position embeddings, 1024 by 768.
play.ipynb: plt.imshow of transformer.wpe.weight (1024 × 768). Horizontal structure = learned positional features.

Build

Implement the network

Write a decoder-only GPT in about 100 lines, load the Hugging Face weights, sample the “Hello, I’m a language model,” prefix, then overfit Tiny Shakespeare until the backward path is trusted. Weight tying and residual init close the architecture.

SECTION 1: implementing the GPT-2 nn.Module

Decoder-only Transformer: drop the encoder and its cross-attention. Two GPT-2 changes: pre-norm (LayerNorms move before attention/MLP, plus a final ln_f) and a clean residual stream so the loss gradient flows unchanged down to the tokens. Attention is the reduce; MLP is the map. GELU (tanh) matches GPT-2's historical TensorFlow quirk. Multi-head attention is one fused module. Target: the GPT-2 124M net in about 100 readable lines versus Hugging Face's ~2000.

Loading the Hugging Face / GPT-2 parameters

He matches Hugging Face names — wte, wpe, h.0…, ln_f, lm_head, c_attn — so GPT-2 124M tensors copy over. Some TensorFlow weights are transposed and must be flipped. Buffers such as the causal attn.bias mask are ignored. from_pretrained("gpt2") is a classmethod that returns his module with OpenAI's weights installed. weight tying will later show up as wte and lm_head sharing storage; loading without a crash is the first confidence check.

Implementing the forward pass to get logits

Input idx is integers of shape (B, T) with T ≤ block size 1024. Position indices 0..T−1 are created on the same device as idx. Lookups from wte plus wpe (broadcast across batch) go through the twelve blocks, ln_f, then lm_head. The residual stream carries activations; pre-norm sits inside each block. Output logits are (B, T, 50257) — a score for every next-token at every position. Softmax is still ahead.

Sampling init, prefix tokens, tokenization

tiktoken's gpt2 encoding tokenizes "Hello, I'm a language model," into 8 token ids. Those 8 tokens are stacked 5 times to make x of shape (5, 8) — lookups into wte, with wpe added by position. The GPT-2 124M module is set to eval() — though this net has no Dropout, so eval may be a no-op — then moved to CUDA. This is the same prefix Hugging Face's docs use; the loop will apply top-k 50.

Sampling loop

Loop under torch.no_grad. Forward, keep only the last-position logits, softmax, then top-k 50 (the Hugging Face pipeline default) so rare tokens are never sampled. Append the new column until length 30. Completions are coherent English but do not match the HF pipeline snippet — leftover pipeline defaults, not a bug. Replicating the same sampler on the HF model matches, so the residual stream, GELU (tanh), and pre-norm path are loading correctly.

Sample, auto-detect the device

Swap from_pretrained for GPT(GPTConfig()) — PyTorch's default random init. Samples are garbage token salad, as they should be. He autodetects device: CUDA, else MPS (Apple Silicon), else CPU. Tensors created inside forward use idx.device so CPU/GPU don't mix. A CPU override still works (~6s for this tiny sample). Serious multi-GPU DDP comes later; you can follow on a laptop until then. The GPT-2 124M architecture is unchanged; Tiny Shakespeare will be the debug set.

Let's train: data batches (B,T) → logits (B,T,C)

Dataset is Tiny Shakespeare (~1M ASCII characters). The GPT-2 124M tokenizer compresses ~3:1, so 1000 chars ≈ 300 tokens. A 1-D token stream is .view(B, T) into a batch. Labels Y are the same stream shifted +1, so he loads B·T+1 tokens: X is buf[:-1], Y is buf[1:]. Debug batch: B=4, T=32. Token 198 is newline. wte is still a separate table; weight tying is not yet applied.

Cross entropy loss

F.cross_entropy wants 2-D logits, so (B, T, vocab) flattens to (B·T, vocab) and Y to (B·T,). At random init, a uniform next-token distribution over the GPT-2 124M vocab should have loss ≈ −ln(1/50257) ≈ 10.82. He measures ~11 — close enough that the model is not confidently wrong. Tiny Shakespeare batches feed that vocab. Next: AdamW on a single batch; the residual stream still has default init.

Optimization loop: overfit a single batch

AdamW at lr 3e-4. Must zero_grad because backward is += into .grad. Bug: a CPU token buffer was never moved — tensor.to(device) is not in-place; you must assign. Overfitting one Tiny Shakespeare batch drives loss toward 0, which proves the GPT-2 124M backward path can memorize. gradient accumulation is not needed yet.

Data loader lite

Simple loader: tokenize the whole file, walk it in strides of B·T, fetch B·T+1 tokens, wrap to 0 at EOF. Tokens stay on CPU and ship to GPU per batch. 338,025 tokens; at B=4, T=32 that's ~2,640 batches per epoch. Fifty AdamW steps only reach ~6.6 loss — expected: lots of unused GPT-2 124M vocab ids can have their biases driven down, but you cannot memorize Tiny Shakespeare in 50 steps. DDP is later.

Parameter sharing: wte and lm_head

In the HF (and OpenAI TF) checkpoint, wte and lm_head.weight are the same storage — identical data_ptr. weight tying comes from Attention Is All You Need (and Press & Wolf 2017): tokens that are close in embedding space should also get similar output logits. He sets transformer.wte.weight = lm_head.weight so one tensor is used twice; backward gradients add from both ends. Saves 768×50257 ≈ 40M params, ~30% of GPT-2 124M. wpe stays separate.

Model initialization: std 0.02, residual init

OpenAI's model.py inits Linear/Embedding weights ~ N(0, 0.02), biases 0; they used 0.01 for pos emb, he keeps 0.02. 0.02 is roughly 1/sqrt(fan_in) at GPT-2 124M widths. residual init scales residual projection weights by 1/sqrt(N) with N = 2 · n_layer (attn + MLP per layer). Weight-tied wte/lm_head get initialized twice at 0.02 — fine. GELU (tanh) is already in the MLP; pre-norm is already in the block.

Speed

Make it fast

Same network, much faster step. Precision (TF32, then bfloat16), torch.compile, Flash Attention, and a padded vocab — timed on his 8×A100 box at B=16, T=1024. Most of the work is memory-bound; the table is his machine, not a laptop.

Karpathy SSH’d into the 8×A100 training box, A100 spec table open.
SSH’d into the training box.
On his 8×A100 box, B=16, T=1024 — not your laptop.
ChangeStep time
FP32 baseline~1000 ms
TF32~333 ms
BF16 autocast~300 ms
torch.compile~130 ms
Flash Attention~96 ms
nice vocab (50304)~93 ms
FusedAdamW~90 ms

SECTION 2: Let's make it fast. GPUs, mixed precision, 1000ms

Hardware first. His box: 8× A100 80GB SXM (Lambda). Default PyTorch is FP32. A100 peak: ~19.5 TFLOPS FP32 vs ~156 TF TF32 vs ~312 TF bfloat16 (dense). Most well-tuned training is memory-bound — 60% utilization is excellent. He times a realistic GPT-2 124M step at B=16, T=1024: baseline ~1000 ms/iter, ~163k tokens/s. Use cuda.synchronize() before the clock. torch.compile comes later. If you OOM, drop B; keep B a nice number.

Tensor Cores, timing the code, TF32 precision, 333ms

A tensor core is a 4×4 matmul instruction. Almost all GPT work is matmuls inside Linear layers; the lm_head (768→50257) dominates at this size. TF32 keeps FP32's 8-bit exponent and sign, crops the mantissa (23→10 bits) inside the instruction only. Enable with torch.set_float32_matmul_precision("high"). No dtype changes in your tensors — wte still looks like FP32. Promised 8×, observed ~3× (1000→333 ms) because the workload is still moving FP32 bytes. Empirically you cannot tell TF32 from FP32 in GPT-2 124M training.

float16, gradient scalers, bfloat16, 300ms

To go faster you must move fewer bits. FP16 shrinks the exponent range, so you need loss/gradient scaling. bfloat16 keeps FP32's exponent and only crops mantissa — same range, no scaler. Use torch.autocast(device_type="cuda", dtype=torch.bfloat16) around forward + loss only. After autocast, logits are BF16 but wte.weight stays FP32. Softmax, LayerNorm, and the loss stay FP32. Result: ~300 ms, ~55k tok/s. GPT-2 124M trains fine this way.

torch.compile, Python overhead, kernel fusion, 130ms

torch.compile is a compiler for the model graph. It reduces Python overhead and fuses kernels so you don't round-trip to HBM between every tiny op (e.g. GELU (tanh) + Linear). One line: model = torch.compile(model). First iteration pays a compile cost; steady state drops to about 130 ms. The residual stream graph becomes one compiled region. GPT-2 124M is now a compiled module.

Flash Attention, 96ms

Standard attention materializes a T×T matrix (1024×1024) then softmax. Flash Attention is an IO-aware algorithm that never writes that full matrix to HBM — it tiles softmax in SRAM. Same math, much less memory traffic. Step time ~96 ms. The residual stream still adds the attention output; pre-norm is unchanged. GPT-2 124M context is 1024, so the savings are real. torch.compile and Flash Attention stack.

Nice/ugly numbers. vocab 50257 → 50304, 93ms

CUDA kernels like dimensions with many factors of 2. 50257 is an ugly number (50,000 BPE merges + 256 byte tokens + 1 EOT). He pads vocab to 50304 — nice numbers. Extra logits are unused; kernels run faster. ~93 ms. wte and the classifier grow by the pad rows; weight tying still holds. GPT-2 124M vocab semantics don't change. Flash Attention is already on.

Run / Morning

Train like they trained

Match the papers: AdamW, gradient clipping, warmup-cosine, weight decay, gradient accumulation, DDP. Swap the Shakespeare toy for FineWeb-EDU, hold out a val split, score HellaSwag the completion way, and start the overnight run.

SECTION 3: hyperparameters, AdamW, gradient clipping

Architecture is done and fast. Now match the papers. GPT-2 is vague; GPT-3 is the concrete optimizer/schedule document. AdamW with gradient-norm clipping (clip exploding steps). This is the start of "train like they trained," not "overfit Tiny Shakespeare." The GPT-2 124M architecture is finished. cosine + warmup comes next; DDP after that.

Learning rate scheduler: warmup + cosine decay

Linear warmup from 0, then cosine decay toward a floor — cosine + warmup, copied from the GPT-3 recipe so early AdamW steps don't smash a randomly-initialized net. In the final script: max_lr = 6e-4, min_lr = 0.1·max_lr, warmup_steps = 715, max_steps = 19073 (~1 epoch). GPT-2 124M starts random; residual init already scaled the projections.

Batch size schedule, weight decay, FusedAdamW, 90ms

GPT-3 ramps batch size in the early part of training. Weight decay is applied to 2-D parameters (the matrices), not to biases or LayerNorms. FusedAdamW is a single CUDA kernel for the AdamW update — less overhead. ~90 ms. cosine + warmup is already the schedule. GPT-2 124M 2-D weights include wte.

Gradient accumulation

GPT-3 124M used a huge token batch (~0.5M tokens). That will not fit in one step, so gradient accumulation over several micro-batches, stepping AdamW only after the last. Math: scale the loss by 1/n_micro so the accumulated grad equals the large-batch grad. GPT-2 124M is the architecture; this is the GPT-3 batch size. DDP will wrap this.

Distributed data parallel (DDP)

torchrun launches one process per GPU. Each rank has a rank id and world_size. DDP all-reduces gradients so every replica stays in sync. Combined with gradient accumulation: use no_sync() on the micro-steps that are not the last, so you don't all-reduce N times. Only rank 0 should print, log, and sample. This is how the 8×A100 box actually trains GPT-2 124M. AdamW steps stay synchronized.

Datasets: GPT-2, GPT-3, FineWeb-EDU

GPT-2 124M trained on WebText (scraped outbound Reddit links). GPT-3 mixed Common Crawl with books/Wikipedia/etc. For a modern reproduction he uses FineWeb-EDU (educational filtered web). Data is tokenized once and written as shards (typically .npy) so the loader just memory-maps arrays. Tiny Shakespeare was only a debug toy. wte still covers the same vocab; HellaSwag eval is later.

Validation data split, validation loss, sampling revive

Hold out a val split and periodically measure val loss — the same next-token metric as training, on unseen documents. Sampling is revived during the run so you can read generations, not just stare at a loss curve. Keep sampling RNG separate from training RNG so generating text does not perturb the optimization. FineWeb-EDU is the train source; top-k 50 stays the sampler; HellaSwag is next. Only the DDP rank-0 process should sample.

Evaluation: HellaSwag, starting the run

HellaSwag is a commonsense 4-way ending pick. He scores it the completion way: condition on the context, sum/average log-probs of each ending, pick the most likely — not free generation. use_compile = False in the final run because torch.compile can interfere with HellaSwag eval. Logging, checkpointing, and eval cadence are wired up. Then he starts the overnight run on FineWeb-EDU. AdamW and cosine + warmup are already running.

Morning

Results, llm.c, the repo

After the overnight run: min train 2.922, min val 3.0726, max HellaSwag 0.3068. Beats the OpenAI GPT-2 124M val-loss line and its HellaSwag (~0.295). GPT-3 124M HellaSwag ~0.335 is still above. Then llm.c, and the build-nanogpt repo.

SECTION 4: results in the morning! GPT-2, GPT-3 repro

Overnight numbers from his plots: min train loss 2.922, min val 3.0726, max HellaSwag 0.3068. Beats the OpenAI GPT-2 124M val-loss dashed line and its HellaSwag (~0.295). GPT-3 124M HellaSwag ~0.335 is still above. Empty file → competitive 124M. He reads some amusing generations. FineWeb-EDU was the data; llm.c is the encore.

Morning loss and HellaSwag plots from the overnight 124M run.
Morning plots. Min train 2.922 / min val 3.0726 / max HellaSwag 0.3068. Beats the OpenAI GPT-2 124M val-loss and HellaSwag (~0.295) lines; GPT-3 124M HellaSwag ~0.335 still above.

Shoutout to llm.c

The same algorithm exists as llm.c — a from-scratch C/CUDA implementation that is faster than the PyTorch teaching code. PyTorch was for understanding; llm.c is for squeezing the hardware. Flash Attention and TF32 ideas still apply; GPT-2 124M is the model. DDP-style multi-GPU is the deployment shape.

Summary, phew, build-nanogpt GitHub repo

Recap of the four acts. The GitHub repo karpathy/build-nanogpt is the artifact: walk the commits in order and you walk the video. By the end it is about 90% of nanoGPT. GPT-2 124M from an empty file. FineWeb-EDU, HellaSwag, and llm.c are the modern extras.

Glossary

The same hover terms used in the chapters, listed.

GPT-2 124M
Smallest GPT-2. 12 layers, 12 heads, 768 dim, context 1024. Paper table of counts is wrong; 124M is the corrected number.
wte
Token embedding table, vocab×768. Also reused as the softmax classifier (weight tying).
wpe
Learned absolute position embeddings, 1024×768. GPT-2 trains these; the original Transformer used fixed sinusoids.
pre-norm
LayerNorm before attention and before MLP, plus a final ln_f. GPT-2's main diagram change vs Attention Is All You Need.
residual stream
The x = x + f(x) highway. Clean (no norms on the highway) so gradients travel to the tokens unchanged.
GELU (tanh)
Smooth ReLU. GPT-2 used the tanh approximation because TensorFlow's erf was slow in 2018.
top-k 50
At sample time, keep only the 50 most likely tokens, zero the rest, renormalize. HF pipeline default; keeps text on the rails.
Tiny Shakespeare
~1M-character debug corpus. Not the real pretraining set.
weight tying
Share wte and lm_head.weight. From AIAYN / Press & Wolf. Saves ~40M params.
residual init
Scale c_proj weights by 1/sqrt(2·n_layer) so residual-stream variance stays ~1.
TF32
Ampere tensor-core mode that crops FP32 mantissa inside matmuls. One line of PyTorch; tensors still look like FP32.
bfloat16
16-bit float with FP32's exponent. Mixed precision without gradient scaling.
torch.compile
PyTorch graph compiler. Fuses kernels, kills Python overhead.
Flash Attention
IO-aware attention that never writes the full T×T matrix to HBM.
nice numbers
Pad ugly dims (50257) to GPU-friendly ones (50304).
AdamW
Adam with decoupled weight decay. Default LLM optimizer.
cosine + warmup
LR rises linearly, then cosine-decays. GPT-3 schedule.
gradient accumulation
Several micro-batches, one optimizer step. Fakes a larger batch.
DDP
One process per GPU; all-reduce grads. torchrun + rank + world_size.
FineWeb-EDU
Modern educational web crawl used here instead of unavailable WebText.
HellaSwag
Commonsense 4-way ending pick. Scored by completion log-prob, not generation.
llm.c
Karpathy's C/CUDA rewrite of this stack, faster than the teaching PyTorch.

Compiled as a study book for the lecture; numbers and claims are those Karpathy states in the video. Walk the commits in build-nanogpt in the same order.