Jesse A. Green / notebook

autoresearch Administrator & Developer Guide

Generated April 2026 from karpathy/autoresearch at 228791f

AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it.

This guide is the architectural and operational companion to the User Guide. It is written for:

  • Researchers forking autoresearch to extend it (new metric, new optimizer, multi-GPU, alternative datasets).
  • Operators running autonomous research at scale — multiple agents, multiple GPUs, shared caches, observability.
  • Reviewers who want to know exactly what the 80M-parameter GPT baseline is doing, byte-for-byte, before trusting agent-driven val_bpb deltas.

autoresearch is small (10 files, ~41 KB of source), so this guide is correspondingly tight. There is no API surface to document, no plugin system, no service layer — every meaningful design decision lives in two Python files.


1. Architecture Overview

autoresearch is a deliberately minimal, single-GPU LLM pretraining sandbox. It is a stripped-down derivative of nanochat, reduced to the smallest configuration that still produces a fair, reproducible val_bpb (bits-per-byte) score on a fixed validation shard, in a fixed wall-clock budget. The whole repo is designed around the assumption that an AI coding agent, prompted with program.md, will edit train.py autonomously, run experiments, and keep what works.

1.1 ASCII component diagram

 ┌─────────────────────────────────────────────────────────────┐
 │                  AUTONOMOUS RESEARCHER                      │
 │   (Claude Code / Codex / etc., prompted with program.md)    │
 └────────────────┬────────────────────────────────────────────┘
                  │ git commit · uv run · grep · git reset
                  ▼
 ┌─────────────────────────────────────────────────────────────┐
 │                   REPOSITORY (branch)                       │
 │                                                             │
 │  prepare.py  (read-only)    train.py  (AGENT-EDITED)        │
 │  program.md  (human-edited) pyproject.toml  (locked)        │
 │  results.tsv (untracked, append-only per-run)               │
 │  analysis.ipynb (post-hoc)  progress.png  (cosmetic)        │
 └────────────────┬────────────────────────────────────────────┘
                  │ uv run
                  ▼
 ┌─────────────────────────────────────────────────────────────┐
 │             SINGLE-GPU TRAINING PROCESS                     │
 │                                                             │
 │  Tokenizer (tiktoken + rustbpe-trained merges)              │
 │  GPT (DEPTH=8, n_embd=768, GQA, sliding window attention)   │
 │  MuonAdamW optimizer (Muon for matrices, AdamW elsewhere)   │
 │  Flash Attention 3 kernel (loaded via `kernels` package)    │
 │                                                             │
 │  while total_training_time < TIME_BUDGET (300 s):           │
 │      forward → backward → optimizer.step() → schedules      │
 │  evaluate_bpb on pinned validation shard (shard_06542)      │
 └────────────────┬────────────────────────────────────────────┘
                  │ stdout
                  ▼
 ┌─────────────────────────────────────────────────────────────┐
 │  STDOUT → run.log → grep ^val_bpb → results.tsv row         │
 └─────────────────────────────────────────────────────────────┘

1.2 Component inventory

Component Path Role
Data + tokenizer prep prepare.py Fixed constants, shard download, BPE training, dataloader, evaluate_bpb, Tokenizer runtime class
Training script train.py GPT model, MuonAdamW, hyperparameters, training loop, final eval
Agent prompt program.md Programs the autonomous research loop
Lockfile uv.lock Pinned dependency graph (uv sync source of truth)
Project metadata pyproject.toml Python version, deps, pytorch-cu128 index
Pinned Python version .python-version 3.10
Analysis notebook analysis.ipynb Plots results.tsv over commits
README image progress.png Cosmetic loss-curve teaser
Repo README README.md Human-readable narrative
Git ignore .gitignore Standard Python exclusions

1.3 Technology stack

Layer Technology Notes
Language Python 3.10+ Pinned in .python-version
Framework PyTorch 2.9.1 (cu128) Pinned in pyproject.toml
Attention Flash Attention 3 Loaded at runtime via kernels from HF Hub
Tokenizer training rustbpe Rust-backed BPE trainer
Tokenizer runtime tiktoken OpenAI's BPE encoder, instantiated from rustbpe merges
Data format Apache Parquet via pyarrow One row group at a time
Project manager uv Dependency resolution, virtualenv, run wrapper
Optimizer MuonAdamW (in-tree) Muon for 2-D matrices, AdamW elsewhere

