Summary by TwinMind · 4h 1m 26s

Let's reproduce GPT-2 (124M)

Andrej Karpathy

Andrej Karpathy1.38M subscribers

Introduction to GPT-2 Reproduction

Goal: Reproduce GPT-2 (124 million parameter version)
OpenAI released GPT-2 in 2019 with blog post, paper, and GitHub code
GPT-2 consists of models at different sizes (miniseries)
- Smallest: 124M parameters (12 layers, 768 dimensions)
- Largest: 1.5B parameters
Different model sizes help establish scaling laws
- Plotting model size vs downstream metrics shows performance trends
Paper had error in parameter counts (corrected in GitHub repo)
Modern reproduction takes ~1 hour, costs ~$10 on cloud compute

OpenAI released model weights, unlike many other models
GPT-2 paper lacks training details; will reference GPT-3 paper for some parameters
Loading pre-trained GPT-2 model using Hugging Face Transformers
- Model architecture includes transformer blocks, embeddings, and output head
Testing pre-trained model with text generation
- Example prompt: "Hello I'm a language model"
Creating custom PyTorch implementation of GPT-2 for better understanding
- Implementing transformer blocks, attention mechanism, and MLP components
- Using GELU activation function (specifically the approximate version with tanh)
Successfully loading OpenAI's weights into custom implementation
- Verifying with text generation using same prompt

GPT-2 is a decoder-only transformer
- No encoder or cross-attention components
- Uses pre-normalization (layer norms before attention/MLP blocks)
- Additional layer normalization after final self-attention block
Block structure:
- Self-attention mechanism with multiple heads
- MLP with GELU activation
- Residual connections around each component
Attention is communication operation (tokens exchange information)
MLP operates independently on each token (no information exchange)
Transformer = repeated application of map-reduce pattern
- Attention = reduce operation
- MLP = map operation

Setting up training data (initially using tiny Shakespeare dataset)
Creating data loader for batching sequences
- Handling token sequences and targets for next-token prediction
Implementing loss function using cross-entropy
Setting up optimizer (AdamW)
Initial training with random initialization
- Expected loss at initialization: ~10.82 (1/vocab_size)
- Training loop with gradient calculation and parameter updates
Weight tying between token embeddings and output projection
- Reduces parameter count by ~30% (40M parameters)
- Same weights used at input embedding and output classification
- Supported by research showing benefits for similar tokens

Implementing initialization following GPT-2 code:
- Linear layers: normal distribution with std=0.02
- Biases: zeros
- Layer norms: weight=1.0, bias=0.0
Special initialization scaling for residual paths
- Scale weights by 1/sqrt(number of residual layers)
- Prevents variance explosion in residual stream
Setting up learning rate scheduler
- Cosine decay with warmup following GPT-3 paper
- Warmup over first 375M tokens, decay to 10% of max value
Gradient clipping at global norm of 1.0
Weight decay of 0.1 (not applied to biases and 1D tensors)

Using tensor cores with reduced precision
- TF32: 8x theoretical speedup with minimal accuracy impact
- Achieved ~3x speedup (memory bandwidth limited)
Mixed precision training with BFloat16
- Maintains full dynamic range unlike FP16
- No need for gradient scaling
- Further improved performance
Implementing torch.compile
- Reduced Python overhead and GPU read/writes
- 2.3x improvement with single line of code
- Kernel fusion optimization explained
Using Flash Attention
- 27% improvement over standard attention
- Avoids materializing full attention matrix
- Uses online softmax calculation
Optimizing vocabulary size for GPU efficiency
- Changed from 50,257 to 50,304 tokens (power-of-two friendly)
- 4% speedup despite doing more computation
- Helps with GPU kernel efficiency

Implementing Distributed Data Parallel (DDP)
Setting up process groups with torch.distributed
Modifying data loader to handle multiple processes
- Each GPU processes different parts of dataset
Implementing gradient accumulation
- Simulating larger batch sizes than what fits in GPU memory
- Careful handling of loss normalization
Synchronizing gradients across GPUs
Handling validation and metrics across processes

Switching from Shakespeare to Fine-Web-EDU dataset
- 10B token sample of high-quality educational content
- Pre-processed and tokenized into shards
Setting up training for full run
- Batch size matching GPT-3 paper (0.5M tokens)
- Processing ~1.5M tokens/second with 8 GPUs
- Full training run: ~1.7 hours for 10B tokens

Implementing validation loss tracking
Adding HellaSwag evaluation
- Multiple-choice benchmark for natural language understanding
- Using native token completion approach
Implementing text sampling to check generation quality
Results after training:
- Surpassed GPT-2 124M validation loss
- Achieved 33.2% on HellaSwag (surpassing original GPT-2's 29.5%)
- Approaching GPT-3 124M performance with only 40B tokens (vs 300B)
- Generated coherent text completions

Added checkpointing functionality
Reference to LLM.c - a pure CUDA implementation
- Faster than PyTorch version
- Same numerical results
Noted limitations and improvements:
- Data shuffling improvements needed
- Torch.compile compatibility issues with evaluation
- Potential for higher learning rates