# autoresearch Administrator & Developer Guide

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`](https://github.com/karpathy/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

```text
 ┌─────────────────────────────────────────────────────────────┐
 │                  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 checkout`s a working branch, edits `train.py`, `git commit`s.
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:

```python
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:

```bash
# 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:
>
> ```bash
> 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.log`s. Useful greps:

```bash
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.md`s 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

```bash
# 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:**

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

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

```bash
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:**

```bash
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:

```python
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

```python
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.rename`d 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

```python
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`

```python
@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:

```python
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

```python
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:

```python
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:

```python
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`:

```python
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:

```python
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:

```text
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.

---

_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. Unofficial; not affiliated with or reviewed by the project._