1.4 Process flow for one experiment

  1. prepare.py (one-time) — populate ~/.cache/autoresearch/{data,tokenizer}.
  2. Agent git checkouts a working branch, edits train.py, git commits.
  3. uv run train.py > run.log 2>&1:
    • Build Tokenizer from cache.
    • Build GPT(config) on the meta device, then to_empty(device="cuda").
    • init_weights() — custom initialization (uniform on attention/MLP matrices, zero on c_proj, embeddings cast to bf16).
    • setup_optimizer() returns a MuonAdamW over per-shape parameter groups.
    • torch.compile(model, dynamic=False).
    • Prefetch first batch, then loop until total_training_time >= TIME_BUDGET.
    • On exit, evaluate_bpb runs over EVAL_TOKENS = 40 * 524288 validation tokens.
    • Print summary block.
  4. Agent reads ^val_bpb and ^peak_vram_mb via grep, writes a row to results.tsv, decides keep vs revert.

Note

The first 10 training steps are excluded from total_training_time (if step > 10: total_training_time += dt). This is so that torch.compile graph capture and autotune do not eat into the 300-second steady-state budget. Without this, slower compile passes would unfairly punish experiments that compile slightly differently.


2. System Requirements

2.1 Hardware

Resource Minimum Recommended Notes
GPU NVIDIA Ampere (sm_80) NVIDIA Hopper H100/H200 Code hardcodes torch.device("cuda")
VRAM ~16 GB (with DEPTH=2, DEVICE_BATCH_SIZE=8 reductions) 80 GB (default config peaks at ~44 GB) See User Guide §5.4 for shrinking
CPU Any modern x86_64 — multiprocessing.Pool is used only during shard download
RAM 16 GB 32 GB+ Parquet row-groups fit comfortably
Disk 5 GB free 20 GB+ for full dataset Each shard is ~30 MB
Network Outbound HTTPS 100 Mbps+ First run downloads ~330 MB; full dataset is hundreds of GB

2.2 Software

Component Version Pinned where
Python >= 3.10 .python-version and pyproject.toml
uv any recent Not pinned
NVIDIA driver CUDA 12.8-compatible Required by pytorch-cu128 wheel
torch 2.9.1 pyproject.toml, ==
kernels >= 0.11.7 pyproject.toml, >=
pyarrow >= 21.0.0 pyproject.toml
tiktoken >= 0.11.0 pyproject.toml
rustbpe >= 0.1.0 pyproject.toml
requests >= 2.32.0 pyproject.toml
numpy >= 2.2.6 pyproject.toml
pandas >= 2.3.3 pyproject.toml (notebook only)
matplotlib >= 3.10.8 pyproject.toml (notebook only)

2.3 Network and external services

Endpoint Purpose
huggingface.co/datasets/karpathy/climbmix-400b-shuffle Training and validation Parquet shards
huggingface.co (kernels API) Flash Attention 3 wheel resolution (varunneal/flash-attention-3 on Hopper, kernels-community/flash-attn3 elsewhere)
download.pytorch.org/whl/cu128 PyTorch wheels via pytorch-cu128 index

Important

No accounts, API keys, or paid services are required to run the upstream repo. The climbmix-400b-shuffle dataset is public on the Hugging Face Hub. The autonomous loop only requires network access during the one-time prep — once ~/.cache/autoresearch/ is populated, training is fully offline.


3. Configuration Reference

There is no config.yaml, no .env, no environment-variable surface beyond two assignments in train.py itself (set inside the script before import torch.cuda runs).

3.1 Environment variables set by train.py

Variable Value Why
PYTORCH_ALLOC_CONF expandable_segments:True Reduces fragmentation in the CUDA caching allocator on long runs
HF_HUB_DISABLE_PROGRESS_BARS 1 Keeps kernels quiet during FA3 kernel resolution

These are set via os.environ[...] at the top of train.py, before the relevant libraries are imported.

3.2 CLI flags (only prepare.py)

Flag Type Default Description
--num-shards int 10 Training shards to download. -1 = all 6,542. The pinned validation shard is always added.
--download-workers int 8 Parallel download workers (multiprocessing.Pool).

3.3 prepare.py fixed constants (immutable per program.md)

Constant Value Role
MAX_SEQ_LEN 2048 Context length used by training and evaluate_bpb
TIME_BUDGET 300 Seconds of steady-state training
EVAL_TOKENS 40 * 524288 (20,971,520) Tokens scored by the final BPB evaluation
CACHE_DIR ~/.cache/autoresearch Root for data/ and tokenizer/ subdirs
BASE_URL https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle/resolve/main Dataset source
MAX_SHARD 6542 Index of the last shard in the dataset
VAL_SHARD 6542 (= MAX_SHARD) Pinned validation shard index
VAL_FILENAME shard_06542.parquet Computed from VAL_SHARD
VOCAB_SIZE 8192 Includes 4 reserved special tokens
SPLIT_PATTERN GPT-4-style regex with \p{N}{1,2} (max 2-digit number runs) BPE pre-tokenization regex
SPECIAL_TOKENS ["<|reserved_0|>", ..., "<|reserved_3|>"] 4 reserved specials
BOS_TOKEN "<|reserved_0|>" Beginning-of-sequence token

Important

