scratch-gpt

A from-scratch GPT built so the real mathematics is visible, in two layers:

1. Hand-written autograd (pure Python/NumPy, no torch.autograd) — a micrograd-style scalar Value engine with manually-derived backward() for + * ** exp log tanh relu softmax cross-entropy, used to train a tiny MLP. Correctness is proven by gradient-checking against numerical finite differences (max abs error < 1e-4).

2. A nanoGPT-style transformer in PyTorch — token embeddings, multi-head causal self-attention (softmax(QKᵀ/√d_k)·V with a causal mask), pre-norm transformer blocks with residuals, an MLP block, a tied linear head, cross-entropy loss, and AdamW. Trains char-level or with a from-scratch subword tokenizer on Tiny Shakespeare or your own text.

   Modern, selectable upgrades over the toy baseline:

   • BPE tokenizer (src/scratchgpt/tokenizer.py) — byte-level Byte-Pair Encoding implemented by hand (no tiktoken), so the model works with subword tokens like real GPTs instead of single characters.
   • RoPE rotary positional embeddings and RMSNorm (both LLaMA-style), selectable via GPTConfig(pos=..., norm=...).
   • Cosine LR schedule with warmup and gradient clipping in training, and nucleus (top-p) sampling alongside top-k.
   • An interactive REPL (scratchgpt.chat) that streams completions.

See MATH.md for the math at the kernel level, mapped to the code.

Install

Requires uv and Python 3.11+.

uv sync --extra dev

This installs numpy, torch, and pytest. CUDA is auto-detected; the code falls back to CPU when no GPU is present. The default model (n_layer=4, n_head=4, n_embd=128, block_size=128, batch_size=32, ~0.8M params) fits comfortably in 4GB of VRAM; on a CUDA OOM it auto-halves the batch size.

Test

uv run pytest -q

The acceptance suite covers the autograd gradient check, MLP learning (XOR), causal-attention shape/causality, GPT forward + loss, single-batch overfit (end-to-end backprop), tokenizers (char + BPE), an end-to-end train→checkpoint→reload integration test, and sampling.

Development

Dev tooling is configured in pyproject.toml. A Makefile wraps the common tasks; run the full quality gate (lint, format, types, tests + coverage) with:

make install      # uv sync --extra dev
make check        # ruff + ruff format --check + mypy + coverage gate

Individual targets: make lint, make format, make typecheck, make test, make cov. The coverage step (scripts/coverage_report.py) reruns the suite and regenerates COVERAGE.md plus an HTML drill-down in htmlcov/, and fails if total coverage drops below the fail_under threshold (85%). GitHub Actions (.github/workflows/ci.yml) runs the same gate on push/PR.

Train

Downloads Tiny Shakespeare to data/input.txt (synthesizes a small corpus if offline), trains, logs train/val loss, writes checkpoints/loss.csv and checkpoints/ckpt.pt, and prints a generated sample.

# Full-ish run (char-level Shakespeare, RoPE + RMSNorm)
uv run python -m scratchgpt.train

# Quick run (loss visibly decreases)
uv run python -m scratchgpt.train --max-iters 200

# Subword (BPE) tokenizer on your own corpus (a file or a folder of .txt)
uv run python -m scratchgpt.train --tokenizer bpe --vocab-size 512 \
    --data-path path/to/your_text.txt

Useful flags: --max-iters --batch-size --block-size --n-layer --n-head --n-embd --lr --min-lr --warmup-iters --grad-clip --weight-decay --tokenizer {char,bpe} --vocab-size --data-path --pos {rope,learned} --norm {rms,layer} --device {cuda,cpu}.

Sample

uv run python -m scratchgpt.sample --prompt "ROMEO:" --max-new-tokens 300

Flags: --ckpt --prompt --max-new-tokens --temperature --top-k --top-p --device.

Chat (interactive)

A REPL that loads the checkpoint once and streams a completion for whatever you type (it continues your text — it is not an instruction-following assistant).

uv run python -m scratchgpt.chat --temperature 0.7

In-REPL commands: /temp <f>, /tokens <n>, /help, /quit.

Layout

src/scratchgpt/
  autograd.py   # layer 1: hand-written reverse-mode autodiff + tiny MLP
  tokenizer.py  # char + from-scratch byte-level BPE tokenizers
  model.py      # layer 2: GPT (attention, RoPE, RMSNorm, blocks, generate)
  data.py       # corpus loading (file/dir/Shakespeare) + tokenization + batching
  train.py      # training loop (cosine LR + warmup, grad clip), checkpointing
  sample.py     # load checkpoint, generate text
  chat.py       # interactive streaming REPL
tests/          # acceptance tests
MATH.md         # the mathematics, mapped to the code