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
2. Architecture Overview
3. The GPT Model (Internals)
4. The Optimizer: MuonAdamW
5. Precision: Explicit dtype & FP8
6. Attention: FA3 with SDPA Fallback
7. The Inference Engine & Tool Use
8. Data Pipeline
9. Checkpoints & the Base Directory
10. Module Reference
11. Requirements & Dependencies
12. Configuration Reference
13. Deploying the Chat Server
14. Distributed Training Operations
15. Security Analysis
16. Code Quality & Conventions
17. Testing Strategy
18. Potential Improvements
19. Contributing
20. Quick Reference

1. Executive Summary

nanochat covers the entire LLM lifecycle in one repository:

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

                         ┌───────────────────────────────────────────────────────────┐
                         │                      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):

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:

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):

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:

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:

$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:

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

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.

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):

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

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)

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

Built-in abuse limits

Enforced in validate_chat_request:

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)

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"]).

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:

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, GitHub Discussions, and the #nanochat Discord.

20. Quick Reference

Pipeline order

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

Model sources

Key files

Critical knobs