Changing any value above invalidates leaderboard comparisons against anyone else running the same code. program.md declares these read-only because the agent's optimizer (val_bpb) is only meaningful if the metric definition is fixed.

3.4 train.py hyperparameter constants (intended to be edited)

Constant Default Tunable by agent
ASPECT_RATIO 64 yes (sets model_dim = DEPTH * ASPECT_RATIO)
HEAD_DIM 128 yes (target per-head dim; n_head = model_dim // HEAD_DIM)
WINDOW_PATTERN "SSSL" yes (string of S/L; last layer always overridden to L)
TOTAL_BATCH_SIZE 2**19 (524,288) yes (must remain a multiple of DEVICE_BATCH_SIZE * MAX_SEQ_LEN)
EMBEDDING_LR 0.6 yes
UNEMBEDDING_LR 0.004 yes
MATRIX_LR 0.04 yes
SCALAR_LR 0.5 yes
WEIGHT_DECAY 0.2 yes (decays linearly to 0 over training)
ADAM_BETAS (0.8, 0.95) yes
WARMUP_RATIO 0.0 yes
WARMDOWN_RATIO 0.5 yes
FINAL_LR_FRAC 0.0 yes
DEPTH 8 yes
DEVICE_BATCH_SIZE 128 yes (lower on OOM)

The training script has fixed seeds:

torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.set_float32_matmul_precision("high")

Note

Hyperparameters are tuned around n_embd = 768 (the default DEPTH=8, ASPECT_RATIO=64, HEAD_DIM=128). When the model dimension changes, setup_optimizer rescales AdamW LRs by (model_dim / 768) ** -0.5, and _step_muon adds a max(1.0, rows / cols) ** 0.5 factor per parameter shape. This is automatic — agents do not need to recompute LRs by hand.

3.5 pyproject.toml excerpt (dependency surface)

uv.lock is the source of truth. The relevant pieces of pyproject.toml:

  • Python: >= 3.10
  • Custom index: pytorch-cu128 (pulls cu128 PyTorch wheels)
  • Direct dependencies: the table in §2.2

Adding new packages requires editing pyproject.toml and re-running uv sync. The agent's program.md explicitly forbids new dependencies, because reproducible runs mean reproducible environments.


4. Deployment & Scaling

autoresearch deploys as a single Python process per GPU. There is no daemon, no service manager, no API server. "Deployment" is git clone + uv sync + uv run prepare.py once, then arbitrarily many uv run train.py invocations.

4.1 Single-host single-GPU (the default)

The intended unit of work. One agent, one GPU, one branch (autoresearch/<tag>). Capacity is bounded by the 5-minute training budget — about 12 experiments per hour.

4.2 Single-host multi-GPU

Not natively supported (the script hardcodes torch.device("cuda"), which is cuda:0). To use additional GPUs, run independent processes — one per GPU — each in its own checkout, each pinned to a different device:

# Terminal 1
cd /workspace/autoresearch-gpu0
CUDA_VISIBLE_DEVICES=0 uv run train.py > run.log 2>&1

# Terminal 2
cd /workspace/autoresearch-gpu1
CUDA_VISIBLE_DEVICES=1 uv run train.py > run.log 2>&1

Use a separate autoresearch/<tag>-gpu0, autoresearch/<tag>-gpu1, … branch per agent so their commits and results.tsv files do not collide.

Tip

Use git worktrees to share a single object store across N agents while giving each its own working tree:

git worktree add ../autoresearch-gpu1 autoresearch/mar5-gpu1

The shared object store cuts disk usage compared to N full clones, and git fetch runs once.

4.3 Shared dataset cache

~/.cache/autoresearch/ is per-user, not per-checkout. Multiple checkouts on the same host share data and tokenizer artifacts automatically. Across hosts, the simplest production pattern is:

  1. Run prepare.py once on a primary host.
  2. rsync the cache directory to peers.
  3. Each peer skips the download (idempotency check sees the files exist).

Total cache size at default --num-shards 10: roughly 330 MB of data plus a few KB of tokenizer.

4.4 Observability

There is none, by design. Each train.py invocation prints a single carriage-returning log line per step plus a final summary block. To monitor a fleet, use any standard log shipper to tail run.logs. Useful greps:

grep -E "^(val_bpb|training_seconds|peak_vram_mb|mfu_percent|num_params_M):" run.log

The agent does effectively the same thing inside the loop.

Note

There is no W&B, no TensorBoard, no Prometheus exporter. The whole point is to reduce the metric surface to a single floating-point number (val_bpb) that the agent can act on. Keep observability out-of-process.

4.5 Capacity planning

Parameter Effect
TIME_BUDGET = 300 s Per-experiment wall-clock floor
Startup + compile ~25 s on H100
evaluate_bpb final pass ~10 s on H100
Per-experiment total ~5 min 30 s
Throughput ~11–12 experiments/hour/GPU
8-hour overnight ~90–100 experiments/GPU

Tip

If you parallelize across N GPUs with N agents and N branches, expect superlinear iteration speed only if your prompts diverge. N agents on identical program.mds explore the same neighbourhood. Different program.md strategies per GPU are how you actually parallelize the search.


5. Operations

5.1 First-time setup checklist

  • [x] Confirm nvidia-smi reports a Hopper or Ampere GPU.
  • [x] uv --version returns a recent build.
  • [x] git clone https://github.com/karpathy/autoresearch.git.
  • [x] cd autoresearch && uv sync.
  • [x] uv run python -c "import torch; print(torch.cuda.is_available())" returns True.
  • [x] uv run prepare.py completes with Done! Ready to train.
  • [x] uv run train.py finishes a 5-minute run and prints val_bpb:.
  • [ ] Configure your coding agent (Claude Code, Codex) with permissions disabled.
  • [ ] Open program.md, review for project-specific edits.
  • [ ] Pick a run tag and prompt the agent.

5.2 Daily operations

There are no daily operations. Once the agent loop is running, you intervene only to:

  • Inspect results.tsv for progress.
  • Kill the process when you want to stop (no graceful-stop signal handler).
  • Cherry-pick a particularly interesting commit out of autoresearch/<tag> if you want to share it.

5.3 Backups

There is nothing meaningful to back up:

  • Code: lives in git on GitHub.
  • Data: rebuildable via uv run prepare.py (~5 minutes on a fast link).
  • Tokenizer: rebuildable in ~40 seconds.
  • results.tsv: intentionally untracked. Copy out before deleting the working tree if you care about the log.
  • Model weights: never persisted. Each experiment trains from scratch.

Warning

The agent uses git reset on the working branch when an experiment fails. Anything uncommitted in the working tree at that point is destroyed. Treat the agent's branch as ephemeral; never share a working tree with manual edits-in-progress.

5.4 Health checks

A run is "healthy" when:

Indicator Healthy value
mfu_percent (final summary) ≥ 35% on H100, ≥ 25% on A100
training_seconds within 1 s of TIME_BUDGET (= 300 s)
val_bpb numeric, > 0, not nan
peak_vram_mb < device VRAM (no OOM)
Exit code 0

If mfu_percent collapses by 50% between two runs that did not change FA-related code, suspect kernel-loading regression — re-run uv sync --reinstall.

5.5 Cache management

# Clear data + tokenizer (forces full re-prep)
rm -rf ~/.cache/autoresearch

# Clear only tokenizer (e.g. after changing VOCAB_SIZE manually)
rm -rf ~/.cache/autoresearch/tokenizer

# Free disk by deleting unused training shards (keeps val shard)
ls ~/.cache/autoresearch/data/shard_*.parquet | head -n -2 | xargs rm -f

Tip

When a coding agent is going to iterate overnight, run uv run train.py once manually first. That fills the torch.compile cache (~/.cache/torch/inductor/) and shaves seconds off every subsequent run. Cold compiles after a pip reinstall can otherwise add ~20 s to each experiment.

5.6 Common runbooks

Reset to a known-good baseline:

git checkout master
git pull
git checkout -b autoresearch/$(date +%Y%m%d)

Recover from agent confusion (results.tsv corrupted, branch in odd state):

git stash drop --all
git reset --hard HEAD
cp results.tsv /tmp/results-backup.tsv  # if salvageable
git checkout master

Redo a single experiment outside the loop:

git checkout <commit-hash>
uv run train.py > /tmp/replay.log 2>&1
grep "^val_bpb:" /tmp/replay.log

Note

Replays will not exactly match the original val_bpb if the torch.compile cache, FA3 kernel choice, or driver version changed since the original run. They will be close — within the noise of bf16 numerics — but not bit-identical.


6. Performance

6.1 Reference numbers (H100, default config)

Metric Value
num_params_M ≈ 80.2
mfu_percent ≈ 39.8%
tok/sec ≈ 500K
total_tokens_M ≈ 500
num_steps ≈ 950
peak_vram_mb ≈ 45,060
val_bpb (baseline) ≈ 1.0
Wall-clock total ≈ 326 s (300 training + ~25 startup + ~1 eval)

H100 BF16 peak FLOPs are hardcoded at 989.5e12 for the MFU calculation:

H100_BF16_PEAK_FLOPS = 989.5e12
mfu = 100 * num_flops_per_token * TOTAL_BATCH_SIZE / dt / H100_BF16_PEAK_FLOPS

6.2 What dominates wall-clock time

  • Forward + backward (≈97%): dominated by FA3 attention and matmuls inside MLP (c_fc is 4 * n_embd, c_proj projects back).
  • Optimizer step (≈2%): MuonAdamW.step() runs Muon for matrices and AdamW for embeddings/scalars.
  • Dataloader (≈<1%): make_dataloader packs documents into rows with best-fit; pinned-CPU staging plus a non-blocking cuda copy hide the cost.

6.3 Techniques in use

Technique Where Effect
Flash Attention 3 train.py via kernels.get_kernel(...) Memory-efficient, hardware-accelerated attention
torch.compile(dynamic=False, fullgraph=True) Both fused steps and the model Single fused graph; no per-step Python overhead
BF16 autocast torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16) Throughput vs FP32
BF16 embeddings init_weights() casts wte and value embeds to bf16 Lower memory, faster
Pinned-memory dataloader cpu_buffer = torch.empty(..., pin_memory=True) Faster H2D copies
GC disabled mid-training gc.freeze(); gc.disable() after step 0 Avoids ~500 ms GC stalls
Sliding-window attention WINDOW_PATTERN = "SSSL" Half-context layers are cheaper than full-context
GQA (n_kv_head <= n_head) Default n_kv_head = n_head = 6 Allows future GQA experiments without code changes
Best-fit BPE packing make_dataloader packs docs to fill rows 100% utilization, no padding

