Auto Research

Autoresearch - In-Depth Codebase Guide

Created by Andrej Karpathy. An autonomous AI agent research platform that lets AI agents run LLM training experiments unattended — modify code, train for 5 minutes, keep improvements, discard failures, repeat.

---

Table of Contents

1. What Is This Project?
2. Core Philosophy
3. Project Structure
4. Setup & Getting Started
5. The Experiment Loop
6. Data Pipeline
7. Model Architecture
8. Training Pipeline
9. Optimizer: MuonAdamW
10. Evaluation & Metrics
11. Key Hyperparameters
12. Performance Optimizations
13. How the Agent Loop Works
14. Results Tracking & Analysis
15. Glossary

---

1. What Is This Project?

Autoresearch is an autonomous research platform where AI agents (like Claude or GPT) conduct neural network experiments on their own. The idea is simple but powerful:

• You go to sleep
• An AI agent modifies the training code (train.py), trains a small LLM for 5 minutes, checks if the result improved
• If it improved → keep the change. If not → revert
• Repeat ~12 times per hour, ~100 experiments overnight

When you wake up, you have a better model and a log of what the agent tried, what worked, and what didn’t.

Target hardware: Single NVIDIA GPU (tested on H100).

---

2. Core Philosophy

Principle	What It Means
Minimal	Only 5 files. No config systems, no distributed training, no abstractions
Self-contained	Everything in prepare.py (data) and train.py (model + training)
Fixed time budget	Every experiment runs for exactly 5 minutes of wall-clock training
Single metric	val_bpb (validation bits-per-byte) — lower is better
Git-based versioning	Each experiment is a git commit; reverted if it didn’t improve

The human writes a program.md file (instructions for the agent). The agent modifies train.py. That’s the entire workflow.

---

3. Project Structure

autoresearch/
├── prepare.py       # Data download, tokenizer training, dataloader, evaluation
│                    # READ-ONLY — the agent must NOT modify this file
│
├── train.py         # Model architecture, optimizer, hyperparameters, training loop
│                    # THIS is what the AI agent modifies each experiment
│
├── program.md       # Instructions for the AI agent (what to try, constraints)
│                    # Written by the human researcher
│
├── pyproject.toml   # Project metadata and dependencies
├── analysis.ipynb   # Jupyter notebook to visualize results.tsv
├── results.tsv      # Experiment log (git-ignored, created during runs)
└── .python-version  # Python version (3.12)

Data cache (created by prepare.py):

~/.cache/autoresearch/
├── data/            # Downloaded parquet shards from HuggingFace
└── tokenizer/       # Trained BPE tokenizer files
    ├── tokenizer.pkl
    └── token_bytes.pt

---

4. Setup & Getting Started

Prerequisites

• NVIDIA GPU (H100 recommended)
• Python 3.10+
• uv package manager

Step 1: Prepare Data & Tokenizer (one-time, ~2 min)

uv run prepare.py

This does two things:

1. Downloads training data — 10 parquet shards (by default) from karpathy/climbmix-400b-shuffle on HuggingFace
2. Trains a BPE tokenizer — 8192 vocab size, saved to ~/.cache/autoresearch/tokenizer/

Optional flags:

• --num-shards N — download more shards (default 10, max 6542)
• --download-workers N — parallel download threads (default 8)

Step 2: Run a Single Training Experiment

uv run train.py

This trains for exactly 5 minutes and prints metrics like:

val_bpb:          1.234
training_seconds: 300.1
peak_vram_mb:     12345
mfu_percent:      45.6

Step 3: Let an Agent Run Experiments

Point an AI agent at program.md and let it loop autonomously (see Section 13).

---

5. The Experiment Loop

Here’s the full cycle that the AI agent runs, repeatedly:

┌─────────────────────────────────────────────────┐
│  1. Read program.md for guidance                │
│  2. Think of a hypothesis (e.g., "try depth=10")│
│  3. Modify train.py                             │
│  4. git commit -m "try depth=10"                │
│  5. Run: uv run train.py > run.log 2>&1         │
│  6. Extract val_bpb from output                 │
│  7. Did it improve?                             │
│     ├─ YES → Record "keep" in results.tsv       │
│     └─ NO  → git reset (revert), record "discard"│
│  8. Go to step 1                                │
└─────────────────────────────────────────────────┘

Each cycle takes ~5-7 minutes (5 min training + startup/compilation overhead).

---

6. Data Pipeline

6.1 Data Source

