$ show karpathy/autoresearch
Autoresearch is a self-contained autonomous ML experimentation framework created by Andrej Karpathy. Its core premise is that an AI agent (e.g., Claude or Codex) modifies a single training script (`train.py`), runs a fixed 5-minute training experiment on a single GPU, evaluates validation bits-per-byte (BPB), and keeps or discards changes based on whether the metric improved — all without human intervention. The human programs the agent via `program.md`, which serves as a lightweight instruction set for the experiment loop. The codebase has two fixed files and one editable file. `prepare.py` handles data downloading (parquet shards from HuggingFace), BPE tokenizer training (via rustbpe), a custom BOS-aligned packed dataloader with best-fit document packing (zero padding), and a fixed BPB evaluation harness. `train.py` contains a GPT model implementation with RoPE, sliding window attention (SSSL pattern), ResFormer value embeddings, RMS normalization, and a hybrid Muon+AdamW optimizer with per-parameter-type learning rate scheduling. The model uses `torch.compile` and runs in bf16 with Flash Attention 3. The `analysis.ipynb` notebook provides post-hoc visualization of experiment results from `results.tsv`. The project is explicitly designed for a single NVIDIA GPU (H100 tested), requires no distributed training infrastructure, and pins all dependencies including PyTorch 2.9.1 with CUDA 12.8. The architecture is deliberately minimal — the entire research loop runs in two Python files with roughly 1,600 lines of non-comment code total.
The LLM identified 3 files worth highlighting for reuse:
The `make_dataloader` function implements a BOS-aligned, best-fit document packing strategy that achieves 100% sequence utilization with zero padding. Every row starts with a BOS token, documents are packed using a greedy best-fit heuristic (largest document that fits remaining space), and when no document fits, the shortest document is cropped to fill exactly. This is a non-trivial, reusable pattern for efficient training data loading that avoids the waste of traditional padding-based approaches. The `evaluate_bpb` function implements a vocab-size-independent bits-per-byte metric that sums per-token cross-entropy in nats and divides by total UTF-8 byte length, excluding special tokens — a clean, reusable evaluation pattern for language model training.
The `MuonAdamW` optimizer implements a hybrid optimization strategy: Muon (Newton-style orthogonalization via polar express) for 2D matrix parameters and AdamW for everything else, with per-parameter-type learning rates, cautious weight decay (only applied when gradient aligns with parameter direction), and NorMuon variance reduction. The fused step functions use `@torch.compile` with CPU scalar tensors to avoid recompilation. The `GPT` model includes ResFormer-style value embeddings with input-dependent gating (alternating layers), RMS normalization with per-layer residual and x0 lambdas, and configurable sliding window attention patterns. The `setup_optimizer` method demonstrates a principled approach to grouping parameters by type and shape with dimension-scaled learning rates — a reusable pattern for any PyTorch training setup.
The training loop implements several production-quality techniques: Python garbage collection is frozen after the first step to avoid ~500ms stalls (a real and non-obvious optimization), loss explosion detection with fast failure, EMA-debiased loss logging, model FLOPS utilization (MFU) tracking against H100 peak bf16 FLOPS, and a time-budget-based LR schedule with configurable warmup and warmdown phases. The loop structure — compile, prefetch first batch, then tight loop with gradient accumulation — is a clean, reusable template for single-GPU training scripts.