6.4 Tuning levers

The agent typically explores along these axes, in roughly this order of impact:

  1. WINDOW_PATTERN (cheap to change, very high impact).
  2. MATRIX_LR, EMBEDDING_LR (cheap, sensitive).
  3. WEIGHT_DECAY, WARMDOWN_RATIO (mid-impact, mid-risk).
  4. DEPTH, ASPECT_RATIO (high impact, but risk OOM).
  5. Optimizer internals — Newton-Schulz steps, momentum schedules (high reward, high risk).

Tip

If you want to baseline a new GPU without an agent, run the default config three times and average val_bpb. Variance run-to-run on identical code is small (<0.005) but not zero.


7. Security

7.1 Threat model

autoresearch is research code, not a service. The relevant security posture has three flavours:

  1. Network → host: the training process makes outbound HTTPS to Hugging Face and downloads Parquet shards, plus pulls a Python wheel for FA3. No inbound exposure.
  2. Dataset → model: the training corpus (climbmix-400b-shuffle) is public web text. Treat the trained model accordingly; do not assume any non-public information is present.
  3. Agent → repo: an autonomous coding agent has shell + git + write access to your working tree. This is the largest attack surface.

Caution

When you launch a coding agent with permissions disabled, you are explicitly handing it the ability to run arbitrary shell commands and rewrite your train.py. Run it only inside a checkout you can afford to lose. Do not run on a host where you have credentials in ~/.aws/, ~/.config/gcloud/, or ~/.ssh/ that the agent could exfiltrate by accident — even though it has no instruction to do so.

