KellerJordan/modded-nanogpt

What it solves

It addresses the challenge of training language models as quickly as possible. Specifically, it aims to find the fastest algorithm to train a model on 8 NVIDIA H100 GPUs that can achieve a specific performance target (3.28 cross-entropy loss on the FineWeb validation set).

How it works

The project is a collaborative "speedrun" that iteratively optimizes a training pipeline. It started from a PyTorch trainer based on NanoGPT and llm.c, and has since integrated dozens of architectural and systems optimizations, including:

• Optimizers: Implementation of the Muon and NorMuon optimizers.
• Architecture: Use of Rotary embeddings, QK-Norm, ReLU², and skip connections from embeddings to every block.
• Precision: Leveraging FP8 for the head and MLP forward pass, and BF16 for cross-entropy computation.
• Attention: Integration of Flash Attention 3, long-short sliding window patterns, and YaRN.
• Systems: Optimizing gradient all-reduce/reduce-scatter and overlapping computation with communication.

Who it’s for

It is for AI researchers, engineers, and enthusiasts interested in extreme training efficiency, LLM optimization, and pushing the limits of hardware utilization on H100 GPUs.

Highlights

• Massive Speedup: Reduced training time from 45 minutes (baseline) to under 2.2 minutes for the same target loss.
• Data Efficiency: Reduced the required tokens from 10 billion to under 400 million.
• Collaborative Evolution: A detailed world record history documenting the progression of techniques used to shave off seconds of training time.