Overview
This tutorial reconstructs the 124-million-parameter GPT-2 model from first principles and turns the implementation into a practical pretraining system. It begins by loading OpenAI's released weights through Hugging Face, inspecting their structure, and using them as a compatibility target for a concise PyTorch model. The architecture is then rebuilt with token and positional embeddings, 12 pre-normalized Transformer blocks, causal multi-head attention, GELU-based MLPs, a final layer normalization, and tied input-output weights. After establishing a correct forward pass, loss calculation, and single-batch overfitting test, the tutorial develops a streaming data loader and aligns initialization and optimization settings with evidence from the GPT-2 source and GPT-3 paper. Performance is progressively improved through TF32, bfloat16 autocasting, torch.compile, FlashAttention, vocabulary padding, fused AdamW, gradient accumulation, and distributed data parallelism across eight A100 GPUs. The final system preprocesses a 10-billion-token FineWeb-Edu sample, tracks validation loss and HellaSwag accuracy, generates samples, and saves checkpoints. A roughly two-hour run surpasses the original GPT-2 124M checkpoint on the demonstrated HellaSwag setup, while a longer run approaches the comparable GPT-3 result. However, possible benchmark contamination, unshuffled data, periodic loss behavior, and torch.compile incompatibilities prevent treating the outcome as an exact historical reproduction.
Sections
Architecture and Training Configuration
Concrete implementation choices used to construct and train the GPT-2-class model.
- Use 12 Transformer layers, 12 attention heads, a 768-dimensional residual stream, a maximum sequence length of 1,024, and the GPT-2 tokenizer with 50,257 valid tokens.
- Build a decoder-only, pre-normalized block with causal multi-head self-attention followed by an MLP using the tanh approximation of GELU; add a final layer normalization before the language-model head.
- Tie the token embedding weight to the language-model head weight so both modules reference the same tensor.
- Initialize linear and embedding weights from a normal distribution with standard deviation 0.02, zero linear biases, and scale residual-output projections by 1/sqrt(2 × number_of_layers).
- Calculate next-token cross-entropy by flattening logits from [B, T, vocabulary] to [B×T, vocabulary] and targets from [B, T] to [B×T].
- Use AdamW with beta values 0.9 and 0.95, epsilon 1e-8, weight decay 0.1 on matrix and embedding weights, no decay on one-dimensional biases and normalization parameters, and global gradient-norm clipping at 1.0.
- Apply a linear learning-rate warmup followed by cosine decay toward 10% of the maximum rate; the GPT-3-derived maximum for this model size is 6e-4.
- Use a total batch of 524,288 tokens and compute gradient accumulation steps as total_batch_tokens divided by microbatch_size × sequence_length × distributed_world_size.
- Accelerate training with TF32 matrix multiplications, bfloat16 autocasting around the forward pass and loss, torch.compile where compatible, scaled-dot-product attention for FlashAttention, and fused AdamW.
- Pre-tokenize FineWeb-Edu into NumPy uint16 shards containing 100 million tokens each; reserve shard 0000 for validation and use the remaining shards for training.
Execution Checklist and Remaining Work
Recommended steps for reproducing the run and addressing known limitations.
- First verify checkpoint compatibility by loading Hugging Face GPT-2 weights into the custom model and confirming that direct generation produces coherent text.
- Sanity-check the random initialization loss against approximately 10.82, then confirm the model can overfit one tiny batch before launching full pretraining.
- Choose the largest hardware-friendly microbatch that fits GPU memory, while using gradient accumulation to preserve the intended total token batch.
- Measure tokens per second after each optimization rather than assuming theoretical FLOP improvements will translate directly into runtime gains.
- Shuffle documents and shards at epoch boundaries so repeated epochs do not preserve identical document adjacency and ordering.
- Investigate the periodic training-loss pattern to determine whether it originates from FineWeb-Edu ordering, shard construction, or loader reset behavior.
- Resolve the torch.compile failures affecting generation and HellaSwag evaluation so compiled training does not require disabling reader-facing evaluations.
- Evaluate saved checkpoints with a broader external harness and multiple task families before claiming parity with GPT-2 or GPT-3.
- Save optimizer state, random-number-generator state, step counters, and loader position in addition to model weights if exact training resumption is required.
Core Concepts
Terms required to understand the implementation and performance work.
- Decoder-only Transformer: a Transformer that uses causal self-attention without an encoder or encoder-decoder cross-attention, predicting each next token from prior tokens.
- Pre-normalization: placing layer normalization before the attention or MLP transformation so the residual stream retains a direct additive path.
- Weight tying: reusing one parameter tensor for the input token embeddings and output classifier weights.
- Gradient accumulation: summing normalized gradients from multiple microbatches before one optimizer step to simulate a larger total batch.
- TF32: an NVIDIA Tensor Core computation mode that keeps the FP32 exponent range but truncates mantissa precision inside matrix multiplications.
- Bfloat16: a 16-bit floating-point format that preserves the eight-bit exponent range of FP32 while reducing mantissa precision, generally avoiding FP16-style gradient scaling.
- Kernel fusion: combining several operations into one GPU kernel so intermediate values remain on-chip and require fewer high-bandwidth-memory reads and writes.
- FlashAttention: an exact attention algorithm organized around the GPU memory hierarchy that avoids materializing the full attention matrix in high-bandwidth memory.
- DistributedDataParallel: PyTorch's multi-process training mechanism in which workers compute local gradients and synchronize them across devices.
- HellaSwag: an adversarial sentence-completion benchmark in which a model must assign the highest likelihood to the natural continuation among four options.
Tools, Papers, and Data
Implementations, references, datasets, and infrastructure explicitly used or recommended.
- OpenAI GPT-2 repository: the original TensorFlow implementation and released GPT-2 weights.
- GPT-2 paper: the primary architectural reference, though it provides limited training detail.
- GPT-3 paper: the source for more explicit optimizer, learning-rate, batch-size, and evaluation settings.
- Hugging Face Transformers: provides a readable PyTorch GPT-2 implementation and converted pretrained weights.
- tiktoken: OpenAI's tokenizer library used to encode and decode GPT-2 tokens.
- FineWeb-Edu sample-10BT: a filtered educational web corpus containing approximately 10 billion tokens.
- HellaSwag: the sentence-completion benchmark used to compare the reproduced model with GPT-2 and GPT-3 checkpoints.
- EleutherAI Language Model Evaluation Harness: infrastructure suggested for broader and more standardized checkpoint evaluation.
- nanoGPT and build-nanogpt: concise PyTorch reference implementations related to the tutorial's model.
- llm.c: a specialized C/CUDA GPT-2 and GPT-3 training implementation that matches the PyTorch reference numerically while running faster in the demonstrated comparison.
- Lambda Labs: the cloud GPU provider used for the eight-A100 training machine and identified as a sponsor.