7.2 Secrets handling

There are none in the repo:

  • requests.get(url, ...) for shards uses no auth.
  • kernels.get_kernel(repo) does not require a Hugging Face token for the public FA3 wheels listed.
  • No OpenAI/Anthropic/W&B keys are referenced in the codebase.

If you wire up an external API key (for example, to log results to a cloud service), keep it out of the working tree. Set it in your shell once and pass it via os.environ from a wrapper script outside train.py. Otherwise the agent might accidentally include it in a commit.

7.3 Credential and data residency

Concern Mitigation
Hugging Face token Not required for upstream operation. If you set one (e.g. for a private fork of the dataset), keep it in the shell, not in pyproject.toml.
Local cache ~/.cache/autoresearch/ is a per-user directory. On shared hosts, set restrictive umask before first prep so other users cannot read your shards.
Trained weights Never persisted. There is no model checkpoint to leak.
results.tsv Untracked; survives git reset. Contains commit hashes, val_bpb, peak VRAM, descriptions. Treat it like any other research log.

7.4 Supply-chain considerations

  • uv sync resolves from PyPI plus the pytorch-cu128 index (download.pytorch.org). Both are first-party.
  • kernels.get_kernel("varunneal/flash-attention-3") and kernels.get_kernel("kernels-community/flash-attn3") pull arbitrary code from Hugging Face Hub at runtime. You implicitly trust those repos; an attacker who compromises either Hub repo can execute code on your GPU box during training startup.
  • rustbpe is a compiled native dependency. It runs only during prepare.py (tokenizer training) — not on the training hot path.

Warning

The Flash Attention 3 kernel is fetched fresh from Hugging Face Hub on every train.py start, unless your kernels cache already has the wheel. If you are running in a regulated environment, mirror the kernel wheel internally and set HF_HUB_OFFLINE=1 after the first run.

7.5 Agent boundary

The clearest line in program.md is prepare.py is read-only. That keeps the metric definition honest, but it is not enforced by the filesystem — it is enforced by the prompt. If you cannot trust your agent to follow that rule, you cannot trust its val_bpb deltas either. Sanity-check by git diff master -- prepare.py after a run and confirming the file is unchanged.

