$ show karpathy/nanochat
nanochat is a minimal, hackable LLM training and inference framework focused on reproducing GPT-2 performance quickly on 8×H100 GPUs. It provides an end-to-end pipeline from tokenizer training and dataset preparation through pretraining, supervised fine-tuning, reinforcement learning (GRPO/REINFORCE), and a chat CLI. The core is a highly optimized GPT model with advanced features including a custom MuonAdamW hybrid optimizer, FP8 training with tensorwise dynamic scaling, unified Flash Attention 3/SDPA interface, sliding window attention, value embeddings, and BOS-aligned best-fit packing dataloader. Hyperparameters are automatically derived from scaling laws based on model depth, enabling compute-optimal training. The codebase is research-oriented, with extensive experiment logging, leaderboard tracking, and scaling law analysis notebooks. While primarily an application (CLI-driven training/evaluation framework), it has strong data_research elements including dataset pipelines, evaluation tasks (CORE, MMLU, GSM8K, etc.), and analytical notebooks.
The LLM identified 7 files worth highlighting for reuse:
MuonAdamW optimizer: a sophisticated fused optimizer combining Muon (orthogonalized updates via Polar Express/Newton-Schulz) for matrix parameters and AdamW for embeddings/scalars, with ZeRO-2 style sharding, async communication, momentum/weight-decay schedulers, cautious weight decay, and NorMuon variance reduction. Highly reusable as a standalone optimizer for transformer training.
Float8Linear and conversion utilities: a drop-in replacement for nn.Linear that uses FP8 matrix multiplication with tensorwise dynamic scaling. Includes module filtering (dims divisible by 16, min 128 features), disable_fp8 context manager for evaluation, and meta-device construction pattern. Demonstrates non-obvious FP8 training integration.
Unified Flash Attention 3 / SDPA interface: auto-detects GPU capability and selects the optimal attention kernel (FA3 on Hopper/Ada, SDPA fallback). Provides a clean API with support for sliding window attention patterns. Demonstrates non-obvious hardware-aware kernel selection and fallback handling.
BOS-aligned best-fit packing dataloader: implements bin-packing of token sequences into fixed-length batches with BOS token alignment, minimizing crop waste (~34.6%). Supports distributed training with resume capability via stateful dataloader. Demonstrates non-obvious data loading optimization for transformer training.
Inference engine with KV cache and tool-use state machine: provides streaming generation with support for calculator and Python execution tools. Includes FA3-compatible KV cache with prefill and replication. Reusable pattern for interactive LLM inference.
HubDataset: a minimal HuggingFace-like dataset wrapper using pyarrow, avoiding the `datasets` dependency. Supports sharded parquet loading, shuffling, and slicing. Demonstrates a lightweight, reusable dataset abstraction for large-scale data loading.
Task, TaskMixture, TaskSequence: composable dataset abstractions for evaluation tasks. Each task defines its own dataset, prompt formatting, and metric computation. Supports mixtures and sequences for multi-task evaluation. Reusable pattern for few-shot evaluation frameworks.