• Dataset: karpathy/climbmix-400b-shuffle on HuggingFace
• Format: Parquet files with a "text" column
• Total shards: 6,542 (by default only 10 are downloaded)
• Validation shard: shard_06542 (always pinned, never used for training)

6.2 Tokenizer

• Type: BPE (Byte Pair Encoding) via rustbpe library
• Vocab size: 8,192 tokens
• Special tokens: 4 reserved tokens (<|reserved_0|> through <|reserved_3|>)
• BOS token: <|reserved_0|> — prepended to every document
• Split pattern: GPT-4 style regex that handles contractions, numbers, whitespace

# The tokenizer split pattern (same as GPT-4)
SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ..."""

6.3 Dataloader (make_dataloader in prepare.py)

The dataloader is carefully designed for 100% token utilization:

Document A (300 tokens)  Document B (500 tokens)  Document C (1200 tokens)
    ↓                        ↓                        ↓
┌──────────────────────────────────────────────────────────┐
│ [BOS] Doc_A | [BOS] Doc_B | [BOS] Doc_C...              │  ← Row 1 (2048 tokens)
│ ...Doc_C continued | [BOS] Doc_D | [BOS] Doc_E...       │  ← Row 2 (2048 tokens)
└──────────────────────────────────────────────────────────┘

Key properties:

• BOS-aligned: Every document starts with a BOS token
• Best-fit packing: Documents are packed greedily to fill rows exactly
• No padding: If a row isn’t full, the shortest remaining document is cropped to fit
• Infinite iterator: Loops through data endlessly, yielding (inputs, targets, epoch)
• Async GPU transfer: Data is pre-loaded to GPU with non_blocking=True

6.4 Constants

MAX_SEQ_LEN = 2048       # Context window (sequence length)
EVAL_TOKENS = 40 * 524288  # ~20.9M tokens for validation

---

7. Model Architecture

The model is a GPT-style transformer with several modern improvements. Defined entirely in train.py.

7.1 High-Level Structure

Input Token IDs
       ↓
┌─────────────────┐
│  Token Embedding │  (wte: vocab_size → n_embd)
└────────┬────────┘
         ↓
┌─────────────────┐
│  Transformer     │  × n_layer (e.g., 12 layers)
│  Block           │
│  ┌─────────────┐ │
│  │ RMSNorm     │ │
│  │ Attention   │ │  ← Flash Attention 3 + RoPE + Sliding Window
│  │ + Residual  │ │
│  │ + Skip (x0) │ │
│  ├─────────────┤ │
│  │ RMSNorm     │ │
│  │ MLP         │ │  ← ReLU² activation
│  │ + Residual  │ │
│  │ + Skip (x0) │ │
│  └─────────────┘ │
└────────┬────────┘
         ↓
┌─────────────────┐
│  RMSNorm        │
│  LM Head        │  (n_embd → vocab_size)
│  Softcap        │  tanh(logits/15) * 15
└────────┬────────┘
         ↓
   Token Probabilities

7.2 Attention Mechanism

# Multi-head self-attention with:
# - Rotary Position Embeddings (RoPE)
# - Flash Attention 3 (hardware-optimized)
# - Sliding Window pattern: "SSSL"
 
# Window pattern explanation:
#   S = short window (half of MAX_SEQ_LEN = 1024 tokens)
#   L = long window  (full MAX_SEQ_LEN = 2048 tokens)
#   Pattern repeats: layers 0,1,2 use short, layer 3 uses long, etc.
#   Last layer ALWAYS uses long window regardless of pattern

Why sliding window? It saves memory and compute — most information is local. Only every 4th layer needs to “see” the full context.

7.3 MLP Block

class MLP:
    fc:   Linear(n_embd → 4 * n_embd)   # expand
    act:  ReLU().square()                 # ReLU² activation (unusual but effective)
    proj: Linear(4 * n_embd → n_embd)    # project back

ReLU² (ReLU-squared): max(0, x)². This creates sparser activations than GELU/SiLU, which can improve training efficiency.

7.4 Residual Connections (with Skip and Scaling)

Each block has two types of residual connections:

# Standard residual (scaled):
x = resid_lambdas[layer] * x + block_output
 
# Skip residual to original input:
x = x + x0_lambdas[layer] * x0   # x0 is the output of the embedding layer

• resid_lambdas (initialized to 1.0): scales the residual stream
• x0_lambdas (initialized to 0.1): adds a small fraction of the original embeddings back

This is inspired by ResFormer — it helps with gradient flow in deeper models.

7.5 Value Embeddings (ResFormer)

Alternating layers include learnable value embeddings:

# In even-numbered layers:
gate = 2 * sigmoid(linear(x[:, :gate_channels]))  # input-dependent gate
v = v + gate * value_embedding                      # mix in value embedding

• Gate weights initialized to 0 → sigmoid(0) * 2 = 1.0 → neutral start
• The model learns when and how much to use these extra embeddings

7.6 Softcap (Output Clamping)

logits = tanh(logits / 15.0) * 15.0

Prevents logits from exploding beyond [-15, 15]. This stabilizes training, especially with aggressive learning rates.

7.7 Model Size Derivation

The model size is controlled by a single knob — DEPTH:

DEPTH = 8                           # Main size control
ASPECT_RATIO = 64                   # Fixed ratio
HEAD_DIM = 128                      # Target head dimension
 
base_dim = DEPTH * ASPECT_RATIO     # 8 * 64 = 512
model_dim = round_up_to(base_dim, HEAD_DIM)  # 512 → 512 (already aligned)
n_heads = model_dim // HEAD_DIM     # 512 / 128 = 4
n_layers = DEPTH * (ASPECT_RATIO // HEAD_DIM)  # 8 * 0.5 = ... (see actual formula)

Increasing DEPTH makes the model both wider and deeper proportionally.

---

8. Training Pipeline

8.1 Initialization Sequence

1. Set random seeds (42)
2. Load tokenizer from ~/.cache/autoresearch/tokenizer/
3. Derive model config from DEPTH hyperparameter
4. Create model on "meta" device (no memory allocation)
5. Call init_weights() for careful initialization
6. Move model to GPU
7. Compile model with torch.compile(dynamic=False)
8. Set up MuonAdamW optimizer with parameter groups
9. Create infinite dataloader

8.2 Weight Initialization (Important!)

Different components get different initialization strategies:

Component	Initialization	Why
Token embeddings (wte)	N(0, 1.0)	Standard embedding init
Q, K, V projections	Uniform(-s, s), s = √3/√n_embd	Kaiming-style
Output projections (attn, MLP)	Zeros	Residual connection starts as identity
resid_lambdas	1.0	Full residual initially
x0_lambdas	0.1	Small skip connection initially
Value embedding gates	Zeros	Neutral gate (sigmoid(0)*2 = 1.0)

Why zero-init output projections? At the start of training, each transformer block is an identity function (output = input). The model gradually learns to add useful transformations. This prevents the “signal collapse” problem in deep networks.

8.3 Training Loop

while (step <= 10) or (training_time < 300 seconds):
    # 1. Learning rate schedule
    lr_multiplier = warmup_then_plateau_then_cooldown(step)
 
    # 2. Forward pass (bf16 mixed precision)
    with torch.amp.autocast("cuda", torch.bfloat16):
        logits = model(inputs)
        loss = cross_entropy(logits, targets)
 
    # 3. Backward pass (with gradient accumulation)
    loss.backward()
    if (step % grad_accum_steps == 0):
        optimizer.step()
        optimizer.zero_grad()
 
    # 4. Fast-fail check
    if loss > 100 or isnan(loss):
        print("FAILED: loss exploded")
        exit(1)
 
    # 5. Memory management
    if step == 1: gc.disable()          # GC pauses are ~500ms
    if step % 5000 == 0: gc.collect()   # Periodic cleanup

8.4 Learning Rate Schedule

LR
 ↑
 │         ┌────────────────────┐
 │        /│    Plateau         │\
 │       / │                    │ \
 │      /  │                    │  \
 │     /   │                    │   \
 │    /    │                    │    \
 │   /     │                    │     \
 │──/──────┴────────────────────┴──────\──→ Steps
 │ Warmup                       Warmdown
 │ (WARMUP_RATIO)               (WARMDOWN_RATIO = 0.5)

Default: no warmup, plateau for first half, linear decay to 0 in second half.

---

9. Optimizer: MuonAdamW

This is a hybrid optimizer that uses different algorithms for different parameter types:

9.1 Parameter Groups

Group	Parameters	Optimizer	Learning Rate
Matrix params	2D weight matrices (attention, MLP)	Muon	0.04 (scaled by √aspect_ratio)
Embeddings	Token embeddings (wte), value embeddings	AdamW	0.6
Unembedding	LM head (output projection)	AdamW	0.004
Scalars	resid_lambdas, x0_lambdas	AdamW	0.5

9.2 What Is Muon?

Muon is a novel optimizer specifically designed for matrix-valued parameters:

Standard gradient descent:  W = W - lr * gradient
Muon:                       W = W - lr * orthogonalize(momentum)

Key steps:

1. Compute gradient
2. Update exponential moving average (momentum)
3. Orthogonalize the momentum matrix using Newton-Schulz iteration (5 steps)
4. Apply NorMuon variance reduction
5. Apply update with cautious weight decay

Why orthogonalize? It keeps weight matrices well-conditioned, preventing training instabilities. Think of it as constantly “straightening out” the weight matrices.

9.3 Momentum Schedule

Muon’s momentum starts low and ramps up:

Steps 0-300:   momentum = 0.85 → 0.95 (linear ramp)
Steps 300+:    momentum = 0.95 (constant)

9.4 Weight Decay Schedule

Weight decay linearly decays to 0 by the end of training:

Start: WEIGHT_DECAY = 0.2
End:   0.0

---

10. Evaluation & Metrics

10.1 The Key Metric: val_bpb (Bits Per Byte)

BPB measures how efficiently the model compresses text, independent of vocabulary size:

BPB = total_cross_entropy_nats / (ln(2) * total_utf8_bytes)

Why not just loss? Different tokenizers produce different token counts for the same text. BPB normalizes by the actual byte count, so you can fairly compare models with different vocab sizes.

Lower is better. A BPB of 1.0 means the model uses 1 bit per byte of text on average.

10.2 Evaluation Process

def evaluate_bpb(model, tokenizer):
    total_nats = 0
    total_bytes = 0
    for batch in validation_data:      # ~20.9M tokens
        logits = model(batch)
        per_token_loss = cross_entropy(logits, targets, reduction='none')
        total_nats += sum(per_token_loss)
        total_bytes += sum(utf8_byte_counts[tokens])  # skip special tokens
    return total_nats / (log(2) * total_bytes)

Special tokens (byte count = 0) are excluded from the BPB calculation.

10.3 Output Metrics

After each training run, train.py prints:

val_bpb:          1.234    ← THE metric to optimize
training_seconds: 300.1    ← always ~300s (5 minutes)
total_seconds:    320.5    ← includes startup/compilation
peak_vram_mb:     12345    ← GPU memory used
mfu_percent:      45.6     ← Model FLOPs Utilization (% of H100 peak)
total_tokens_M:   150.2    ← millions of tokens processed
num_steps:        286      ← optimizer steps taken
num_params_M:     45.3     ← millions of parameters
depth:            8        ← model depth

---

11. Key Hyperparameters

All hyperparameters live at the top of train.py and are meant to be modified by the agent:

Model Architecture

DEPTH = 8              # THE main knob — controls both width and depth
ASPECT_RATIO = 64      # model_dim = DEPTH * ASPECT_RATIO (rounded to HEAD_DIM)
HEAD_DIM = 128         # dimension per attention head
WINDOW_PATTERN = "SSSL"  # S=short window (1024), L=long window (2048)

Batch Size

DEVICE_BATCH_SIZE = 128      # sequences per forward pass
TOTAL_BATCH_SIZE = 2**19     # ~524K tokens per optimizer step
# grad_accum_steps = TOTAL_BATCH_SIZE / (DEVICE_BATCH_SIZE * MAX_SEQ_LEN)

Learning Rates

EMBEDDING_LR = 0.6      # Token embeddings — high LR
UNEMBEDDING_LR = 0.004  # Output projection — low LR
MATRIX_LR = 0.04        # Weight matrices (Muon) — moderate LR
SCALAR_LR = 0.5         # Per-layer scalars — high LR

Optimization

WEIGHT_DECAY = 0.2           # Cautious decay (Muon only), decays to 0
ADAM_BETAS = (0.8, 0.95)     # AdamW momentum parameters
WARMUP_RATIO = 0.0           # Fraction of time for LR warmup (none by default)
WARMDOWN_RATIO = 0.5         # Fraction of time for LR cooldown
FINAL_LR_FRAC = 0.0          # Final LR as fraction of peak (0 = decay to zero)

---

12. Performance Optimizations

Flash Attention 3

Uses hardware-optimized attention kernels via the kernels library. Automatically detects GPU architecture (Hopper/H100 = compute capability 9.0).

torch.compile

The entire model is compiled with torch.compile(model, dynamic=False) for maximum kernel fusion and optimization. First step is slow (compilation), subsequent steps are fast.

Mixed Precision (bfloat16)

All training runs in bfloat16 via torch.amp.autocast. Float32 is used selectively (e.g., softcap logits computation).

Garbage Collection Management

# Python's GC causes ~500ms stalls during training
gc.disable()                    # Disabled after first step
if step % 5000 == 0: gc.collect()  # Manual collection periodically

Memory Optimizations

• Pre-allocated GPU buffers for data loading
• non_blocking=True for async CPU→GPU transfers
• PYTORCH_ALLOC_CONF="expandable_segments:True" for better memory allocation
• Embedding weights cast to bf16 after initialization

---

13. How the Agent Loop Works

The Agent’s Perspective

The AI agent (Claude, GPT, etc.) operates as follows:

1. Read program.md — contains human-written instructions like:

   • “Try increasing DEPTH”
   • “Experiment with different activation functions”
   • “Don’t change prepare.py”

2. Decide on an experiment — based on past results and program.md guidance

3. Modify train.py — change hyperparameters, architecture, or training logic

4. Commit & Run:

   git add train.py
   git commit -m "try depth=10"
   uv run train.py > run.log 2>&1

5. Parse Output — extract val_bpb and peak_vram_mb from run.log

6. Record Results:

   # Append to results.tsv:
   abc1234  1.234  12.3  keep     try depth=10
   

7. Keep or Revert:

   • If val_bpb improved → keep the commit
   • If not → git reset --hard HEAD~1 (revert)
   • If crashed → record as “crash”, revert

Constraints for the Agent

• NEVER modify prepare.py — evaluation must remain consistent
• Must stay within VRAM limits — soft constraint on peak memory
• 5-minute budget is fixed — can’t be changed
• Single metric optimization — only val_bpb matters

---

14. Results Tracking & Analysis

results.tsv Format

commit    val_bpb   memory_gb   status    description
abc1234   1.234     12.3        keep      try depth=10
def5678   1.256     14.1        discard   try depth=12 (regression)
ghi9012   0.000     0.0         crash     try broken activation

Columns:

• commit: 7-character git hash
• val_bpb: validation bits-per-byte (lower = better)
• memory_gb: peak VRAM in GB
• status: keep (improved), discard (regressed), crash (error)
• description: short explanation of the experiment

analysis.ipynb

The Jupyter notebook loads results.tsv and plots:

• BPB frontier over time (best BPB at each point)
• Per-experiment improvements
• Keep/discard/crash rates
• Summary statistics

---

15. Glossary

Term	Meaning
BPB	Bits Per Byte — the primary evaluation metric. Lower = better compression = better model
BOS	Beginning Of Sequence token — prepended to every document
BPE	Byte Pair Encoding — tokenization algorithm that merges frequent byte pairs
Flash Attention	Hardware-optimized attention kernel (much faster than naive attention)
GQA	Grouped Query Attention — KV heads shared across query heads
MFU	Model FLOPs Utilization — % of theoretical GPU peak being used
Muon	Novel optimizer that orthogonalizes momentum for matrix parameters
RMSNorm	Root Mean Square Normalization — simpler alternative to LayerNorm
RoPE	Rotary Position Embeddings — encodes position info into Q/K vectors
Softcap	tanh(x/c)*c — clamps logits to prevent explosion
Sliding Window	Attention that only looks at nearby tokens (saves memory)
ResFormer	Technique adding value embeddings and skip connections to input
Newton-Schulz	Iterative algorithm to orthogonalize matrices (used in Muon)
Warmdown	Gradual LR decrease at end of training

---

Quick Reference: File Responsibilities

prepare.py  →  "I handle data and evaluation. Don't touch me."
train.py    →  "I am the experiment. Modify me."
program.md  →  "I tell the agent what to try. Human writes me."
results.tsv →  "I record what happened. Auto-generated."

---

Dependencies

torch==2.9.1       Core deep learning framework (CUDA 12.8)
kernels>=0.11.7    Flash Attention 3 implementation
rustbpe>=0.1.0     Fast BPE tokenizer (Rust-based)
tiktoken>=0.11.0   OpenAI's tokenizer encoding library
pyarrow>=21.0.0    Parquet file reading
requests>=2.32.0   HTTP downloads for data
numpy>=2.2.6       Numerical operations
pandas>=2.3.3      DataFrame ops for analysis
matplotlib>=3.10.8 Plotting in analysis notebook