# nanochat Admin & Developer Guide

This guide is the deep technical reference for nanochat: its architecture, the GPT model internals, the training/inference pipeline, configuration, deployment of the chat server, security posture, code quality, and improvement opportunities. It is intended for engineers, researchers, and operators who need to understand, deploy, extend, or fork the codebase — not just run it.

nanochat is a single-node LLM training and inference harness. There is no distributed control plane, no database, no service mesh — it is a set of Python scripts that read and write files under one base directory, optionally launched across multiple GPUs on one machine via `torchrun`. The design philosophy is deliberate minimalism: a "strong baseline" codebase that is cohesive, readable, and maximally forkable, with **one dial of complexity** (`--depth`) that determines all other model hyperparameters.

> [!IMPORTANT]
> Everything in this guide is derived directly from the source at commit `dc54a1a`. Where the code carries an explicit caveat (e.g. "not a true security sandbox", "blindly following AdamW theory for Muon"), this guide reproduces it faithfully rather than papering over it.

---

## Table of Contents

1. [Executive Summary](#1-executive-summary)
2. [Architecture Overview](#2-architecture-overview)
3. [The GPT Model (Internals)](#3-the-gpt-model-internals)
4. [The Optimizer: MuonAdamW](#4-the-optimizer-muonadamw)
5. [Precision: Explicit dtype & FP8](#5-precision-explicit-dtype--fp8)
6. [Attention: FA3 with SDPA Fallback](#6-attention-fa3-with-sdpa-fallback)
7. [The Inference Engine & Tool Use](#7-the-inference-engine--tool-use)
8. [Data Pipeline](#8-data-pipeline)
9. [Checkpoints & the Base Directory](#9-checkpoints--the-base-directory)
10. [Module Reference](#10-module-reference)
11. [Requirements & Dependencies](#11-requirements--dependencies)
12. [Configuration Reference](#12-configuration-reference)
13. [Deploying the Chat Server](#13-deploying-the-chat-server)
14. [Distributed Training Operations](#14-distributed-training-operations)
15. [Security Analysis](#15-security-analysis)
16. [Code Quality & Conventions](#16-code-quality--conventions)
17. [Testing Strategy](#17-testing-strategy)
18. [Potential Improvements](#18-potential-improvements)
19. [Contributing](#19-contributing)
20. [Quick Reference](#20-quick-reference)

---

## 1. Executive Summary

nanochat covers the entire LLM lifecycle in one repository:

| Concern | nanochat's answer |
|---------|-------------------|
| **Scope** | Tokenization → pretraining → SFT → RL → evaluation → inference → chat UI. |
| **Scale target** | A single GPU node (1–8 GPUs). GPT-2 capability in ~3 hours / ~$48 on 8×H100. |
| **Configurability** | One dial (`--depth`); all other hyperparameters auto-derived via scaling laws. |
| **State** | The filesystem. One base directory (`~/.cache/nanochat`) holds data, tokenizer, checkpoints, eval output, and reports. No DB, no message broker. |
| **Distribution** | PyTorch DDP via `torchrun`. Single-GPU runs use gradient accumulation. |
| **Inference** | A custom KV-cache `Engine` with batched generation and a calculator tool-use state machine. |
| **Serving** | A FastAPI web server (`chat_web`) with a multi-GPU worker pool, streaming SSE responses, and a static HTML UI. |
| **License** | MIT. |

The codebase is intentionally small (~14 core modules, 8 task modules, 9 scripts) and reads top-to-bottom. There are no abstraction layers you must learn — the training loop, the model, and the optimizer are each a single readable file.

---

## 2. Architecture Overview

```text
                         ┌───────────────────────────────────────────────────────────┐
                         │                      scripts/  (entry points)              │
                         │  tok_train  base_train  chat_sft  chat_rl  chat_eval ...    │
                         └───────┬───────────────┬───────────────┬───────────────┬────┘
                                 │ imports        │ imports        │ imports        │
              ┌──────────────────▼────────────────▼────────────────▼────────────────▼─────────────────┐
              │                                    nanochat/  (core library)                            │
              │                                                                                          │
              │  common.py        gpt.py            optim.py        engine.py        tokenizer.py        │
              │  (device, ddp,    (GPT nn.Module:   (MuonAdamW,     (KV-cache infer, (BPE: rustbpe       │
              │   dtype, dirs,     attn, MLP,        DistMuonAdamW)  tool-use FSM)     train / tiktoken   │
              │   flops table)     scaling-law math)                                  infer)             │
              │                                                                                          │
              │  dataloader.py    dataset.py        checkpoint_     core_eval.py     loss_eval.py        │
              │  (bos-bestfit     (parquet shards,  manager.py      (CORE / DCLM)    (bits per byte)     │
              │   packing)         HF download)     (save/load)                                          │
              │                                                                                          │
              │  flash_attention.py   fp8.py        execution.py    report.py        ui.html / logo.svg  │
              │  (FA3 | SDPA)         (Float8Linear) (sandbox exec)  (markdown report) (web frontend)     │
              └──────────────────────────────────────────────────────────────────────────────────────┬─┘
                                 │ used by                                                              │
              ┌──────────────────▼──────────────────┐                       ┌──────────────────────────▼──┐
              │            tasks/  (eval/SFT)         │                       │   runs/  (orchestration .sh) │
              │  Task base + ARC, MMLU, GSM8K,        │                       │  speedrun, runcpu,           │
              │  HumanEval, SpellingBee, SmolTalk,    │                       │  miniseries, scaling_laws    │
              │  CustomJSON, TaskMixture/Sequence     │                       └──────────────────────────────┘
              └───────────────────────────────────────┘
```

**Layering.** `scripts/` are thin orchestrators that parse arguments and drive a loop; all reusable logic lives in `nanochat/`. `tasks/` defines evaluation and SFT datasets behind a common `Task` interface. `runs/` are bash scripts that chain the Python scripts into full pipelines.

**Data flow.** `dataset.py` downloads parquet shards → `dataloader.py` tokenizes and packs them into `(B, T)` batches → `gpt.py` forwards them → `optim.py` updates weights → `checkpoint_manager.py` persists state → `engine.py` loads a checkpoint and generates → `chat_web.py`/`chat_cli.py` expose it.

**No hidden state.** Every artifact is a file under the base directory. To reset, delete the relevant subdirectory; to resume, point a script at the saved step. This is the same "filesystem is the database" simplicity throughout.

---

## 3. The GPT Model (Internals)

The model lives entirely in `nanochat/gpt.py`. It is a decoder-only transformer with a set of deliberate, modern design choices (the docstring enumerates them):

- **Rotary positional embeddings** (RoPE), no learned position embeddings. Base theta = 100000; the cache is over-computed 10× the sequence length.
- **QK-norm** — queries and keys are RMS-normalized before attention, then scaled by `1.2` (split between Q and K) for sharper attention.
- **ReLU² MLP** — `c_fc → relu(x).square() → c_proj`, hidden dim 4× model dim.
- **Untied weights** for token embedding (`wte`) and output head (`lm_head`).
- **No biases** in linear layers; **no learnable params** in RMSNorm.
- **Group-Query Attention (GQA)** support (`n_kv_head ≤ n_head`).
- **Sliding-window attention** via a per-layer window pattern (see below).
- **Pre-norm residual blocks**: `x = x + attn(norm(x))`, `x = x + mlp(norm(x))`.

### Notable architectural additions

The model goes beyond a vanilla GPT with several learnable mechanisms (inspired by modded-nanoGPT and ResFormer):

| Mechanism | Parameter | What it does |
|-----------|-----------|--------------|
| **Residual scaling** | `resid_lambdas` (per layer) | Scales the residual stream per layer. Init decays `1.15 → 1.05` across depth. |
| **x0 blending** | `x0_lambdas` (per layer) | Blends the initial normalized embedding back in at each layer. Init decays `0.20 → 0.05`. |
| **Token smear** | `smear_lambda`, `smear_gate` | Mixes the previous token's embedding into the current position (cheap bigram-like signal), gated by a small learned network on the first 24 channels. |
| **Backout** | `backout_lambda` | Subtracts the mid-layer (halfway) residual before the final norm, to remove low-level features before logit projection. Init `0.2`. |
| **Value embeddings** | `value_embeds` (alternating layers) | ResFormer-style value residual: per-head input-dependent gate (range 0–3) mixes a learned value embedding into `v`. Present on alternating layers; the last layer always included. |

### Sliding-window attention

`window_pattern` (default `SSSL`) is a string tiled across layers: `L` = full context, `S` = quarter context (ceil-rounded to a 128 tile, e.g. 2048 → 768). The **final layer is always `L`** regardless of the pattern. The per-layer window sizes are precomputed in `_compute_window_sizes()` and passed to Flash Attention as `(left, right)` tuples.

> [!NOTE]
> Under the PyTorch SDPA fallback, sliding-window attention is emulated with explicit masks and is much slower. On non-Hopper GPUs, prefer `--window-pattern L` to avoid this penalty.

### Meta-device initialization (a footgun to know)

`GPT.__init__` runs under `torch.device("meta")` — it only establishes shapes/dtypes, **no real data**. All actual parameter initialization happens in `init_weights()`. The standard build sequence is:

```python
with torch.device("meta"):
    model = GPT(config)      # shapes only
model.to_empty(device=device) # allocate storage (uninitialized)
model.init_weights()          # real initialization
```

> [!WARNING]
> Any computation that needs real tensor values must NOT live in `__init__`. This is a common source of bugs when extending the model — put real initialization in `init_weights()`.

### Forward pass shape

The forward computes logits as `(B, T, padded_vocab_size)`, with the vocab padded up to a multiple of 64 for tensor-core efficiency (outputs are cropped logically). Logits are soft-capped to `[-15, 15]`. Loss uses cross-entropy with `ignore_index=-1` (the masking convention used throughout for padding and non-target tokens).

### FLOPs and scaling-parameter accounting

`estimate_flops()` returns FLOPs/token (forward+backward) as `6 × (matmul params) + attention FLOPs`, accounting for the reduced effective sequence length of sliding-window layers. `num_scaling_params()` returns a breakdown (`wte`, `value_embeds`, `lm_head`, `transformer_matrices`, `scalars`, `total`) so scaling-law analysis can pick a convention. `base_train` uses `transformer_matrices + lm_head` as the "scaling params" for the cleanest scaling laws.

---

## 4. The Optimizer: MuonAdamW

`nanochat/optim.py` defines a combined optimizer with two flavors selected per parameter group via a `kind` field:

- **Muon** (`kind='muon'`) — used for the 2D matrix parameters of the transformer blocks. Muon orthogonalizes the momentum-smoothed gradient (Newton-Schulz iteration, `ns_steps=5`) before the update. Matrix params are grouped by shape so they can be stacked and processed together.
- **AdamW** (`kind='adamw'`) — used for everything else: embeddings, the LM head, value embeddings, and the per-layer scalars (`resid_lambdas`, `x0_lambdas`, smear/backout params).

`DistMuonAdamW` is the distributed variant used under DDP; `MuonAdamW` is the single-process version. `GPT.setup_optimizer()` builds the param groups with all hyperparameters explicit (per-group `lr`, `betas`, `eps`, `weight_decay`), including a `1/√(model_dim/768)` learning-rate scale for the AdamW groups (tuned at a 768-dim reference model).

> [!NOTE]
> The codebase is candid that Muon's scaling rules are extrapolated from AdamW theory: "Note that these papers study AdamW, *not* Muon. We are blindly following AdamW theory for scaling hoping it ~works for Muon too." Treat the LR/weight-decay scaling as empirical, not proven, for Muon.

### Hyperparameter auto-derivation (in `base_train`)

From `--depth` alone, `base_train` derives:

1. **Token horizon** = `target_param_data_ratio × scaling_params`.
2. **Batch size** ∝ `D^0.383` (Power Lines paper), clamped to a power of two.
3. **LR correction** η ∝ `√(B/B_ref)` (AdamW sqrt scaling, reused for Muon).
4. **Weight-decay scaling** to keep the `T_epoch = B/(η·λ·D)` constant.

Schedules: linear LR warmup → constant → linear warmdown; Muon momentum warms 0.85→0.97 then down to 0.90; Muon weight decay follows a cosine decay to zero over training.

---

## 5. Precision: Explicit dtype & FP8

nanochat **does not use `torch.amp.autocast`.** Precision is managed explicitly through a single global `COMPUTE_DTYPE` (in `common.py`), auto-detected per hardware (`_detect_compute_dtype`):

| Hardware | dtype | Reason |
|----------|-------|--------|
| CUDA SM ≥ 80 | `bfloat16` | Native bf16 tensor cores. |
| CUDA SM < 80 | `float32` | No bf16; fp16 requires a GradScaler (available via `NANOCHAT_DTYPE=float16`). |
| CPU / MPS | `float32` | No reduced-precision tensor cores. |

Override with `NANOCHAT_DTYPE` ∈ `{bfloat16, float16, float32}`.

**How mixed precision works without autocast:** master weights stay fp32 (for optimizer precision); the custom `Linear` layer casts weights to the activation dtype in `forward`. Embeddings are stored directly in `COMPUTE_DTYPE` to save memory (except under fp16, which needs fp32 embeddings because GradScaler can't unscale fp16 grads). fp16 training enables a `torch.amp.GradScaler` in the training loops (SFT supports it; RL does not).

### FP8 training (`fp8.py`)

`--fp8` swaps eligible `nn.Linear` layers for a custom `Float8Linear` (a ~150-line replacement for torchao's ~2000-line implementation). It supports only the **tensorwise** dynamic-scaling recipe. Mechanics:

- A single `autograd.Function` (`_Float8Matmul`, marked `@torch._dynamo.allow_in_graph`) quantizes operands to FP8, calls `torch._scaled_mm` (cuBLAS FP8 kernel, ~2× bf16), and dequantizes.
- Forward and weight use `float8_e4m3fn` (more mantissa); gradients use `float8_e5m2` (wider range).
- Only layers with both dims divisible by 16 and `min(in,out) ≥ 128` are converted (`fp8_module_filter`).
- Evaluation temporarily swaps FP8 layers back to bf16 `Linear` (`disable_fp8` context manager) for consistent metrics.

> [!CAUTION]
> FP8 training requires CUDA + an H100-class GPU. On non-CUDA devices `--fp8` is ignored with a warning. The `rowwise` recipe is **not** implemented — requesting it raises `ValueError: Only 'tensorwise' recipe is supported`.

---

## 6. Attention: FA3 with SDPA Fallback

`nanochat/flash_attention.py` exports a `flash_attn` namespace that matches the Flash Attention 3 API exactly but transparently falls back to PyTorch SDPA. Resolution happens once at import:

- FA3 kernels are loaded only on **Hopper (SM 90)** GPUs, via the `kernels` library (`varunneal/flash-attention-3`). Ada (SM 89) and Blackwell (SM 100) fall back until FA3 is recompiled for them.
- Even on Hopper, FA3 is used only in **bf16** (FA3 supports bf16/fp8; fp16/fp32 use SDPA).
- The SDPA fallback (`_sdpa_attention`) handles three cases: full-context equal-length, single-token decode (slices the KV window), and chunked/sliding-window inference (builds an explicit boolean mask). It supports GQA via `enable_gqa`.

`USE_FA3` and `HAS_FA3` are module-level booleans the training scripts read to print warnings. A test override (`_override_impl`) can force `'fa3'` or `'sdpa'` for the fallback test.

---

## 7. The Inference Engine & Tool Use

`nanochat/engine.py` provides efficient generation, decoupled from tokenization (it operates purely on token-id lists).

### KV cache

`KVCache` is built for FA3's `flash_attn_with_kvcache` API: tensors are `(n_layers, B, T, H, D)`, position is tracked per batch element via an int32 `cache_seqlens` tensor, and FA3 updates the cache in place. It also stores `prev_embedding` for the token-smear mechanism during decode. `prefill()` copies a batch-1 prefill cache into a batch-N decode cache so the prompt is processed once and then expanded across samples.

### Generation flow

`Engine.generate()`:
1. Batch-1 prefill of the prompt; take the last-position logits.
2. Replicate the KV cache to `num_samples` rows.
3. Per-row `RowState` tracks the token sequence, a forced-token deque, and tool-block state.
4. Loop: sample next token per row (`sample_next_token` with temperature/top-k; temperature 0 = argmax), update state, run tool logic, yield `(token_column, token_masks)`.

`generate_batch()` wraps this into non-streaming batched generation, returning final sequences with terminal tokens stripped. Termination tokens are `<|assistant_end|>` and `<|bos|>`.

### Tool use: the calculator state machine

The chat special tokens coordinate tool use:

| Token | Role |
|-------|------|
| `<|bos|>` | Beginning of sequence / row terminator. |
| `<|user_start|>`, `<|user_end|>` | Wrap a user message. |
| `<|assistant_start|>`, `<|assistant_end|>` | Wrap an assistant message. |
| `<|python_start|>`, `<|python_end|>` | Wrap a calculator/Python expression the model wants evaluated. |
| `<|output_start|>`, `<|output_end|>` | Wrap the tool's result, force-injected back into the stream. |

When the model emits `<|python_end|>` closing a block, the Engine decodes the captured expression and runs `use_calculator()`, which:

- Strips commas; for pure-math expressions, evaluates with a 3-second timeout via `eval(formula, {"__builtins__": {}}, {})`. The `**` power operator is **disallowed**.
- For string operations, allows only `.count(...)` and blocks dangerous patterns (`__`, `import`, `exec`, `eval`, `open`, `globals`, etc.).
- On success, the result tokens are force-injected wrapped in `<|output_start|>`/`<|output_end|>` (masked as non-sampled, so they aren't trained on in RL).

> [!NOTE]
> `use_calculator` is a constrained arithmetic/counting tool, not a general Python interpreter. It is what lets a tiny model answer "how many r's in strawberry" by counting rather than guessing.

### Separately: the HumanEval sandbox (`execution.py`)

`execution.py` provides `execute_code()` for evaluating model-generated Python (used by the HumanEval task). Each execution runs in its own process with a timeout, a 256 MB memory limit (`RLIMIT_AS/DATA/STACK`), captured stdout/stderr, a temp working directory, and a `reliability_guard()` that nulls out destructive functions (`os.system`, `os.kill`, `shutil.rmtree`, `subprocess.Popen`, etc.).

> [!CAUTION]
> The module's own docstring is explicit: **this is not a true security sandbox.** Network access is not blocked, `ctypes`/dynamic features can bypass restrictions, and there is no kernel-level isolation (no seccomp, containers, or VMs). It protects against accidental destructive behavior during evaluation, **not** against malicious adversarial code. Run evaluations on untrusted models inside a real sandbox (container/VM).

---

## 8. Data Pipeline

### Dataset (`dataset.py`)

The pretraining corpus is **ClimbMix-400B** (`huggingface.co/datasets/karpathy/climbmix-400b-shuffle`), distributed as 6543 parquet shards (`shard_00000.parquet` … `shard_06542.parquet`), each ~250M chars / ~100 MB compressed. Stored under `<base_dir>/base_data_climbmix`. The **last shard (6542) is always the validation split**; everything else is train. Downloads use a multiprocessing `Pool`, write to a `.tmp` file then atomically rename, and retry up to 5 times with exponential backoff.

> [!NOTE]
> Legacy support: if the ClimbMix directory is absent, the code warns about the FinewebEdu-100B → ClimbMix migration (March 4 2026) and falls back to the old `base_data` directory. This fallback path is marked for eventual deletion.

### Dataloader (`dataloader.py`)

`tokenizing_distributed_data_loader_with_state_bos_bestfit` implements **BOS-aligned best-fit packing**:

- Every row starts with `<|bos|>`.
- Documents are tokenized in batches (multithreaded) into a buffer; for each row, the **largest document that fits** the remaining space is placed; when nothing fits, the shortest buffered doc is cropped to fill exactly.
- Result: 100% utilization (no padding), but ~35% of tokens are discarded to cropping. The trade-off buys cleaner contexts (every token can attend back to its document's BOS).
- DDP sharding is by row group; an approximate resume state (`pq_idx`, `rg_idx`, `epoch`) is yielded for checkpointing.

A pre-allocated CPU pinned buffer and a single host-to-device copy minimize transfer overhead.

> [!NOTE]
> SFT uses a different, conversation-aware loader (`sft_data_generator_bos_bestfit` inside `chat_sft.py`) that packs whole conversations and pads rather than crops, masking targets to `-1` for padding, user prompts, and tool outputs (only assistant completions are trained on).

### Tokenizer (`tokenizer.py`)

Two implementations behind a shared interface:

- **`RustBPETokenizer`** — trains via the Rust `rustbpe` library and serves via `tiktoken` at inference (GPT-4 style). This is the default (`get_tokenizer()`).
- **`HuggingFaceTokenizer`** — wraps a HF `tokenizers` model, used for evaluating HF baselines.

Both expose `encode`, `decode`, `encode_special`, `get_bos_token_id`, `render_conversation` (produces token ids + loss mask for SFT), and `render_for_completion` (primes the assistant for generation during eval/RL).

---

## 9. Checkpoints & the Base Directory

`get_base_dir()` returns `$NANOCHAT_BASE_DIR` if set, else `~/.cache/nanochat` (created on demand). Layout:

```text
$NANOCHAT_BASE_DIR/
├── base_data_climbmix/          # downloaded pretraining shards (parquet)
├── tokenizer/                   # trained tokenizer + token_bytes.pt
├── base_checkpoints/<tag>/      # pretrained base models (source="base")
├── chatsft_checkpoints/<tag>/   # SFT models (source="sft")
├── chatrl_checkpoints/<tag>/    # RL models (source="rl")
├── base_eval/                   # CORE eval CSVs
├── eval_bundle/                 # downloaded CORE eval data (eval_bundle.zip)
├── identity_conversations.jsonl # synthetic identity data for SFT
└── report/                      # markdown run reports
```

### Checkpoint format (`checkpoint_manager.py`)

Per step `<step>` in a `<tag>` directory:

| File | Contents |
|------|----------|
| `model_<step>.pt` | Model `state_dict` (rank 0 only). |
| `meta_<step>.json` | Model config, user config (CLI args), val bpb, dataloader state, loop state. |
| `optim_<step>_rank<r>.pt` | Optimizer state, **sharded per rank** (each rank saves its own). |

`load_model(source, ...)` maps `base|sft|rl` to the directory, guesses the model tag (largest `d<N>`, else most recent), guesses the step (last), and rebuilds the model. Two compatibility shims patch old checkpoints: `_patch_missing_config_keys` (defaults `window_pattern="L"`) and `_patch_missing_keys` (defaults `resid_lambdas=1.0`, `x0_lambdas=0.0`). The `_orig_mod.` prefix added by `torch.compile` is stripped on load.

> [!IMPORTANT]
> Optimizer state is per-rank-sharded. To resume distributed training you must restart with the **same world size** so each rank finds its `optim_*_rank<r>.pt` shard.

---

## 10. Module Reference

### `nanochat/` core library

| Module | Responsibility | Key symbols |
|--------|----------------|-------------|
| `common.py` | Device/DDP/dtype init, base dir, file download with lock, GPU peak-FLOPS table, logging. | `compute_init`, `get_base_dir`, `COMPUTE_DTYPE`, `get_peak_flops`, `print0`, `download_file_with_lock` |
| `gpt.py` | The GPT model and scaling-law math. | `GPT`, `GPTConfig`, `Block`, `CausalSelfAttention`, `MLP`, `Linear` |
| `optim.py` | Combined Muon + AdamW optimizers (single + distributed). | `MuonAdamW`, `DistMuonAdamW`, fused step kernels |
| `engine.py` | KV-cache inference + tool-use state machine. | `Engine`, `KVCache`, `RowState`, `use_calculator`, `sample_next_token` |
| `tokenizer.py` | BPE tokenizers (train via rustbpe, infer via tiktoken; HF wrapper). | `RustBPETokenizer`, `HuggingFaceTokenizer`, `get_tokenizer`, `render_conversation` |
| `dataloader.py` | BOS-aligned best-fit distributed pretraining loader. | `tokenizing_distributed_data_loader_*_bos_bestfit` |
| `dataset.py` | Parquet shard download/iteration. | `download_single_file`, `parquets_iter_batched`, `list_parquet_files` |
| `checkpoint_manager.py` | Save/load model, optimizer, metadata. | `save_checkpoint`, `load_checkpoint`, `load_model`, `build_model` |
| `core_eval.py` | CORE / DCLM evaluation (MC, schema, LM tasks). | `evaluate_task`, `evaluate_example`, prompt renderers |
| `loss_eval.py` | Bits-per-byte evaluation. | `evaluate_bpb` |
| `flash_attention.py` | FA3 ↔ SDPA unified attention. | `flash_attn`, `HAS_FA3`, `USE_FA3` |
| `fp8.py` | Minimal tensorwise FP8 training. | `Float8Linear`, `_Float8Matmul`, `convert_to_float8_training` |
| `execution.py` | Sandboxed Python execution for HumanEval. | `execute_code`, `reliability_guard`, `ExecutionResult` |
| `report.py` | Assemble the markdown run report (system info, configs, metrics). | `Report`, `get_report`, `estimate_cost` |

### `tasks/` evaluation & SFT datasets

All tasks subclass `Task` (`tasks/common.py`) with `eval_type` (`generative` or `categorical`), `num_examples`, `get_example`, `evaluate`, and optionally `reward` (for RL). `TaskMixture` and `TaskSequence` compose tasks.

| Task | Type | Purpose |
|------|------|---------|
| `ARC` | categorical | Multiple-choice science questions (Easy/Challenge). |
| `MMLU` | categorical | Broad-topic multiple choice. |
| `GSM8K` | generative | Grade-school math (has `reward` for RL). |
| `HumanEval` | generative | Simple Python coding (runs in the sandbox). |
| `SpellingBee` / `SimpleSpelling` | generative | Letter spelling/counting (has `reward`). |
| `SmolTalk` | — | HuggingFace conversation dataset (SFT mixture). |
| `CustomJSON` | — | Arbitrary JSONL conversations (e.g. identity data). |

### `dev/` developer utilities

`gen_synthetic_data.py` (generate identity/synthetic conversations, with `--num/--workers/--output/--append/--save-metadata` flags and an `HF_TOKEN` env var), `repackage_data_reference.py` (reference for how pretraining shards were prepared), notebooks for scaling analysis, and assets.

---

## 11. Requirements & Dependencies

From `pyproject.toml` (`requires-python >= 3.10`, `torch==2.9.1`):

| Dependency | Role |
|------------|------|
| `torch==2.9.1` | Core tensor/NN/distributed framework (CUDA 12.8 or CPU index). |
| `datasets` | HuggingFace datasets (SmolTalk, MMLU, GSM8K, etc.). |
| `tiktoken` | Fast BPE inference tokenization. |
| `tokenizers` | HuggingFace tokenizer backend. |
| `rustbpe` | Rust BPE **training**. |
| `fastapi` + `uvicorn` | The chat web server. |
| `wandb` | Experiment logging (optional; `--run=dummy` disables). |
| `kernels` | Loads the FA3 kernel on Hopper. |
| `psutil` | System info for the report. |

**Dev group** (`--group dev`): `pytest`, `matplotlib`, `ipykernel`, `python-dotenv`, `transformers` (the last required to evaluate HF baselines).

`torch` is resolved from mutually-exclusive `cpu`/`gpu` extras pointing at different PyTorch wheel indexes. The transitive dependency `filelock` is used by `download_file_with_lock` to coordinate concurrent downloads across ranks. `jinja2` is used by `core_eval.py` to render eval prompts, and `pyarrow` to read parquet.

---

## 12. Configuration Reference

nanochat has **no config files** — all configuration is via CLI flags and a few environment variables. There is no `config.yaml`, no `.env` schema, no settings object.

### Environment variables

| Variable | Used in | Effect |
|----------|---------|--------|
| `NANOCHAT_BASE_DIR` | `common.get_base_dir` | Base directory for all artifacts (default `~/.cache/nanochat`). |
| `NANOCHAT_DTYPE` | `common._detect_compute_dtype` | Force compute dtype (`bfloat16`/`float16`/`float32`). |
| `RANK`, `LOCAL_RANK`, `WORLD_SIZE` | `common.get_dist_info` | Set by `torchrun`; control DDP. |
| `OMP_NUM_THREADS` | run scripts | Set to `1` to avoid CPU oversubscription. |
| `HF_TOKEN` | `dev/gen_synthetic_data.py` | HuggingFace API token for synthetic data generation. |
| `WANDB_RUN` | run scripts | Default wandb run name passed to `--run`. |
| `PYTORCH_ALLOC_CONF` | training scripts | Set to `expandable_segments:True` to reduce fragmentation. |

> [!NOTE]
> The training scripts set `PYTORCH_ALLOC_CONF="expandable_segments:True"` programmatically at import — you don't need to export it yourself.

### Model config (`GPTConfig`)

| Field | Default | Meaning |
|-------|---------|---------|
| `sequence_len` | `2048` | Max context length. |
| `vocab_size` | `32768` | Vocabulary size (padded to ×64 internally). |
| `n_layer` | `12` | Transformer depth. |
| `n_head` | `6` | Query heads. |
| `n_kv_head` | `6` | Key/value heads (GQA when `< n_head`). |
| `n_embd` | `768` | Model dimension. |
| `window_pattern` | `SSSL` | Sliding-window pattern (final layer forced to `L`). |

In practice these are computed from `--depth` in `base_train` rather than set directly.

For per-script CLI flags, see the User Guide's [Command & Flag Reference](karpathy__nanochat__UserGuide.md#13-command--flag-reference) and the per-script sections there.

---

## 13. Deploying the Chat Server

`scripts/chat_web.py` is a FastAPI app serving the model. Architecture:

- **`WorkerPool`** loads one full model replica per GPU (`--num-gpus`). Each `Worker` holds a device, an `Engine`, and a tokenizer. Available workers sit in an `asyncio.Queue`; requests acquire a worker (waiting if all busy) and release it after streaming completes.
- **Lifespan**: models load on startup; the server prints `Server ready at http://localhost:<port>`.
- **CORS** is wide open (`allow_origins=["*"]`).

### Endpoints

| Method | Path | Description |
|--------|------|-------------|
| `GET` | `/` | Serves `nanochat/ui.html` (with the API URL rewritten to same-origin). |
| `GET` | `/logo.svg` | The nanochat logo. |
| `POST` | `/chat/completions` | Streaming chat (Server-Sent Events). Accepts `messages`, optional `temperature`, `max_tokens`, `top_k`. |
| `GET` | `/health` | `{status, ready, num_gpus, available_workers}`. |
| `GET` | `/stats` | Worker pool stats (total/available/busy workers, per-worker device). |

### Built-in abuse limits

Enforced in `validate_chat_request`:

| Limit | Value |
|-------|-------|
| Max messages per request | 500 |
| Max characters per message | 8000 |
| Max total conversation characters | 32000 |
| Temperature range | 0.0–2.0 |
| Top-k range | 0–200 (0 = full vocab) |
| Max tokens range | 1–4096 |

Roles must be `user` or `assistant`. Streaming carefully accumulates tokens to avoid emitting partial multi-byte UTF-8 (emoji) sequences.

### Recommended production hardening

> [!CAUTION]
> Out of the box the server has **no authentication**, binds to `0.0.0.0`, and allows all CORS origins. Before exposing it beyond a trusted network:
> - Put it behind a reverse proxy (nginx/Caddy/Traefik) terminating TLS.
> - Add authentication at the proxy (basic auth, OAuth, or an API gateway).
> - Restrict CORS and bind to `127.0.0.1`, letting the proxy handle external traffic.
> - Add rate limiting at the proxy in addition to the built-in per-request caps.
> - Run inside a container with resource limits.

- [ ] TLS termination configured
- [ ] Authentication in front of `/chat/completions`
- [ ] CORS tightened to known origins
- [ ] Rate limiting at the edge
- [ ] Server bound to localhost behind the proxy
- [ ] Resource/GPU memory limits set

Multi-GPU serving scales throughput linearly with replicas (`--num-gpus`), since each request is handled by one worker end to end. Note that multi-GPU serving requires CUDA — CPU/MPS support only a single worker.

---

## 14. Distributed Training Operations

- **Launch**: `torchrun --standalone --nproc_per_node=<N> -m scripts.<name> -- <args>`. The `--standalone` flag runs a single-node rendezvous. `--nproc_per_node` = number of GPUs.
- **The `--` separator** divides torchrun args from script args and is mandatory.
- **Init/cleanup**: `compute_init()` sets seeds (42), enables TF32 matmul precision, initializes the NCCL process group, and pins each rank to its local GPU. `compute_cleanup()` destroys the group. Always paired around a run.
- **Rank-0 conventions**: only rank 0 logs (`print0`), saves model/meta files, and writes reports. All ranks save their own optimizer shard.
- **Gradient accumulation**: `grad_accum_steps = total_batch_size / (device_batch_size × max_seq_len × world_size)`. Single-GPU runs get a large accumulation count automatically.
- **Determinism**: full determinism is intentionally **not** enabled (`torch.use_deterministic_algorithms` is commented out to avoid slowdowns); global seeds are set but most code uses explicit RNG objects.

### Monitoring

The training loop logs per step: loss (debiased EMA), LR multiplier, step time, tokens/sec, and **MFU** (model FLOPS utilization, computed against the hardcoded GPU peak-FLOPS table in `common.py`). With wandb enabled (`--run=<name>`), watch `val/bpb`, `core_metric`, `train/mfu`, and `train/tok_per_sec`.

> [!NOTE]
> If MFU shows `0%`, your GPU isn't in the `get_peak_flops` table (it returns `inf` for unknowns). The metric is cosmetic; training is unaffected.

### Garbage-collection tuning

The training loops manually manage Python GC: after the first step they `gc.collect(); gc.freeze(); gc.disable()`, then `gc.collect()` every 5000 steps. This avoids ~500 ms GC pauses that otherwise recur during long runs.

---

## 15. Security Analysis

nanochat is research/training software, not a hardened production service. Key findings from the source:

### Code execution surfaces

1. **`execution.py` sandbox** — explicitly **not** a security boundary (see §7). Network is open, `ctypes` can escape, no kernel isolation. **Mitigation:** run model-code evaluation (HumanEval) inside a container or VM with no network.
2. **`use_calculator` in `engine.py`** — uses `eval()` but with `{"__builtins__": {}}`, a character allowlist, a dangerous-pattern denylist, disallowed `**`, a 3-second timeout, and acceptance of only arithmetic or `.count(...)`. This is a much tighter surface than the HumanEval sandbox, but it is still `eval`. **Mitigation:** keep the allowlist conservative; do not broaden it without review.

### Web server (`chat_web.py`)

| Concern | Status |
|---------|--------|
| Authentication | **None.** Anyone reaching the port can use the model. |
| CORS | `allow_origins=["*"]` — fully open. |
| Bind address | `0.0.0.0` by default. |
| Input validation | Present (message/length/temperature/token caps). |
| Rate limiting | Only the per-request size caps; no per-client rate limit. |
| Secrets | None handled by the server. |

**Mitigation:** front with an authenticating TLS proxy; see §13.

### Data & supply chain

- Pretraining data and the CORE eval bundle are downloaded over HTTPS from S3 and HuggingFace with retry/backoff; there is **no checksum verification** of downloaded shards or the eval bundle. **Mitigation:** if integrity matters, verify hashes out of band.
- `torch.load` is used to load checkpoints. **Only load checkpoints you trust** — pickle deserialization can execute code. nanochat loads only locally produced checkpoints by default.
- The FA3 kernel is fetched at runtime from the `kernels` hub (`varunneal/flash-attention-3`) on Hopper GPUs — a runtime code-download dependency to be aware of in locked-down environments.

### Secrets

The only secret-like value is `HF_TOKEN` (used by `dev/gen_synthetic_data.py`) and optional `wandb` credentials (via `wandb login`). Neither is stored in the repo.

> [!IMPORTANT]
> There is no malicious code in this repository. The security notes above describe **operational** risks of running training/research software (untrusted code execution, an unauthenticated dev server), not vulnerabilities in nanochat itself. The codebase is candid about these limits in its own docstrings.

---

## 16. Code Quality & Conventions

**Strengths.**
- **Radical readability.** Each major component (model, optimizer, training loop, engine) is one self-contained file you can read top to bottom. Heavy inline commentary explains *why*, not just *what* (e.g. the FP8 module's comparison to torchao, the scaling-law derivations in `base_train`).
- **Explicit over implicit.** No autocast, no hidden config objects, explicit param groups, explicit dtype management. The "one dial" philosophy keeps the surface tiny.
- **Honest comments.** The code flags its own approximations and footguns ("blindly following AdamW theory for Muon", "this is dumb, but we need to init the rotary embeddings", "not a true security sandbox").

**Conventions.**
- Scripts parse argparse at module top level and run a `while`/`for` loop; reusable logic is imported from `nanochat/`.
- Hyphenated CLI flags (`--device-batch-size`) map to underscored attrs.
- `print0` gates console output to rank 0; `DummyWandb` provides a no-op logging interface.
- Masking convention: `-1` is the cross-entropy ignore index for padding, prompts, and tool outputs throughout.
- The `d<depth>` model-tag convention threads through checkpoints and run scripts.

**Rough edges (acknowledged in-code).**
- Device/dtype tracking is partly hardcoded ("cuda → bf16, everything else → fp32") and the Engine inherits this assumption; the KVCache allocates eagerly. The author notes a larger refactor is warranted.
- `GPT.__init__` running on meta device is a documented footgun.
- Some logging-extraction in run scripts relies on `grep`-ing training log lines (fragile to format changes).

---

## 17. Testing Strategy

Tests live in `tests/` and run via `pytest` (config in `pyproject.toml`: `testpaths=["tests"]`, `slow` marker for slow tests, `python_files=["test_*.py"]`).

| Test file | Covers |
|-----------|--------|
| `tests/test_engine.py` | The inference `Engine` (KV-cache correctness, generation). |
| `tests/test_attention_fallback.py` | FA3 ↔ SDPA equivalence (using the `_override_impl` hook). |

The engine module also contains an inline `__main__` self-test that verifies `Engine.generate()` matches the naive `model.generate()` token-for-token. Coverage is intentionally light — this is a research codebase where correctness is largely validated by end-to-end metrics (CORE, bpb) rather than exhaustive unit tests.

> [!TIP]
> Run `pytest -m "not slow"` for a fast check, or the full suite before submitting changes. The attention-fallback test is especially valuable when modifying the model or attention path.

---

## 18. Potential Improvements

Drawn from in-code TODOs, community issues, and the documented rough edges:

| Area | Opportunity | Source |
|------|-------------|--------|
| Device/dtype | Refactor the hardcoded "cuda→bf16" assumption out of `Engine`/`KVCache`; allocate KV cache lazily. | `engine.py` comment |
| Attention | Adopt official PyTorch FA3 builds; extend fast attention to Ada/Blackwell. | Issue #737, `flash_attention.py` |
| Chunked inference | Fix token-smear for chunked prefill (first token of a chunk misses its cross-chunk predecessor). | Issue #756 |
| Flash Attention | Performance/error-handling improvements to the FA module. | Issue #752 |
| MPS | Add BF16 support on Apple Silicon. | Issue #685 |
| Memory leak | `base_eval` HellaSwag gets progressively slower / leaks memory on small models (Mac Studio). | Issue #427 |
| CORE parity | SQuAD task under-scores vs. reference (31% vs 37%) — investigate. | `core_eval.py` TODO |
| RoPE | Consider bumping the rotary base theta above 100K (more common recently). | `gpt.py` TODO |
| Model re-init | `build_model` calls `init_weights()` only to initialize rotary buffers — wasteful; refactor. | `checkpoint_manager.py` TODO |
| Run scripts | Replace fragile `grep`-based log parsing with structured metric output. | `runs/*.sh` |
| Legacy code | Remove the FinewebEdu→ClimbMix dataset fallback once migration is complete. | `dataset.py` |
| RL fp16 | SFT supports fp16 GradScaler; RL does not — add it for consistency. | README |

---

## 19. Contributing

The project's stated goal: improve the state of the art in **micro models** trainable end to end on budgets under $1000, optimizing for both cost and *cognitive* simplicity. Contributions must keep the codebase a single, cohesive, minimal, hackable "strong baseline" — no giant config objects, model factories, or if-then-else monsters.

- **Principled across depth.** Any change must work for *all* settings of `--depth`, since that single dial drives every derived hyperparameter. The reference quality target is improving the time-to-GPT-2 (CORE > 0.256525).
- **AI disclosure policy.** When submitting a PR, declare any parts with substantial LLM contribution that you did not write or do not fully understand.
- **Validate with the speedrun.** `runs/speedrun.sh` is the canonical recipe; the leaderboard (`dev/LEADERBOARD.md`) tracks wall-clock-to-GPT-2 on 8×H100.
- **License**: MIT. Cite as the `@misc{nanochat}` entry in the README if used in research.

Community channels: [DeepWiki](https://deepwiki.com/karpathy/nanochat), GitHub Discussions, and the `#nanochat` Discord.

---

## 20. Quick Reference

### Pipeline order

```text
nanochat.dataset  →  tok_train  →  tok_eval  →  base_train  →  base_eval
                                                     │
                                                     ▼
                                                 chat_sft  →  chat_eval  →  [chat_rl]  →  chat_cli / chat_web
```

### Model sources

| Source | Directory | Produced by |
|--------|-----------|-------------|
| `base` | `base_checkpoints/` | `base_train` |
| `sft` | `chatsft_checkpoints/` | `chat_sft` |
| `rl` | `chatrl_checkpoints/` | `chat_rl` |

### Key files

| Path | What |
|------|------|
| `nanochat/gpt.py` | The model. |
| `nanochat/optim.py` | MuonAdamW optimizer. |
| `nanochat/engine.py` | Inference + tool use. |
| `scripts/base_train.py` | Pretraining loop + scaling-law derivation. |
| `scripts/chat_web.py` | FastAPI chat server. |
| `runs/speedrun.sh` | Reference full pipeline. |
| `pyproject.toml` | Dependencies, Python/torch versions. |

### Critical knobs

| Knob | Where | Effect |
|------|-------|--------|
| `--depth` | training scripts | The one dial — sets model size and all derived hyperparameters. |
| `--device-batch-size` | training scripts | VRAM fit (reduce on OOM; results unchanged). |
| `--fp8` | `base_train` | Faster pretraining on H100+. |
| `NANOCHAT_DTYPE` | env | Force compute precision. |
| `NANOCHAT_BASE_DIR` | env | Relocate all artifacts. |
| `--num-gpus` | `chat_web` | Serving throughput (replicas). |

---

_Generated June 2026 from karpathy/nanochat at dc54a1a. 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._