Important

A misbehaving agent that "improves" val_bpb by editing evaluate_bpb to return a lower number is the most plausible failure mode. Always inspect the diff before promoting an agent's branch.


8. Developer Internals

8.1 prepare.py deep dive

8.1.1 Data download

def download_single_shard(index):
    # 5 attempts, exponential backoff (2^attempt seconds)
    # Atomic temp-file → rename pattern

Each shard is downloaded into {filepath}.tmp and atomically os.renamed to the final name. A partial file therefore cannot be mistaken for a complete shard. Failure modes are caught via requests.RequestException and IOError; .tmp artifacts are cleaned up on each failure.

download_data(num_shards, download_workers=8) builds the index set [0..num_train) ∪ {VAL_SHARD}, filters out already-downloaded shards, then maps over the remainder with multiprocessing.Pool. Workers are capped at min(8, needed).

Note

Issue #215 ("Harden downloaded dataset shard cache in prepare.py") tracks adding checksum verification to the cache, which would let prepare.py detect a corrupt-but-non-empty shard. Today, integrity is "the file exists, non-empty, and not .tmp".

8.1.2 Tokenizer training

def train_tokenizer():
    # 1. Skip if both tokenizer.pkl and token_bytes.pt exist.
    # 2. rustbpe.Tokenizer().train_from_iterator(text_iterator(), ...)
    # 3. Convert merges to tiktoken.Encoding.
    # 4. Pickle tiktoken.Encoding to tokenizer.pkl.
    # 5. Build token_bytes (torch.int32 tensor of UTF-8 byte lengths).
    # 6. Sanity check: encode/decode roundtrip on a multilingual string.

Two artifacts are written:

  • tokenizer.pkl — pickled tiktoken.Encoding.
  • token_bytes.pt — torch.int32 tensor of length n_vocab where token_bytes[i] = len(enc.decode([i]).encode("utf-8")) for normal tokens and 0 for the 4 reserved specials.

token_bytes is the lookup that lets evaluate_bpb divide cross-entropy nats by total UTF-8 bytes regardless of vocab size.

8.1.3 Dataloader

_document_batches(split) is an infinite iterator over batches of raw text from Parquet row groups. It excludes the validation shard for split="train" and uses only the validation shard for split="val". Every iteration yields (batch, epoch) so the training loop can report epoch in the streaming log.

make_dataloader(tokenizer, B, T, split, buffer_size=1000) wraps _document_batches with:

  • BPE encoding via tokenizer.encode(...), prepending the BOS token.
  • A document buffer that refills to at least buffer_size documents.
  • Best-fit packing: for each row of length T+1, pick the largest document that fits the remaining space. If none fits, crop the shortest document to fill the row exactly.
  • Pre-allocated CPU pinned and CUDA tensors for inputs and targets.

This achieves 100% utilization (no padding tokens) at the cost of cropping the worst-case document. inputs = row[:, :-1], targets = row[:, 1:], standard next-token-prediction.

8.1.4 evaluate_bpb

@torch.no_grad()
def evaluate_bpb(model, tokenizer, batch_size):
    token_bytes = get_token_bytes(device="cuda")
    val_loader = make_dataloader(tokenizer, batch_size, MAX_SEQ_LEN, "val")
    steps = EVAL_TOKENS // (batch_size * MAX_SEQ_LEN)
    total_nats = 0.0
    total_bytes = 0
    for _ in range(steps):
        x, y, _ = next(val_loader)
        loss_flat = model(x, y, reduction='none').view(-1)
        y_flat = y.view(-1)
        nbytes = token_bytes[y_flat]
        mask = nbytes > 0          # exclude special tokens
        total_nats += (loss_flat * mask).sum().item()
        total_bytes += nbytes.sum().item()
    return total_nats / (math.log(2) * total_bytes)

Notes:

  • Uses reduction='none' so per-token nats are recovered before masking.
  • Special tokens have token_bytes == 0; the mask zeroes them in both numerator and denominator.
  • nats / (ln(2) * total_bytes) converts to bits/byte.
  • EVAL_TOKENS is fixed; steps is therefore deterministic given batch_size.

Note

The comment in source says "DO NOT CHANGE — this is the fixed metric". program.md mirrors this rule. Anything that changes nats-per-token semantics or how byte counts are aggregated invalidates the leaderboard.

8.2 train.py deep dive

8.2.1 Model — GPT

A minimalist GPT with these design choices:

