README

Modded-NanoGPT

This repository hosts the NanoGPT speedrun, in which we (collaboratively|competitively) search for the fastest algorithm to use 8 NVIDIA H100 GPUs to train a language model that attains 3.28 cross-entropy loss on the FineWeb validation set.

(Note: Besides the main track, there is also an optimization track where we try to minimize steps subject to fixed arch/data/bsz and with unlimited wallclock budget.)

The target (3.28 validation loss on FineWeb) follows Andrej Karpathy's GPT-2 replication in llm.c, which attains that loss after running for 45 minutes.
The speedrun code also descends from llm.c's PyTorch trainer, which itself descends from NanoGPT, hence the name of the repo.
Thanks to the efforts of many contributors, this repo now contains a training algorithm which attains the target performance in:
 Under 40 seconds on 8xH100 (the llm.c GPT-2 replication needed 45 minutes)
 under 330M tokens (the llm.c GPT-2 replication needed 10B)

This improvement in training speed has been brought about by the following techniques:
 Modernized architecture: Rotary embeddings, QK-Norm, and ReLU²
 The Muon optimizer [writeup] [repo]
 Use FP8 for head, and asymmetric rescale and softcap logits
 FP8 for the MLP forward and backward passes and the QKV projections
 Initialization of projections to zero (muP-like)
 Skip connections from embedding to every block as well as from block 3 to 6
 Extra embeddings which are mixed into the values in attention layers (inspired by Zhou et al. 2024)
 Flash Attention 3 with long-short sliding window attention pattern (inspired by Gemma 2) and window size warmup with YaRN
 Align training batch starts with EoS and set a max document length
 Accumulate gradients for 2 steps for embedding and lm_head before updating parameters
 Single activation input for last 3 attention layers
 Polar Express implementation in Muon
 Smear module to enable 1 token look back
 Sparse attention gate
 NorMuon
 Cautious Weight Decay w/ schedule tied to LR
 Exponential decay of residual stream
 Batch size schedule
 Max seq length schedule
 Partial Key Offset
 Multi token prediction
 Untie embed and lm_head at 2/3 of training
 Additional gating on value embeddings and skip connection
 Paired head attention
 Hashed bigram + trigram embedding table (84.6M rows, 65B parameters) w/ sign trick, sharded across GPUs
 MUDD skip connections to residual stream and attention values
 Learnable XSA
 Lightweight Dynamically Composable MHA
 Prefix token prediction auxiliary loss
 ANVIL optimizer: twin-rail (fast + slow) momentum, re-derived orthogonalization maps, cautious decay gated on the slow rail
 Sampled softmax over a shared candidate set for most of training
 Sparse row-wise embedding updates: only touched rows are exchanged and updated (beta1 = 0, one second-moment value per row), every 4 steps
 Depth reduction and mixed-width attention (64-dim QK on most layers, half-width V on two layers)
 Tail weight averaging (EMA) with a raised LR fl