Summary by TwinMind · 4h 1m 26s

# Let's reproduce GPT-2 (124M)

Andrej Karpathy

[Let's reproduce GPT-2 (124M)](https://www.youtube.com/watch?v=l8pRSuU81PU) [Andrej Karpathy](https://www.youtube.com/channel/UCXUPKJO5MZQN11PqgIvyuvQ)

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

## GPT-2 Architecture and Loading Weights

- 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

## Model Architecture Details

- 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

## Training the Model from Scratch

- 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

## Initialization and Training Details

- 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)

## Performance Optimizations

- 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

## Scaling to Multiple GPUs

- 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

## Training on High-Quality Data

- 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

## Evaluation and Results

- 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

## Advanced Implementation and Future Work

- 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