Choice Reason
RMSNorm via F.rms_norm Lighter than LayerNorm; standard in modern decoder-only LMs
Pre-norm + residual x = x + attn(norm(x), ...)
Rotary embeddings Precomputed once for sequence_len * 10, registered as non-persistent buffers
Grouped-query attention support n_kv_head <= n_head; defaults to n_kv_head == n_head
Sliding-window attention Per-layer window size driven by WINDOW_PATTERN
ReLU² in MLP x = F.relu(x).square(); c_fc to 4 * n_embd, c_proj back
Logits softcap softcap=15, softcap * tanh(logits / softcap) for stability
resid_lambdas, x0_lambdas Per-layer mixing scalars: x = resid_lambdas[i] * x + x0_lambdas[i] * x0. Allows the model to learn how much of the original input to keep at each layer.
Value embeddings (ResFormer) Alternating layers (and always the last) get a learned Embedding(vocab_size, kv_dim) whose output is mixed into v via an input-dependent gate of the first 32 channels.

The value-embedding gate:

gate = 2 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
v = v + gate.unsqueeze(-1) * ve

Initialized with zeros (init_weights() zeros ve_gate.weight), so 2 * sigmoid(0) = 1.0 at start — the value-residual pathway is identity at init and the network learns to bend it.

has_ve(layer_idx, n_layer) returns True for alternating layers, with the last layer always included. So at DEPTH=8 (even), layers 1, 3, 5, 7 carry value embeddings.

8.2.2 MuonAdamW optimizer

class MuonAdamW(torch.optim.Optimizer):
    """Combined optimizer: Muon for 2D matrix params, AdamW for others."""

setup_optimizer builds parameter groups by kind:

  • kind='adamw' for: lm_head weights, wte weights, value-embed weights, resid_lambdas, x0_lambdas.
  • kind='muon' for: every distinct shape inside transformer.h parameters.

The split is per-shape so _step_muon can stack same-shape tensors, run Muon on the stack, and unbind back into the parameter list:

stacked_grads = torch.stack([p.grad for p in params])
stacked_params = torch.stack(params)
muon_step_fused(stacked_grads, stacked_params, ...)
torch._foreach_copy_(params, list(stacked_params.unbind(0)))

Both adamw_step_fused and muon_step_fused are @torch.compile(dynamic=False, fullgraph=True). To keep torch.compile from recompiling on hyperparameter changes, all scalar hyperparameters (LR, betas, eps, weight decay, momentum) are pre-allocated as 0-D CPU tensors (self._adamw_lr_t, etc.) and .fill_()-ed before each step. The fused functions take tensors, not Python floats, so they retain their compiled graph across LR-schedule changes.

8.2.3 Muon step internals

Muon performs orthogonalization of the gradient via polar-express (a.k.a. Newton-Schulz-style) iteration:

polar_express_coeffs = [
    (8.156554524902461, -22.48329292557795, 15.878769915207462),
    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
]

Five iteration coefficients, applied in order (ns_steps=5 by default). Each iteration computes A = X.mT @ X (or X @ X.mT, whichever is smaller) and updates X = a * X + X @ B where B = b * A + c * (A @ A).

After orthogonalization comes NorMuon variance reduction: the step is rescaled by the running second moment of the (normalized) gradient. The result is a Muon update that respects per-row or per-column variance like AdamW does.

Cautious weight decay applies only when (g * params) >= 0:

mask = (g * stacked_params) >= 0
stacked_params.sub_(lr * g + lr * wd * stacked_params * mask)

This is the "cautious update" trick — weight decay does not push parameters in the direction of -g when they are already moving the same way.

8.2.4 Training loop schedule

Three schedules track progress = total_training_time / TIME_BUDGET:

Schedule Logic
get_lr_multiplier(progress) Linear warmup over WARMUP_RATIO of training, flat at 1.0, then linear cooldown over WARMDOWN_RATIO of training to FINAL_LR_FRAC.
get_muon_momentum(step) Linear from 0.85 to 0.95 over the first 300 steps.
get_weight_decay(progress) WEIGHT_DECAY * (1 - progress). Decays linearly to 0.

The training loop:

while True:
    for micro_step in range(grad_accum_steps):
        with autocast_ctx:
            loss = model(x, y)
        train_loss = loss.detach()
        loss = loss / grad_accum_steps
        loss.backward()
        x, y, epoch = next(train_loader)

    # update schedules → optimizer.step() → zero_grad
    ...
    if step > 10 and total_training_time >= TIME_BUDGET:
        break

grad_accum_steps = TOTAL_BATCH_SIZE // (DEVICE_BATCH_SIZE * MAX_SEQ_LEN) — at default config that is 2.

8.3 Tests

There are no tests in the repository (no tests/, no pytest config, no CI workflow visible). The de-facto regression test is "the baseline run prints val_bpb ≈ 1.0". Forks adding multi-GPU or non-NVIDIA support typically introduce their own validation harness.

8.4 Code patterns and conventions

Pattern Where
if __name__ == "__main__": only in prepare.py train.py runs imperatively at module level — there is no main()
Module-level constants with ALL_CAPS for tunables train.py hyperparameters block
Single-file module per concern prepare.py for harness; train.py for model + loop
Atomic-rename file writes prepare.py:download_single_shard
@torch.compile(dynamic=False, fullgraph=True) All hot kernels (adamw_step_fused, muon_step_fused)
torch.cuda.synchronize() around t0/t1 Ensures wall-clock measurements are not optimistic
Pre-allocated 0-D tensors for compile-stable hparams MuonAdamW.__init__
torch.no_grad() decorator on eval and init evaluate_bpb, GPT.init_weights, optimizer step

8.5 Improvement opportunities

These are the visible gaps that a future agent (or human contributor) might pursue:

  • Issue #215 — Harden the dataset cache: checksum each shard so corrupt files are detected and re-downloaded.
  • Issue #35 — Investigate / document the FLOPs-vs-time-budget tradeoff. The current design fixes time; some users prefer a fixed FLOPs budget for hardware-independent comparisons.
  • Issue #519 — Add the Ascend-compatible fork to the README's notable-forks list.
  • Multi-GPU support (intentionally absent) — A future agent could reintroduce a thin DDP wrapper without breaking the metric, because evaluate_bpb is a global aggregate.
  • Lighter forward path for evaluation — evaluate_bpb re-uses the training graph. A separate eval-only graph could skip the compiled-loss wrapping and run faster, freeing time budget for training.

Note

Any "improvement" that touches prepare.py requires retraining the tokenizer and re-downloading shards. The branch contract is that prepare.py lives outside the experiment loop. Improvements there should arrive as one-off PRs from a human contributor, not as agent commits.

8.6 Contributing

This is a personal research repo on karpathy/autoresearch. Pull requests are accepted but the maintainer's bar is high — the project has stayed deliberately small. Forks for other hardware backends are welcome and tend to be merged into the README's notable-forks list rather than the main branch.

If you submit a PR:

  1. Keep prepare.py byte-identical unless you are explicitly changing the eval harness in coordination with the maintainer.
  2. Run the baseline to confirm val_bpb is unchanged on default config (or document the change clearly).
  3. Avoid adding dependencies. The 9-package surface is intentional.

9. Integration Points

autoresearch is not designed to integrate into other systems. It has no API, no CLI surface beyond prepare.py, no on-disk format intended for external consumption. The integration story is exclusively about the agent loop.

9.1 The agent contract

The agent must:

  • Read program.md for the experimental protocol.
  • Stay on a autoresearch/<tag> branch; never touch master.
  • Never modify prepare.py, pyproject.toml, uv.lock, .python-version, or program.md.
  • Run uv run train.py > run.log 2>&1 without modification.
  • Parse ^val_bpb: and ^peak_vram_mb: from the run log.
  • Append a row to results.tsv in the documented format.

9.2 results.tsv format

Column Type Meaning
commit 7-char short git hash The commit the experiment was run on
val_bpb %.6f Final BPB; 0.000000 on crash
memory_gb %.1f Peak VRAM in GB (i.e. peak_vram_mb / 1024); 0.0 on crash
status keep / discard / crash What the agent did with this experiment
description text (no tabs) Short experiment description

Example:

commit	val_bpb	memory_gb	status	description
a1b2c3d	0.997900	44.0	keep	baseline
b2c3d4e	0.993200	44.2	keep	increase matrix_lr to 0.06
c3d4e5f	1.005000	44.0	discard	switch MLP activation to GeLU
d4e5f6g	0.000000	0.0	crash	double model width (OOM)

Important

results.tsv is intentionally untracked. Do not commit it. The agent depends on it surviving git reset so it can keep its log even when the working tree is rolled back.

9.3 External services touched

Service When Required?
Hugging Face datasets (huggingface.co) prepare.py shard downloads Yes for first prep
Hugging Face Hub kernels API train.py startup, FA3 wheel resolution Yes unless cached + HF_HUB_OFFLINE=1
download.pytorch.org/whl/cu128 uv sync Yes for first install
GitHub git push / git pull Optional, only if you mirror branches off-host

9.4 Embedding autoresearch into a larger pipeline

The intended embedding is the agent itself, not a programmatic API. Practical patterns:

  • CI runner per GPU. A self-hosted runner picks a tag, runs the agent for N hours, archives the resulting results.tsv, then resets the working tree.
  • Cron-driven baselines. Run uv run train.py nightly on the master branch, append a single row to a long-running baselines.tsv outside the repo. Use this to detect regressions caused by upstream PyTorch / FA3 / driver changes.
  • Multi-tag tournaments. Spawn N branches with N different program.md strategies on N GPUs, then diff their results.tsv files at the end of the run to compare strategies.

Tip

If you want a lightweight programmatic entry point, write your own thin Python script that imports Tokenizer, make_dataloader, and evaluate_bpb from prepare, builds a model, and runs evaluation only. This works because prepare.py exposes those exact symbols at module level. It does not require any code changes inside the repo.