diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..ba4dce5 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,110 @@ +name: CI + + +on: + # PRs into main are the merge gate (TDR-014). Pushes to any branch also run, + # so breakage on Linux or an older Python is found while the work is in + # progress rather than at PR time -- this workflow sat unexecuted for its + # whole life because nothing had been pushed to main or opened as a PR. + push: + branches: ["**"] + pull_request: + branches: [main] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +env: + # Floor is 3.12 (see the comment on requires-python in pyproject.toml). + PYTHON_VERSION: "3.12" + +jobs: + lint: + name: Lint (ruff) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ env.PYTHON_VERSION }} + - name: Install ruff + run: pip install ruff + - name: Ruff check + run: ruff check . + - name: Ruff format (check only) + run: ruff format --check . + + typecheck: + name: Type check (mypy) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ env.PYTHON_VERSION }} + cache: pip + - name: Install package + type stubs + run: pip install -e ".[dev]" + - name: mypy + run: mypy tokenizer dataset model attention training inference optimization benchmark + + test: + name: Tests (pytest, py${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + # 3.14 is the version this project is developed on; 3.12 is the floor. + python-version: ["3.12", "3.13", "3.14"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - name: Install package + run: pip install -e ".[dev]" + - name: Run tests with coverage + run: pytest --cov=tokenizer --cov=dataset --cov=model --cov=attention --cov=training --cov=inference --cov=optimization --cov=benchmark --cov-report=term-missing + + import-smoke: + name: Import smoke + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ env.PYTHON_VERSION }} + cache: pip + - name: Install package + run: pip install -e . + - name: Modules import cleanly + run: | + python -c "import tokenizer, dataset, model, attention, training, inference, optimization, benchmark; print('all modules import OK')" + python -c "from tokenizer import BPE; t=BPE(vocab_size=256); ids=t.encode('hello'); assert t.decode(ids)=='hello'; print('tokenizer smoke OK')" + python -c " + import torch + from attention import CausalAttention + from dataset import DataEngine + e = DataEngine.from_config({'dataset': {'source': 'synthetic', 'seq_len': 16, 'batch_size': 2}}, vocab_size=64) + b = e.next_batch() + assert b['input_ids'].shape == (2, 16) + a = CausalAttention(d_model=32, n_heads=4, n_kv_heads=2, max_seq_len=64).eval() + out, cache = a(torch.randn(2, 16, 32), use_cache=True) + assert out.shape == (2, 16, 32) and len(cache) == 16 + print('dataset + attention smoke OK') + " + - name: Training entry point runs + # Bounded deliberately. This step is a wiring check, not a training run: + # the shipped config is 2000 steps of a 32.5M model, which on a CPU + # runner is ~1.8 minutes and writes ~1.8 GB of checkpoints. + run: | + python -m training.train --config configs/default.yaml \ + dataset.batch_size=2 dataset.seq_len=32 \ + training.steps=5 training.warmup_steps=1 \ + training.ckpt_every=0 training.log_every=1 \ + training.out_dir="${RUNNER_TEMP:-/tmp}/ci-checkpoints" diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f612cdd --- /dev/null +++ b/.gitignore @@ -0,0 +1,125 @@ +# --------------------------------------------------------------------------- +# Python +# --------------------------------------------------------------------------- +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# --------------------------------------------------------------------------- +# Testing / Linting / Type-checking caches +# --------------------------------------------------------------------------- +.mypy_cache/ +.pytest_cache/ +.ruff_cache/ +.coverage +.coverage.* +htmlcov/ +.tox/ +.nox/ +.hypothesis/ +nosetests.xml +coverage.xml + +# --------------------------------------------------------------------------- +# Virtual environments / package managers +# --------------------------------------------------------------------------- +.venv/ +venv/ +env/ +ENV/ +env.bak/ +venv.bak/ +.conda/ +.python-version +poetry.lock.persona +uv.lock.persona + +# --------------------------------------------------------------------------- +# Model checkpoints / training artifacts (large binaries, never commit) +# --------------------------------------------------------------------------- +checkpoints/* +!checkpoints/.gitkeep +runs/ +logs/ +wandb/ +*.pt +*.pth +*.ckpt +*.safetensors +*.bin +*.onnx + +# --------------------------------------------------------------------------- +# Data files / datasets (large, regenerable, or licensed) +# --------------------------------------------------------------------------- +dataset/data/ +data/ +*.npy +*.npz +*.memmap +*.h5 +*.hdf5 +*.parquet +*.arrow + +# --------------------------------------------------------------------------- +# Raw / source text dumps +# --------------------------------------------------------------------------- +arxiv.txt + +# --------------------------------------------------------------------------- +# Secrets & environment files (NEVER commit API keys / tokens) +# --------------------------------------------------------------------------- +.env +.env.* +*.env +.env.local +.env.production +secrets/ +credentials/ +*.pem +*.key +*.p12 +config/api_keys*.yaml + +# --------------------------------------------------------------------------- +# OS / editor / editor-config files +# --------------------------------------------------------------------------- +.DS_Store +.DS_Store? +._* +.Spotlight-V100 +.Trashes/ +ehthumbs.db +Thumbs.db +.idea/ +.vscode/ +*.swp +*.swo +*~ +*.orig +*.rej +.vscode/ + +# --------------------------------------------------------------------------- +# Notebooks / scratch experiments (work-in-progress, not part of the package) +# --------------------------------------------------------------------------- +test_agent.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..ba2ac1c --- /dev/null +++ b/README.md @@ -0,0 +1,624 @@ +
+

Language Model Engine ๐Ÿง โšก

+

A decoder-only GPT-style language model built entirely from scratch in PyTorch โ€” custom byte-level BPE tokenizer, memory-mapped data engine, grouped-query attention with RoPE, a hand-written tiled attention kernel, a hand-written AdamW, and cached autoregressive generation. No pretrained weights, no HuggingFace.

+
+ + + + + + + +
+
+ +--- + +### Project Overview + +*This project builds a complete language model system from first principles, with the emphasis on low-level systems engineering rather than on training a model quickly.* + +- **Everything is hand-written:** BPE merges, rotary embeddings, grouped-query attention, the KV cache, a tiled online-softmax attention kernel *with its own recomputing backward*, and AdamW are all implemented directly โ€” the only thing borrowed from PyTorch is tensor math. +- **Three interchangeable attention kernels:** a readable `manual` reference, torch's fused `sdpa`, and our own `flash` tiled implementation. All three are asserted to produce identical output, so the reference acts as a correctness oracle for the optimised paths. +- **Memory-mapped data engine:** the corpus is a bare `uint16` array on disk read through `mmap`, so window count comes from `stat()` alone and no bulk copy ever enters RAM โ€” even with multiprocessing workers. +- **GQA that actually saves memory:** 8 query heads share 2 KV heads, shrinking the KV cache **4ร—** (2048 KiB โ†’ 512 KiB per layer at 512 tokens) and the attention parameters by 37.5%. +- **Test-driven throughout:** every module's exit test is written before its implementation. **359 tests** currently pass across unit and integration suites. +- **Runs on Apple Silicon:** device resolution is `mps โ†’ cuda โ†’ cpu`, with bf16 autocast on MPS and no `GradScaler` (which is CUDA-only). + +--- + +### System Architecture + +```mermaid +flowchart TB + RAW["๐Ÿ“„ Raw text corpus"] + + subgraph BUILT ["Implemented โ€” 359 tests"] + TOK["tokenizer/
byte-level BPE
train ยท encode ยท decode"] + DS["dataset/
mmap + streaming
windowing ยท sharding"] + ATT["attention/
GQA + RoPE + KV cache
manual ยท sdpa ยท flash"] + MODEL["model/
RMSNorm ยท SwiGLU
CausalLM"] + TRAIN["training/
AdamW ยท cosine ยท AMP
grad accum ยท checkpoints"] + INFER["inference/
greedy ยท top-k ยท top-p
cached generation"] + OPT["optimization/
int8 / int4 quant
KV arena"] + BENCH["benchmark/
latency ยท throughput ยท memory"] + end + + CKPT[("๐Ÿ’พ checkpoints/")] + + RAW --> TOK --> DS --> TRAIN + ATT --> MODEL --> TRAIN + TRAIN --> CKPT --> INFER + CKPT --> OPT --> INFER + BENCH -.measures.-> DS + BENCH -.measures.-> ATT + BENCH -.measures.-> TRAIN + BENCH -.measures.-> INFER + + classDef done fill:#1b5e20,stroke:#4CAF50,stroke-width:2px,color:#fff + classDef todo fill:#37474f,stroke:#78909c,stroke-width:1px,color:#cfd8dc,stroke-dasharray: 4 4 + classDef store fill:#4a148c,stroke:#ab47bc,color:#fff + class TOK,DS,ATT,MODEL,TRAIN,INFER,OPT,BENCH done + class CKPT,RAW store +``` + +| Module | Responsibility | State | +|---|---|---| +| `tokenizer/` | Byte-level BPE, GPT-2 pre-tokenization, piece cache, incremental training index | โœ… 54 tests | +| `dataset/` | `uint16` corpus format, mmap/streaming readers, corpus download, document split, message prep | โœ… 76 tests | +| `attention/` | RoPE, grouped-query causal attention, KV cache, three kernels + recomputing backward | โœ… 43 tests | +| `model/` | RMSNorm, SwiGLU, pre-norm blocks, `CausalLM` with weight tying and depth-scaled init | โœ… 25 tests | +| `training/` | Hand-written AdamW, cosine schedule, AMP, grad accumulation, checkpoints, validation | โœ… 58 tests | +| `inference/` | Greedy / top-k / top-p sampling, cached generation, batched generation with key padding | โœ… 29 tests | +| `optimization/` | int8 / int4 per-channel weight quantization, preallocated KV arena | โœ… 38 tests | +| `benchmark/` | Device-synced timing harness, p50/p90/p95/p99, exact memory metrics, baseline report | โœ… 25 tests | + +--- + +### Data Pipeline + +```mermaid +flowchart LR + A["๐Ÿ“„ raw text"] --> B["BPE.encode()"] + B --> C["append <|endoftext|>
per document"] + C --> D[("corpus.bin
flat uint16, no header")] + + D --> E["stat() โ†’ n_tokens
window count without opening"] + D --> F["np.memmap
lazy, per-process handle"] + + F --> G["window i โ†’ tokens[iยทT : iยทT+T+1]"] + G --> H["input_ids = w[:-1]
labels = w[1:]"] + + H --> I["DataLoader
workers ยท prefetch ยท pin_memory"] + I --> J["batch [B, T]"] + J --> K["๐Ÿ–ฅ๏ธ .to(device)"] + + classDef disk fill:#4a148c,stroke:#ab47bc,color:#fff + classDef proc fill:#0d47a1,stroke:#42a5f5,color:#fff + classDef out fill:#1b5e20,stroke:#4CAF50,color:#fff + class D disk + class B,C,E,F,G,H,I proc + class J,K out +``` + +Each window reads **`T + 1`** ids and splits them, rather than reading `T` and shifting inside the window. That one extra token means the final position gets a *real* label instead of padding โ€” shifting inside would silently discard 1/512 of the training signal and teach one position from a fake target. + +--- + +### Attention Architecture + +```mermaid +flowchart TB + X["x โ€” [B, T, d_model]"] + + X --> QP["q_proj โ†’ [B, 8, T, 64]"] + X --> KP["k_proj โ†’ [B, 2, T, 64]"] + X --> VP["v_proj โ†’ [B, 2, T, 64]"] + + QP --> RQ["RoPE rotate
offset = cache length"] + KP --> RK["RoPE rotate
offset = cache length"] + VP --> NOV["never rotated
RoPE encodes where, not what"] + + RK --> CACHE[("KV cache
stores 2 heads, not 8")] + NOV --> CACHE + + CACHE --> REP["repeat_kv ร—4
query head h โ†’ kv head hรท4"] + + RQ --> KERNEL + REP --> KERNEL + + subgraph KERNEL ["kernel โ€” all three must agree"] + M["manual
oracle"] + S["sdpa
fused, default"] + FL["flash
tiled, O(T) memory"] + end + + KERNEL --> O["o_proj โ†’ [B, T, d_model]"] + + classDef gqa fill:#e65100,stroke:#ffb74d,color:#fff + classDef cache fill:#4a148c,stroke:#ab47bc,color:#fff + classDef kern fill:#0d47a1,stroke:#42a5f5,color:#fff + class KP,VP,REP gqa + class CACHE cache + class M,S,FL kern +``` + +#### KV cache: prefill then decode + +```mermaid +sequenceDiagram + autonumber + participant U as Caller + participant A as CausalAttention + participant C as KVCache + + Note over U,C: Prefill โ€” the whole prompt in one parallel pass + U->>A: forward(prompt[1, 6, C], use_cache=True) + A->>A: RoPE at positions 0..5 + A->>C: update(k[1,2,6,64], v[1,2,6,64]) + C-->>A: full K, V (len 6) + A-->>U: out[1, 6, C], cache(len=6) + + Note over U,C: Decode โ€” one token at a time, O(T) not O(Tยฒ) + U->>A: forward(token[1, 1, C], kv_cache, use_cache=True) + A->>A: RoPE at position 6 โ€” from cache length, not slice index + A->>C: update(k[1,2,1,64], v[1,2,1,64]) + C-->>A: full K, V (len 7) + Note right of A: q_len == 1 โ†’ no mask needed
one query legally sees every key + A-->>U: out[1, 1, C], cache(len=7) +``` + +The RoPE offset in step 7 is the single easiest bug to ship in this module: rotating a decode token at position 0 instead of `cache_len` leaves prefill looking perfect while every generated token believes it is the start of the sequence. `test_cached_decode_matches_full_forward` pins it by asserting that token-by-token decoding reproduces one full parallel forward exactly. + +#### Tiled ("flash") attention + +```mermaid +flowchart TB + START(["q, k, v โ€” [B, H, T, D]"]) --> QT["for each query tile
block_q rows"] + QT --> INIT["m = -inf ยท l = 0 ยท acc = 0
running max, denominator, output"] + INIT --> KT{"next key tile?"} + + KT -->|"j > last query pos"| DONE["break โ€” every later tile
is fully masked"] + KT -->|yes| SC["scores = q_tile ยท k_tileแต€ ยท scale
only [block_q ร— block_k] resident"] + SC --> MASK["mask keys > query position"] + MASK --> NEWM["m_new = max(m, rowmax)
clamp -inf rows โ†’ 0, else NaN"] + NEWM --> RESCALE["rescale = exp(m โˆ’ m_new)
l = lยทrescale + ฮฃexp(s โˆ’ m_new)
acc = accยทrescale + exp(sโˆ’m_new)ยทv_tile"] + RESCALE --> KT + + DONE --> OUT["out_tile = acc / l"] + OUT --> NEXT{"more query tiles?"} + NEXT -->|yes| QT + NEXT -->|no| FIN(["out โ€” [B, H, T, D]"]) + + classDef hot fill:#e65100,stroke:#ffb74d,color:#fff + classDef safe fill:#1b5e20,stroke:#4CAF50,color:#fff + class SC,RESCALE hot + class DONE,NEWM safe +``` + +The full `T ร— T` score matrix is never allocated โ€” peak memory is O(T) rather than O(Tยฒ). This is asserted rather than claimed: `test_flash_never_materialises_the_full_score_matrix` monkey-patches `torch.matmul` and fails if any intermediate exceeds the tile width. + +--- + +### Training on Real Data + +```bash +python -m dataset.prepare_corpus --corpus tinyshakespeare --vocab-size 8192 +python -m training.train --config configs/tinyshakespeare.yaml +``` + +`prepare_corpus` downloads, fits a vocabulary on a bounded sample, splits **by +document**, and writes `train.bin` / `val.bin` beside the `vocab.json` they were +encoded with. Training then reports held-out perplexity every `eval_every` steps. + +Two properties worth knowing: + +- **The split is by document, never by token.** Windows are contiguous slices of + one flat array, so cutting a *tokenised* stream puts validation tokens inside + training windows and the reported perplexity comes out quietly too good. + Documents are assigned by a hash of their own text, so the split is + reproducible without storing a seed, and adding documents never reshuffles + existing ones across the boundary. +- **Corpora are never committed.** `data/` is gitignored, and every entry records + its licence, because "where did the training data come from" is unpleasant to + answer retroactively. + +Behind a TLS-inspecting corporate proxy (Zscaler and similar), `requests` fails +where `curl` succeeds โ€” the proxy re-signs certificates with a corporate root CA +that certifi does not carry. `pip install -e ".[data]"` pulls in `truststore`, +which routes verification through the OS keychain. Verification is never +disabled. + +#### Email and chat corpora + +`dataset/messages.py` prepares mail and message records โ€” stripping quoted +replies, signatures and confidentiality footers, optionally adding +`<|from|>` / `<|subject|>` / `<|body|>` structure tokens, and masking recognisable +identifier shapes. It ships **no connector to any mailbox or workspace**; +the caller supplies the records. + +Read the module docstring before pointing it at real mail. Three things are true +and none of them are obvious: + +1. **A model memorises its training data.** This project measured a 32.5M-param + model reaching loss 0.017 on a 10,677-token corpus โ€” reproducing it. A + checkpoint trained on mail *is* a copy of that mail and inherits its handling + obligations. +2. **`redact()` is best-effort, not a compliance control.** It catches shapes + (addresses, phone numbers, long digit runs). It cannot catch a name in prose + or a deal codename โ€” `test_redaction_cannot_catch_a_name_in_prose` pins that + limit deliberately. +3. **A mailbox is a fine-tuning corpus, not a pretraining one.** ~1โ€“10M tokens + against the ~650M this model size wants; trained alone it memorises. + +--- + +### Optimization Strategy + +| Concept | Implementation | Measured effect | +|---|---|---| +| **Grouped-Query Attention** | 8 query heads share 2 KV heads; the cache stores the *unexpanded* heads and `repeat_kv` expands only on read | KV cache **2048 KiB โ†’ 512 KiB** per layer at 512 tokens; attention params 1,048,576 โ†’ 655,360 | +| **KV cache** | Prefill computes all K/V once, then each decode step appends a single position | Generation drops from O(Tยฒ) recomputation to O(T) | +| **Tiled attention** | Block-wise online softmax with running max and denominator; fully-masked tiles are skipped entirely | Score memory O(Tยฒ) โ†’ O(block_q ร— block_k) | +| **Memory-mapped corpus** | Bare `uint16` array; `np.memmap` opened lazily *per process* and dropped on pickle | No bulk copy into RAM; each DataLoader worker gets its own handle instead of the corpus through a pipe | +| **RoPE as buffers** | cos/sin tables registered with `persistent=False` | Derived constants stay out of every checkpoint | +| **Header-free token format** | File is exactly `n_tokens ร— 2` bytes | Window count from `stat()` โ€” no open, no parse, no read | +| **Per-channel weight quantization** | Symmetric int8/int4 with one scale per output row; int4 packs two codes per byte | 128.03 MiB โ†’ 41.21 MiB (int8) / 26.71 MiB (int4), perplexity within 10% | +| **Preallocated KV arena** | One `[B, H, max_seq_len, D]` allocation written in place, replacing per-step concatenation | Zero allocations per decode step; reserved-vs-used bytes reported separately | + +--- + +### Pipeline Workflow + +**(1)** Tokenization: +- Text is encoded to UTF-8 bytes, so there are no unknown tokens for any input โ€” emoji and rare scripts included. +- BPE merges are learned greedily, with deterministic tie-breaking so training is reproducible. +- Special tokens occupy the lowest ids and are atomic โ€” never merged across. + +**(2)** Corpus packing: +- Documents are encoded and joined with `<|endoftext|>` so windows cannot silently splice one document onto the next. +- Ids are validated against the `uint16` range and written as a flat array. An out-of-range id raises rather than wrapping (`65536 โ†’ 0`), which would corrupt the corpus in a way that only surfaces much later as inexplicable loss. + +**(3)** Windowing: +- `stat()` gives the token count; `(n_tokens โˆ’ 1) // seq_len` gives the window count. +- Each window reads `seq_len + 1` ids and splits into `input_ids` / `labels`. +- Workers shard strided (`order[worker_id::num_workers]`) so no window is emitted twice. + +**(4)** Attention: +- Q/K/V are projected, split into heads, and Q/K rotated by RoPE at their absolute positions. +- KV heads are cached pre-expansion, then expanded to full head count on read. +- The kernel masks according to the `q_len` / `k_len` relationship โ€” full triangle, no mask, or an explicit bottom-right triangle for chunked prefill. + +--- + +### Setting up the project in your machine + +#### Prerequisites + +- Python 3.12+ +- pip +- Git + +#### Clone the repository + +```bash +git clone https://github.com/555aaditya/Language-Model.git +cd Language-Model +``` + +#### Install dependencies + +```bash +python3 -m venv .venv +source .venv/bin/activate +pip install -e ".[dev]" +``` + +#### Verify the install + +```bash +pytest +``` + +*Expected output:* +``` +359 passed +``` + +#### Run the entry point + +```bash +python -m training.train --config configs/default.yaml +``` + +*Expected output (on Apple Silicon):* +``` +device: mps | amp dtype: torch.bfloat16 +dataset engine OK: input_ids (16, 512), labels (16, 512) on mps:0 +``` + +Any config key can be overridden from the command line with dotted notation: + +```bash +python -m training.train --config configs/default.yaml training.lr=3e-4 model.n_layers=12 attention.impl=flash +``` + +--- + +### Testing + +```bash +pytest # everything (359 tests) +pytest tests/unit -q # unit only +pytest tests/integration -q # cross-module seams +pytest tests/unit/test_attention.py -q # one module +``` + +Reproduce the CI gate locally: + +```bash +ruff check . +ruff format --check . +mypy tokenizer dataset model attention training inference optimization benchmark +pytest --cov=tokenizer --cov=dataset --cov=attention --cov-report=term-missing +``` + +| Suite | Tests | What it pins | +|---|---|---| +| `test_tokenizer.py` | 14 | Encode/decode round trips, unicode, merge ordering, special tokens, serialisation | +| `test_dataset.py` | 29 | `uint16` bounds, window arithmetic, next-token shift, worker sharding, DataLoader trap avoidance | +| `test_attention.py` | 39 | Causal masking, RoPE relative-distance property, GQA group mapping, KV-cache equivalence, kernel agreement | +| `test_model.py` | 25 | RMSNorm vs LayerNorm behaviour, SwiGLU gating, pre-norm residuals, parameter count, weight tying, cached decode | +| `test_optimizer.py` | 20 | AdamW matching `torch.optim.AdamW` step for step, decoupled decay, no-decay groups, cosine schedule shape | +| `test_trainer.py` | 18 | Overfit-one-batch, grad accumulation equivalence, clipping, LR schedule, checkpoint resume | +| `test_train_config.py` | 20 | YAML loading, dotted-key overrides, type preservation across overrides | +| `test_inference.py` | 29 | Greedy determinism, top-k/top-p set selection, cross-device seeded sampling, cached-vs-uncached generation | +| `test_tokenizer_dataset_attention.py` | 5 | Cross-module seams: tokenizer โ†’ `.bin` โ†’ batches โ†’ attention | +| `test_quantization.py` | 21 | int8/int4 round-trip error bounds, per-channel scales, real byte savings, perplexity tolerance | +| `test_kv_pool.py` | 17 | Arena never reallocates, returns only the filled prefix, generates identically to the concat cache | +| `test_benchmark.py` | 25 | Device sync, interpolated percentiles, undersampling detection, exact memory arithmetic, baseline reproduction | +| `test_pretokenize.py` | 40 | Pattern tiles every input, no cross-piece merges, cache soundness, sub-quadratic scaling | +| `test_messages.py` | 22 | Quoted-reply/signature/disclaimer stripping, redaction limits, structure tokens | +| `test_corpora_and_eval.py` | 25 | Document-split disjointness and stability, validation perplexity, batched generation | +| `test_end_to_end.py` | 5 | Full pipeline learns real text; trained model's cache stays equivalent; resume converges | + +--- + +### A Measured Run + +Training this repo's own `docs/TDD.md` as a corpus, on an M5 Pro via MPS: + +| | | +|---|---| +| Corpus | 39,922 chars โ†’ **10,677 tokens** (3.74 chars/token) | +| Tokenizer | 2,048 vocab, 1,791 merges, trained in 4.0 s | +| Model | 4.33M params โ€” `d_model` 256, 4 layers, 8 heads / 2 KV heads | +| Throughput | **~75,000 tokens/sec** (bf16 autocast) | +| 250 steps | 6.9 s wall clock | +| Loss | 7.679 โ†’ 0.017 (baseline `ln(2048)` = 7.625) | + +**Read that loss honestly:** 4.33M parameters against 10,677 tokens is a ~400:1 +overparameterisation, so 0.017 is memorisation, not language modelling. The +generations show it โ€” correct markdown structure, real vocabulary from the +document, no coherent syntax. What the run demonstrates is that the *pipeline* +is sound end to end: the tokenizer round-trips, the corpus packs, gradients +flow, the schedule fires, and the cache generates. Producing a model worth +evaluating needs a corpus several orders of magnitude larger, which is what the +`benchmark/` stage exists to measure. + +--- + +### Benchmark Results + +`python -m benchmark.report --config configs/default.yaml` + +The report deliberately splits into two halves, because they deserve different +levels of trust. **Deterministic** metrics are exact arithmetic over tensor +shapes โ€” identical on every machine, and asserted against a checked-in baseline +in CI. **Measured** metrics are wall-clock and reproducible only in shape. + +#### Deterministic โ€” 32.51M parameters, `configs/default.yaml` + +| Precision | Size | Compression | +|---|---|---| +| fp32 | 128.03 MiB | 1.00ร— | +| fp16 | 64.02 MiB | 2.00ร— | +| int8 | 41.21 MiB | **3.11ร—** | +| int4 | 26.71 MiB | **4.79ร—** | + +int8 lands at 3.11ร—, not the 4ร— the bit-width suggests โ€” because the token +embedding is tied to the LM head and is deliberately **not** quantised (see the +design table below). At 2.1M parameters it stays fp32 and sets a floor on how +small the model can get. Quoting 4ร— here would be arithmetic that the code +doesn't do. + +| KV cache @ 1024 tokens | Size | +|---|---| +| MHA (8 KV heads) | 32.0 MiB | +| GQA (2 KV heads) | **8.0 MiB** โ€” 4.00ร— smaller | + +#### Measured โ€” Apple M5 Pro, MPS, torch 2.13, bf16 + +| Phase | runs | p50 | p90 | p95 | p99 | max | stdev | +|---|---|---|---|---|---|---|---| +| Prefill (32 tok) | 100 | 2.80 ms | 2.88 ms | 2.93 ms | 2.99 ms | 3.00 ms | 0.05 ms | +| Decode (32 tok) | 20 | 77.86 ms | 78.21 ms | 78.39 ms | 78.73 ms | 78.81 ms | 0.35 ms | + +Throughput: **8,026 tok/s** prefill, **411 tok/s** decode (2.43 ms/token), +**~75,000 tok/s** training. + +Sample counts are set per phase rather than globally, because a percentile is +only as good as the samples behind it: with `n` runs the finest resolvable tail +probability is `1/n`, so a p99 from 10 runs is just the maximum wearing a +different label. Every timing therefore carries `resolvable_percentile` +alongside it โ€” p99.0 for prefill, p95.0 for the slower decode loop. + +KV cache speedup, measured rather than assumed (TDR-012): + +| New tokens | Cached | Uncached | Speedup | +|---|---|---|---| +| 16 | 40.6 ms | 51.8 ms | 1.27ร— | +| 64 | 154.3 ms | 208.1 ms | 1.35ร— | +| 256 | 620.2 ms | 1150.0 ms | **1.85ร—** | + +The speedup grows with sequence length, which is the O(Tยฒ) โ†’ O(T) claim showing +up. It is also **modest in absolute terms** โ€” at 32M parameters, per-step Python +dispatch and MPS kernel-launch overhead are a large share of each decode step, so +the asymptotic win is partly masked by constant factors. The benchmark reports +the ratio whichever way it falls, including the short-sequence regime where the +cache can genuinely lose. + +--- + +### CI / Infrastructure + +```mermaid +flowchart LR + DEV["๐Ÿ‘ฉโ€๐Ÿ’ป feature branch"] --> PR["Pull Request โ†’ main"] + PR --> CI{"GitHub Actions
free tier"} + + CI --> L["Lint
ruff check
ruff format --check"] + CI --> T["Type check
mypy ยท 8 packages"] + CI --> P["Tests
pytest + coverage
py3.12 ยท 3.13 ยท 3.14"] + CI --> S["Import smoke
modules import
+ end-to-end pipeline"] + + L --> GATE{"all green?"} + T --> GATE + P --> GATE + S --> GATE + + GATE -->|no| BLOCK["๐Ÿšซ merge blocked"] + GATE -->|yes| MERGE["โœ… merge to main"] + MERGE --> LOCKED[("๐Ÿ”’ main
always stable")] + + classDef job fill:#0d47a1,stroke:#42a5f5,color:#fff + classDef good fill:#1b5e20,stroke:#4CAF50,color:#fff + classDef bad fill:#b71c1c,stroke:#ef5350,color:#fff + class L,T,P,S job + class MERGE,LOCKED good + class BLOCK bad +``` + +`main` is locked (TDR-013): all work happens on feature branches and merges only through a reviewed PR with every CI job green. + +--- + +### Build Order + +```mermaid +flowchart LR + S0["0 ยท tokenizer
BPE"] --> S1["1 ยท dataset
mmap engine"] + S0 --> S2["2 ยท attention
GQA + RoPE"] + S2 --> S3["3 ยท model
CausalLM"] + S1 --> S4["4 ยท training
AdamW + cosine"] + S3 --> S4 + S4 --> S5["5 ยท inference
sampling"] + S4 --> S6["6 ยท optimization
quantization"] + S6 --> S5 + S5 --> S7["7 ยท benchmark
metrics harness"] + + classDef done fill:#1b5e20,stroke:#4CAF50,stroke-width:2px,color:#fff + classDef next fill:#e65100,stroke:#ffb74d,stroke-width:2px,color:#fff + classDef todo fill:#37474f,stroke:#78909c,color:#cfd8dc,stroke-dasharray: 4 4 + class S0,S1,S2,S3,S4,S5,S6,S7 done +``` + +Each stage has an exit test written *before* the implementation. See [`docs/BUILD_ORDER.md`](docs/BUILD_ORDER.md) for the stable interface contracts. + +--- + +### Tools and Technologies + +| Tool / Technology | Purpose | +|---|---| +| **PyTorch 2.13** | Autograd, tensor math, and device management โ€” every layer above that is hand-written | +| **NumPy** | `memmap` corpus reader and the `uint16` on-disk token format | +| **Apple MPS** | Primary development accelerator; bf16 autocast, fused `scaled_dot_product_attention` | +| **pytest** | 359-test TDD suite โ€” exit tests written before each module | +| **ruff** | Linting and formatting, 100-char lines, `E/F/I/W/UP/B` rule set | +| **mypy** | Static type checking across all eight packages | +| **PyYAML** | Config format, reused as the CLI override parser so types can never diverge | +| **GitHub Actions** | Free-tier CI: lint, type check, tests on three Pythons, import smoke | + +--- + +### Key Design Decisions + +Full rationale and alternatives for each is recorded as a Technical Decision Record in [`docs/TDD.md` ยง13](docs/TDD.md). + +| Decision | Rationale | +|---|---| +| **Byte-level BPE** (TDR-002) | Zero out-of-vocabulary tokens for arbitrary UTF-8. A GPT-2 style `bytes_to_unicode` table renders every byte to a printable character, so token keys never collide with special-token text. | +| **Rotary position embeddings** (TDR-015) | Attention scores become a function of *relative* offset, removing a `max_seq_len ร— d_model` parameter table entirely. Supersedes the learnable absolute table originally specified in TDR-006, which was never implemented. | +| **Grouped-query attention** (TDR-016) | The KV cache โ€” not the weights โ€” dominates decode memory, and it scales with `n_kv_heads`. A 4:1 group ratio cuts the cache 4ร— for negligible quality cost, which matters directly on unified memory. | +| **RMSNorm + SwiGLU** (TDR-017) | RMSNorm drops mean-centring and bias for one fewer reduction pass; SwiGLU's multiplicative gate beats a plain GELU MLP at equal parameter count. | +| **Loss lives in the trainer, not the model** | `CausalLM.forward` returns logits only. Computing a loss during inference is wasted work, and keeping the model pure makes the generation path obvious. | +| **Dict batches, not tuples** | Datasets yield `{"input_ids", "labels"}` so a field can be added later without breaking every caller โ€” and so a transposed `(x, y)` can't silently typecheck. | +| **`uint16` token format** | Caps the vocabulary at 65,536, comfortably above GPT-2's 50,257. Out-of-range ids raise instead of widening the format and doubling the page-cache footprint. | +| **Three attention kernels, asserted equal** (TDR-007) | Optimised kernels are only trustworthy against a reference. `manual` is the oracle; `sdpa` and `flash` must match it within floating-point tolerance or the build fails. | +| **Python 3.12 floor** | numpy โ‰ฅ 2.3 ships PEP 695 syntax in its stubs that mypy rejects when targeting older versions. Rather than typecheck against a Python we never run, runtime, ruff, mypy and the CI matrix all share one floor. | +| **MPS as the primary target** (TDR-019) | It is the hardware this is actually built on, so it is the only device whose numbers can be reported honestly. Benchmarks are always reported per-device, never as a general claim. | +| **Locked `main`, PR-gated** (TDR-013) | Every state of the repository stays auditable and revertible; CI green is a merge requirement. | + +--- + +### Repository Layout + +``` +Language-Model/ +โ”œโ”€โ”€ tokenizer/ # byte-level BPE โœ… +โ”‚ โ”œโ”€โ”€ bpe.py +โ”‚ โ””โ”€โ”€ pretokenize.py # GPT-2 style splitter (TDR-020) +โ”œโ”€โ”€ dataset/ # data engine โœ… +โ”‚ โ”œโ”€โ”€ binfile.py # uint16 corpus format +โ”‚ โ”œโ”€โ”€ token_dataset.py # mmap / streaming / synthetic readers +โ”‚ โ”œโ”€โ”€ engine.py # DataEngine โ†’ endless device-ready batches +โ”‚ โ”œโ”€โ”€ corpora.py # download + document-level train/val split +โ”‚ โ”œโ”€โ”€ messages.py # email / chat cleaning and redaction +โ”‚ โ””โ”€โ”€ prepare_corpus.py # preparation CLI +โ”œโ”€โ”€ attention/ # attention โœ… +โ”‚ โ”œโ”€โ”€ rope.py # rotary position embeddings +โ”‚ โ”œโ”€โ”€ kv_cache.py # per-layer K/V cache +โ”‚ โ”œโ”€โ”€ kernels.py # manual ยท sdpa ยท flash +โ”‚ โ””โ”€โ”€ causal.py # CausalAttention (GQA) +โ”œโ”€โ”€ model/ # transformer โœ… +โ”‚ โ”œโ”€โ”€ norm.py # RMSNorm +โ”‚ โ”œโ”€โ”€ ffn.py # SwiGLU +โ”‚ โ”œโ”€โ”€ block.py # pre-norm TransformerBlock +โ”‚ โ””โ”€โ”€ causal_lm.py # CausalLM +โ”œโ”€โ”€ training/ # training engine โœ… +โ”‚ โ”œโ”€โ”€ optimizer.py # hand-written AdamW + param groups +โ”‚ โ”œโ”€โ”€ scheduler.py # cosine decay w/ linear warmup +โ”‚ โ”œโ”€โ”€ trainer.py # AMP, grad accum, clipping +โ”‚ โ”œโ”€โ”€ checkpoint.py # save / resume +โ”‚ โ”œโ”€โ”€ device.py # mps โ†’ cuda โ†’ cpu resolution +โ”‚ โ””โ”€โ”€ train.py # config, overrides, CLI +โ”œโ”€โ”€ inference/ # generation โœ… +โ”‚ โ”œโ”€โ”€ sampling.py # temperature ยท top-k ยท top-p +โ”‚ โ”œโ”€โ”€ generate.py # cached autoregressive loop +โ”‚ โ””โ”€โ”€ batched.py # multi-request generation +โ”œโ”€โ”€ optimization/ # post-training opt โœ… +โ”‚ โ”œโ”€โ”€ quantize.py # int8 / int4 per-channel weights +โ”‚ โ””โ”€โ”€ kv_pool.py # preallocated KV arena +โ”œโ”€โ”€ benchmark/ # metrics harness โœ… +โ”‚ โ”œโ”€โ”€ harness.py # device-synced timing + percentiles +โ”‚ โ”œโ”€โ”€ memory.py # exact footprint arithmetic +โ”‚ โ”œโ”€โ”€ latency.py # prefill / decode / cache speedup +โ”‚ โ”œโ”€โ”€ throughput.py # training tokens per second +โ”‚ โ””โ”€โ”€ report.py # JSON report CLI +โ”œโ”€โ”€ configs/ # default.yaml +โ”œโ”€โ”€ checkpoints/ # serialised state +โ”œโ”€โ”€ docs/ # TDD.md ยท BUILD_ORDER.md +โ”œโ”€โ”€ tests/ # unit/ + integration/ +โ””โ”€โ”€ .github/ # CI workflows +``` + +--- + +### Documentation + +| Document | Contents | +|---|---| +| [`docs/TDD.md`](docs/TDD.md) | Full technical design โ€” architecture, algorithms, memory analysis, tradeoff matrices, and 19 Technical Decision Records | +| [`docs/BUILD_ORDER.md`](docs/BUILD_ORDER.md) | Dependency-aware sequencing and the stable interface contracts between modules | +| `~/.claude/vault/language-model/` | Long-lived engineering notes and gotchas, kept outside the repo (TDR-018) | diff --git a/attention/__init__.py b/attention/__init__.py new file mode 100644 index 0000000..fe2723c --- /dev/null +++ b/attention/__init__.py @@ -0,0 +1,29 @@ +"""Causal grouped-query attention with RoPE, a KV cache, and three kernels. + +See docs/TDD.md ยง3.3 / ยง4 and TDR-007, TDR-008, TDR-015, TDR-016. +""" + +from attention.causal import VALID_IMPLS, CausalAttention +from attention.kernels import ( + causal_block_mask, + flash_attention, + manual_attention, + repeat_kv, + sdpa_attention, +) +from attention.kv_cache import KVCache +from attention.rope import RotaryEmbedding, apply_rotary_emb, rotate_half + +__all__ = [ + "VALID_IMPLS", + "CausalAttention", + "KVCache", + "RotaryEmbedding", + "apply_rotary_emb", + "causal_block_mask", + "flash_attention", + "manual_attention", + "repeat_kv", + "rotate_half", + "sdpa_attention", +] diff --git a/attention/causal.py b/attention/causal.py new file mode 100644 index 0000000..e59aeb4 --- /dev/null +++ b/attention/causal.py @@ -0,0 +1,216 @@ +"""Causal grouped-query attention with RoPE and an optional KV cache. + +Contract (docs/BUILD_ORDER.md): + + CausalAttention.forward(x, *, kv_cache=None, use_cache=False) + -> (out, new_kv_cache) + +The three masking regimes, which are the subtle part: + +| q_len | k_len | situation | masking | +|-------|-------------|---------------------|----------------------------------| +| T | T | training / prefill | standard causal triangle | +| 1 | cache + 1 | decode step | none โ€” one query sees everything | +| n | cache + n | chunked prefill | explicit bottom-right triangle | + +The third row is why we cannot simply pass ``is_causal=True`` everywhere: torch +aligns its implicit triangle to the top-left, which is only equivalent to ours +when ``q_len == k_len``. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from attention.kernels import ( + causal_block_mask, + flash_attention, + manual_attention, + repeat_kv, + sdpa_attention, +) +from attention.kv_cache import KVCache +from attention.rope import RotaryEmbedding, apply_rotary_emb + +VALID_IMPLS = ("manual", "sdpa", "flash") + + +class CausalAttention(nn.Module): + """Multi-head causal attention with grouped KV heads and rotary positions.""" + + def __init__( + self, + d_model: int, + n_heads: int, + n_kv_heads: int | None = None, + *, + max_seq_len: int = 1024, + rope_theta: float = 10000.0, + dropout: float = 0.0, + impl: str = "sdpa", + bias: bool = False, + block_q: int = 64, + block_k: int = 64, + ) -> None: + super().__init__() + n_kv_heads = n_heads if n_kv_heads is None else n_kv_heads + + if d_model % n_heads != 0: + raise ValueError(f"d_model={d_model} must be divisible by n_heads={n_heads}") + if n_heads % n_kv_heads != 0: + raise ValueError( + f"n_heads={n_heads} must be divisible by n_kv_heads={n_kv_heads} " + f"so each KV head serves an equal-sized query group" + ) + if impl not in VALID_IMPLS: + raise ValueError(f"unknown attention impl {impl!r}; expected one of {VALID_IMPLS}") + if impl == "flash" and dropout > 0.0: + # The tiled path never materialises normalised probabilities, so + # there is nothing to drop out of; silently ignoring the setting + # would make `flash` train differently from `manual`. + raise ValueError("attention dropout is not supported by impl='flash'") + + self.d_model = d_model + self.n_heads = n_heads + self.n_kv_heads = n_kv_heads + self.n_rep = n_heads // n_kv_heads + self.head_dim = d_model // n_heads + self.max_seq_len = max_seq_len + self.dropout = dropout + self.impl = impl + self.block_q = block_q + self.block_k = block_k + + # Q is full width; K/V are narrower by the group ratio. This asymmetry + # is why the fused Linear(C, 3C) of a classic MHA block does not apply. + self.q_proj = nn.Linear(d_model, n_heads * self.head_dim, bias=bias) + self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=bias) + self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=bias) + self.o_proj = nn.Linear(n_heads * self.head_dim, d_model, bias=bias) + self.resid_dropout = nn.Dropout(dropout) + + self.rope = RotaryEmbedding(self.head_dim, max_seq_len, theta=rope_theta) + + def forward( + self, + x: torch.Tensor, + *, + kv_cache: KVCache | None = None, + use_cache: bool = False, + key_padding_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, KVCache | None]: + """``key_padding_mask`` is ``[batch, k_len]``, True where a key is real. + + Needed for batched generation over ragged prompts: short prompts are + left-padded, and without masking those pad positions every short request + attends to filler as though it were context. + """ + batch, q_len, _ = x.shape + offset = len(kv_cache) if kv_cache is not None else 0 + + if offset + q_len > self.max_seq_len: + raise ValueError( + f"context length {offset + q_len} exceeds max_seq_len={self.max_seq_len}" + ) + if kv_cache is not None and kv_cache.keys is not None: + if kv_cache.keys.shape[0] != batch: + raise ValueError( + f"batch size changed mid-generation: cache holds " + f"{kv_cache.keys.shape[0]}, got {batch}" + ) + + # [B, T, C] -> [B, H, T, head_dim] + q = self.q_proj(x).view(batch, q_len, self.n_heads, self.head_dim).transpose(1, 2) + k = self.k_proj(x).view(batch, q_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + v = self.v_proj(x).view(batch, q_len, self.n_kv_heads, self.head_dim).transpose(1, 2) + + # Rotate at absolute positions. `offset` is the cache length, so a decode + # step rotates its single token at its true position rather than at 0. + cos, sin = self.rope(q_len, offset=offset) + q = apply_rotary_emb(q, cos, sin) + k = apply_rotary_emb(k, cos, sin) + + # Cache the *pre-expansion* KV heads: storing the repeated copies would + # throw away the entire GQA memory saving. + if use_cache: + kv_cache = kv_cache if kv_cache is not None else KVCache() + k, v = kv_cache.update(k, v) + elif kv_cache is not None: + k, v = kv_cache.update(k, v) + + k = repeat_kv(k, self.n_rep) + v = repeat_kv(v, self.n_rep) + + out = self._attend(q, k, v, offset=offset, key_padding_mask=key_padding_mask) + + out = out.transpose(1, 2).reshape(batch, q_len, self.n_heads * self.head_dim) + out = self.resid_dropout(self.o_proj(out)) + return out, (kv_cache if use_cache else None) + + def _attend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + offset: int, + key_padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + q_len, k_len = q.shape[2], k.shape[2] + + # A single query attends over the whole cache: nothing to mask -- unless + # some of that cache is padding. + no_mask_needed = q_len == 1 and key_padding_mask is None + # Square case only: torch's implicit triangle is top-left aligned. + plain_causal = q_len == k_len and key_padding_mask is None + + if self.impl == "flash": + if key_padding_mask is not None: + # The tiled kernel builds its mask per tile from positions alone. + # Threading a per-batch padding mask through it is real work, and + # a silent fallback to another kernel would hide which code path + # actually ran -- so refuse explicitly. + raise ValueError( + "impl='flash' does not support key_padding_mask; use 'sdpa' or " + "'manual' for batched generation over ragged prompts" + ) + return flash_attention( + q, + k, + v, + causal=not no_mask_needed, + offset=offset, + block_q=self.block_q, + block_k=self.block_k, + ) + + if no_mask_needed: + mask = None + elif plain_causal and self.impl == "sdpa": + mask = None # handled by is_causal below, which skips materialising it + else: + mask = causal_block_mask(q_len, k_len, offset, q.device) + + if key_padding_mask is not None: + if key_padding_mask.shape != (q.shape[0], k_len): + raise ValueError( + f"key_padding_mask must be [batch, k_len] = " + f"{(q.shape[0], k_len)}, got {tuple(key_padding_mask.shape)}" + ) + # Broadcast [B, k_len] -> [B, 1, 1, k_len] and union with the causal + # triangle. A key is blocked if it is in the future OR is padding. + padded = ~key_padding_mask[:, None, None, :] + mask = padded if mask is None else (mask | padded) + + if self.impl == "sdpa": + return sdpa_attention( + q, + k, + v, + mask=mask, + is_causal=plain_causal and not no_mask_needed, + dropout_p=self.dropout, + training=self.training, + ) + return manual_attention(q, k, v, mask=mask, dropout_p=self.dropout, training=self.training) diff --git a/attention/kernels.py b/attention/kernels.py new file mode 100644 index 0000000..63729c9 --- /dev/null +++ b/attention/kernels.py @@ -0,0 +1,322 @@ +"""Three interchangeable attention kernels (TDD ยง4.3). + +All three take ``q, k, v`` of shape ``[B, H, T, D]`` with KV heads already +expanded, and must return identical output within floating-point tolerance โ€” +``test_all_impls_agree`` is the gate. + +- ``manual_attention`` โ€” the readable reference and the correctness oracle. +- ``sdpa_attention`` โ€” dispatches to torch's fused kernels (including on MPS). +- ``flash_attention`` โ€” our own tiled/online-softmax version, which never + materialises the ``T x T`` score matrix. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn.functional as F + + +def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand ``[B, n_kv_heads, T, D]`` to ``[B, n_kv_heads * n_rep, T, D]``. + + Query head ``h`` is paired with KV head ``h // n_rep``, so the repeats of a + given KV head must land in *contiguous* output slots. Getting this ordering + wrong (e.g. tiling instead of interleaving) silently mispairs every query + head with the wrong group: shapes stay valid, the model still trains, it is + just quietly worse. Hence ``test_repeat_kv_maps_query_head_to_its_group``. + """ + if n_rep == 1: + return x + batch, n_kv_heads, seq, head_dim = x.shape + return ( + x[:, :, None, :, :] + .expand(batch, n_kv_heads, n_rep, seq, head_dim) + .reshape(batch, n_kv_heads * n_rep, seq, head_dim) + ) + + +def causal_block_mask(q_len: int, k_len: int, offset: int, device: torch.device) -> torch.Tensor: + """``[q_len, k_len]`` bool mask, ``True`` where attention is **disallowed**. + + Aligned bottom-right via ``offset``: query row *i* sits at absolute position + ``offset + i`` and may see keys ``0 .. offset + i``. This matters whenever + ``q_len != k_len`` (chunked prefill against a warm cache) โ€” torch's + ``is_causal=True`` aligns the triangle *top-left* instead, which would let + early queries see the future and hide most of the cache from later ones. + """ + q_pos = torch.arange(q_len, device=device) + offset + k_pos = torch.arange(k_len, device=device) + return k_pos[None, :] > q_pos[:, None] + + +def manual_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + mask: torch.Tensor | None, + dropout_p: float = 0.0, + training: bool = False, +) -> torch.Tensor: + """Textbook ``softmax(QKแต€ / โˆšd) V``. Allocates the full ``[B,H,Tq,Tk]`` scores.""" + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.shape[-1]) + if mask is not None: + scores = scores.masked_fill(mask, float("-inf")) + probs = F.softmax(scores, dim=-1) + if dropout_p > 0.0 and training: + probs = F.dropout(probs, p=dropout_p) + return torch.matmul(probs, v) + + +def sdpa_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + mask: torch.Tensor | None, + is_causal: bool = False, + dropout_p: float = 0.0, + training: bool = False, +) -> torch.Tensor: + """torch's fused kernel. ``mask`` follows our convention (True = blocked).""" + # sdpa's bool convention is the inverse of ours: True means *participate*. + attn_mask = None if mask is None else ~mask + return F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + is_causal=is_causal, + dropout_p=dropout_p if training else 0.0, + ) + + +def _flash_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + offset: int = 0, + block_q: int = 64, + block_k: int = 64, +) -> tuple[torch.Tensor, torch.Tensor]: + """Tiled forward pass. Returns ``(out, logsumexp)``. + + The log-sum-exp per query row is the *only* extra state the backward pass + needs. It is O(T) โ€” one scalar per query โ€” which is what lets backward + recompute the score tiles instead of the forward saving them (TDR-021). + + Walks query tiles on the outside and key tiles on the inside, carrying a + running max ``m`` and running denominator ``l`` per query row. When a key + tile raises the max, the accumulator is rescaled by ``exp(m_old - m_new)`` + before the new contribution is added โ€” algebraically identical to a single + global softmax, but only one ``[block_q, block_k]`` tile is ever resident, + so peak memory is O(T) rather than O(Tยฒ). + + Two edge cases that produce NaN if unhandled: + + - A key tile entirely beyond the causal frontier leaves the row max at + ``-inf``, and ``exp(-inf - -inf)`` is NaN. Rows are clamped to a finite + max before the exponential. + - Under causal masking, key tiles past the last query position contribute + nothing at all, so the loop breaks out rather than computing them โ€” this + is where the tiled form actually saves work, not just memory. + """ + *_, q_len, head_dim = q.shape + k_len = k.shape[-2] + scale = 1.0 / math.sqrt(head_dim) + out = torch.empty_like(q) + lse = torch.empty((*q.shape[:-1], 1), device=q.device, dtype=q.dtype) + + for i in range(0, q_len, block_q): + q_tile = q[:, :, i : i + block_q] + tile_q = q_tile.shape[2] + q_pos = torch.arange(i, i + tile_q, device=q.device) + offset + + shape = (*q.shape[:2], tile_q, 1) + running_max = torch.full(shape, float("-inf"), device=q.device, dtype=q.dtype) + running_sum = torch.zeros(shape, device=q.device, dtype=q.dtype) + acc = torch.zeros_like(q_tile) + + for j in range(0, k_len, block_k): + # Every later key tile is fully masked too, so stop rather than + # burn work on tiles that contribute nothing. + if causal and j > int(q_pos[-1].item()): + break + + k_tile = k[:, :, j : j + block_k] + v_tile = v[:, :, j : j + block_k] + scores = torch.matmul(q_tile, k_tile.transpose(-2, -1)) * scale + + if causal: + k_pos = torch.arange(j, j + k_tile.shape[2], device=q.device) + scores = scores.masked_fill(k_pos[None, :] > q_pos[:, None], float("-inf")) + + new_max = torch.maximum(running_max, scores.amax(dim=-1, keepdim=True)) + # Rows still fully masked have max -inf; pin them to 0 so the + # exponentials below evaluate to exactly 0 instead of NaN. Safe + # because such rows also have running_sum == 0 and acc == 0, so the + # choice of origin cannot affect the final quotient. + new_max = torch.where(torch.isneginf(new_max), torch.zeros_like(new_max), new_max) + + rescale = torch.exp(running_max - new_max) + probs = torch.exp(scores - new_max) + running_sum = running_sum * rescale + probs.sum(dim=-1, keepdim=True) + acc = acc * rescale + torch.matmul(probs, v_tile) + running_max = new_max + + # Causal attention always leaves a row at least its own key, so the + # denominator is positive; clamp only to keep a degenerate all-masked + # call from producing NaN instead of zeros. + denom = running_sum.clamp_min(torch.finfo(q.dtype).tiny) + out[:, :, i : i + tile_q] = acc / denom + # log-sum-exp in the original (unshifted) score space, so backward can + # rebuild the exact probabilities without knowing the tile order. + lse[:, :, i : i + tile_q] = running_max + torch.log(denom) + + return out, lse + + +def _flash_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + lse: torch.Tensor, + *, + causal: bool = True, + offset: int = 0, + block_q: int = 64, + block_k: int = 64, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Recompute score tiles and accumulate ``dq``, ``dk``, ``dv`` (TDR-021). + + With ``L`` (the saved log-sum-exp) the probabilities are recoverable exactly: + ``p_ij = exp(s_ij - L_i)``, no running max needed. Differentiating + ``o_i = ฮฃ_j p_ij v_j`` then gives + + D_i = ฮฃ_j p_ij (do_i ยท v_j) = do_i ยท o_i + dv_j += ฮฃ_i p_ij do_i + dp_ij = do_i ยท v_j + ds_ij = p_ij (dp_ij โˆ’ D_i) + dq_i += scale ฮฃ_j ds_ij k_j + dk_j += scale ฮฃ_i ds_ij q_i + + ``D_i`` collapses to a row-wise dot product of ``dout`` and ``out``, which is + why the forward only has to keep ``out`` and ``L`` โ€” the ``T ร— T`` matrices + ``p``, ``dp`` and ``ds`` exist one tile at a time and are discarded. + """ + head_dim = q.shape[-1] + scale = 1.0 / math.sqrt(head_dim) + q_len, k_len = q.shape[-2], k.shape[-2] + + dq = torch.zeros_like(q) + dk = torch.zeros_like(k) + dv = torch.zeros_like(v) + row_dot = (dout * out).sum(dim=-1, keepdim=True) # D_i + + for i in range(0, q_len, block_q): + span_q = slice(i, i + block_q) + q_tile = q[:, :, span_q] + dout_tile = dout[:, :, span_q] + lse_tile = lse[:, :, span_q] + dot_tile = row_dot[:, :, span_q] + tile_q = q_tile.shape[2] + q_pos = torch.arange(i, i + tile_q, device=q.device) + offset + dq_tile = torch.zeros_like(q_tile) + + for j in range(0, k_len, block_k): + if causal and j > int(q_pos[-1].item()): + break + + span_k = slice(j, j + block_k) + k_tile = k[:, :, span_k] + v_tile = v[:, :, span_k] + + scores = torch.matmul(q_tile, k_tile.transpose(-2, -1)) * scale + if causal: + k_pos = torch.arange(j, j + k_tile.shape[2], device=q.device) + scores = scores.masked_fill(k_pos[None, :] > q_pos[:, None], float("-inf")) + + probs = torch.exp(scores - lse_tile) + dv[:, :, span_k] += torch.matmul(probs.transpose(-2, -1), dout_tile) + + dprobs = torch.matmul(dout_tile, v_tile.transpose(-2, -1)) + dscores = probs * (dprobs - dot_tile) + + dq_tile += torch.matmul(dscores, k_tile) * scale + dk[:, :, span_k] += torch.matmul(dscores.transpose(-2, -1), q_tile) * scale + + dq[:, :, span_q] = dq_tile + + return dq, dk, dv + + +class FlashAttentionFn(torch.autograd.Function): + """Autograd wrapper that keeps the O(T) memory claim true in *training*. + + Without this, the tiled kernel is plain PyTorch ops, so autograd saves every + tile's intermediates for backward โ€” and the tiles sum back to O(Tยฒ). Measured + before this change: the tiled path saved *more* than the naive one (1.30ร— + at T=512) and both grew quadratically, so the memory advantage existed only + under ``no_grad``. + + Saving ``(q, k, v, out, lse)`` and recomputing the score tiles in backward + trades a second pass over the tiles for O(Tยทd) saved state instead of O(Tยฒ). + """ + + @staticmethod + def forward( # type: ignore[override] + ctx: Any, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + offset: int, + block_q: int, + block_k: int, + ) -> torch.Tensor: + out, lse = _flash_forward( + q, k, v, causal=causal, offset=offset, block_q=block_q, block_k=block_k + ) + ctx.save_for_backward(q, k, v, out, lse) + ctx.causal, ctx.offset = causal, offset + ctx.block_q, ctx.block_k = block_q, block_k + return out + + @staticmethod + def backward(ctx: Any, dout: torch.Tensor) -> tuple[Any, ...]: # type: ignore[override] + q, k, v, out, lse = ctx.saved_tensors + dq, dk, dv = _flash_backward( + dout.contiguous(), + q, + k, + v, + out, + lse, + causal=ctx.causal, + offset=ctx.offset, + block_q=ctx.block_q, + block_k=ctx.block_k, + ) + return dq, dk, dv, None, None, None, None + + +def flash_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + offset: int = 0, + block_q: int = 64, + block_k: int = 64, +) -> torch.Tensor: + """Tiled attention with online softmax and a recomputing backward (TDR-007, TDR-021).""" + out: torch.Tensor = FlashAttentionFn.apply(q, k, v, causal, offset, block_q, block_k) + return out diff --git a/attention/kv_cache.py b/attention/kv_cache.py new file mode 100644 index 0000000..5c4f791 --- /dev/null +++ b/attention/kv_cache.py @@ -0,0 +1,63 @@ +"""Per-layer key/value cache for autoregressive decoding (TDR-008). + +Without a cache, generating token *n* recomputes attention over all *n-1* prior +tokens, so a full generation costs O(Tยฒ). With one, each step computes a single +new K/V pair and reads the rest, making it O(T) overall. + +**Current implementation is concatenation-based**, i.e. each append allocates a +new tensor and copies the old one, which is O(T) per step and O(Tยฒ) copied bytes +across a generation. That is deliberate for now: it is obviously correct and the +copy is not the bottleneck at our context lengths. The preallocated ring buffer +that removes it is scoped as memory pooling in TDD ยง7.2 / build stage 6, where +it can be justified against a measured baseline rather than assumed (TDR-012). +""" + +from __future__ import annotations + +import torch + + +class KVCache: + """Growable K/V store for one attention layer. + + Holds ``[B, n_kv_heads, T, head_dim]`` tensors โ€” note ``n_kv_heads``, not + ``n_heads``: the cache stores the *unexpanded* grouped-query heads, which is + where GQA's memory saving actually lands (TDR-016). Expansion to full head + count happens after the read, on values that are never stored. + """ + + def __init__(self) -> None: + self.keys: torch.Tensor | None = None + self.values: torch.Tensor | None = None + + def update(self, keys: torch.Tensor, values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Append this step's K/V and return the full cached history.""" + if self.keys is None or self.values is None: + self.keys, self.values = keys, values + else: + if keys.shape[0] != self.keys.shape[0]: + raise ValueError( + f"batch size changed mid-generation: cache holds " + f"{self.keys.shape[0]}, got {keys.shape[0]}" + ) + self.keys = torch.cat((self.keys, keys), dim=2) + self.values = torch.cat((self.values, values), dim=2) + return self.keys, self.values + + def __len__(self) -> int: + """Number of positions currently cached.""" + return 0 if self.keys is None else int(self.keys.shape[2]) + + @property + def nbytes(self) -> int: + """Memory held by this cache โ€” the quantity GQA is trying to shrink.""" + if self.keys is None or self.values is None: + return 0 + return ( + self.keys.numel() * self.keys.element_size() + + self.values.numel() * self.values.element_size() + ) + + def reset(self) -> None: + self.keys = None + self.values = None diff --git a/attention/rope.py b/attention/rope.py new file mode 100644 index 0000000..2cd3944 --- /dev/null +++ b/attention/rope.py @@ -0,0 +1,76 @@ +"""Rotary Position Embeddings (TDR-015). + +RoPE rotates each consecutive pair of channels in Q and K by an angle +proportional to the token's absolute position. Because a dot product between +two rotated vectors depends only on the *difference* of their angles, the +resulting attention score is a function of relative distance โ€” absolute +position never appears in the score, only in the rotation applied on the way in. + +Two details that are easy to get wrong: + +- **V is never rotated.** RoPE encodes *where to look*, not *what is there*. + Rotating V corrupts the values being aggregated. +- **Under KV caching the offset is the cache length, not the slice index.** + A decode step passes a 1-token slice; rotating it at position 0 instead of + position ``cache_len`` makes every generated token believe it is the start of + the sequence. Prefill still looks perfect, so this bug survives most tests โ€” + hence ``test_cached_decode_matches_full_forward``. +""" + +from __future__ import annotations + +import torch +from torch import nn + + +class RotaryEmbedding(nn.Module): + """Precomputed cos/sin rotation tables for a given head dimension.""" + + # Declared so the type checker knows these are Tensors; `register_buffer` + # alone leaves them typed as `Tensor | Module`. + cos_cached: torch.Tensor + sin_cached: torch.Tensor + + def __init__(self, head_dim: int, max_seq_len: int, theta: float = 10000.0) -> None: + super().__init__() + if head_dim % 2 != 0: + raise ValueError(f"head_dim must be even to rotate in pairs, got {head_dim}") + self.head_dim = head_dim + self.max_seq_len = max_seq_len + self.theta = theta + + # Channel pair i rotates at frequency theta^(-2i/head_dim): low channels + # turn slowly (they encode long-range position), high channels quickly. + inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + angles = torch.outer(torch.arange(max_seq_len, dtype=torch.float32), inv_freq) + # Duplicated halves so the table lines up with rotate_half's [-x2, x1] layout. + emb = torch.cat((angles, angles), dim=-1) + + # persistent=False: these are pure functions of (head_dim, max_seq_len, + # theta) and would otherwise be dead weight in every checkpoint. + self.register_buffer("cos_cached", emb.cos(), persistent=False) + self.register_buffer("sin_cached", emb.sin(), persistent=False) + + def forward(self, seq_len: int, offset: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(cos, sin)`` of shape ``[seq_len, head_dim]`` for positions + ``offset .. offset + seq_len``.""" + end = offset + seq_len + if end > self.max_seq_len: + raise ValueError( + f"position {end} exceeds max_seq_len={self.max_seq_len}; " + f"rebuild the rotary table or shorten the context" + ) + return self.cos_cached[offset:end], self.sin_cached[offset:end] + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + """Map ``[x1, x2] -> [-x2, x1]`` over the split halves of the last dim.""" + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Rotate ``x`` ``[B, H, T, head_dim]`` by the given cos/sin ``[T, head_dim]``.""" + cos = cos.to(x.dtype) + sin = sin.to(x.dtype) + return x * cos + rotate_half(x) * sin diff --git a/benchmark/__init__.py b/benchmark/__init__.py new file mode 100644 index 0000000..965cd76 --- /dev/null +++ b/benchmark/__init__.py @@ -0,0 +1,41 @@ +"""Benchmark harness: latency, throughput and memory (TDD ยง9, TDR-012). + +Metrics split two ways. ``memory`` is exact arithmetic over tensor shapes and is +identical on every machine, so it can be asserted against a recorded baseline. +``latency`` and ``throughput`` are wall-clock and are reproducible only in +shape โ€” always read them next to ``report.environment()``. +""" + +from benchmark.harness import Timing, synchronize, time_it +from benchmark.latency import ( + benchmark_attention_impls, + benchmark_decode, + benchmark_prefill, + kv_cache_speedup, +) +from benchmark.memory import ( + activation_memory_estimate, + gqa_saving_report, + kv_cache_memory, + model_memory_report, +) +from benchmark.report import build_report, deterministic_report, environment +from benchmark.throughput import benchmark_training_step + +__all__ = [ + "Timing", + "activation_memory_estimate", + "benchmark_attention_impls", + "benchmark_decode", + "benchmark_prefill", + "benchmark_training_step", + "build_report", + "deterministic_report", + "environment", + "gqa_saving_report", + "kv_cache_memory", + "kv_cache_speedup", + "model_memory_report", + "synchronize", + "time_it", +] diff --git a/benchmark/harness.py b/benchmark/harness.py new file mode 100644 index 0000000..d2b7d22 --- /dev/null +++ b/benchmark/harness.py @@ -0,0 +1,132 @@ +"""Timing primitives (TDD ยง9). + +The one thing a GPU benchmark must get right is **synchronisation**. CUDA and +MPS both queue work asynchronously, so a naive ``perf_counter`` around a forward +pass measures how long it took to *submit* the kernels, not to run them โ€” which +reliably produces impossibly fast numbers that then get quoted. Every timed +region here syncs the device on both sides. + +Warmup matters for the same reason: the first call on any device pays for +kernel compilation, lazy module init and allocator growth. Those costs are real +but they are not per-step costs, so folding them into the average +misrepresents steady state. + +Latency is reported as **p50/p90/p95/p99, not just mean**. Step times are +right-skewed โ€” an allocator growth or a page fault produces occasional large +outliers โ€” so a mean quietly reports a number that no individual step achieved. + +A percentile is only as good as the sample count behind it: ``n`` runs resolve +tail probabilities no finer than ``1/n``, so p99 needs ~100 runs to mean +anything a p90 doesn't already say. ``Timing.resolvable_percentile`` reports +that ceiling next to the numbers, because an undersampled p99 looks exactly +like a well-sampled one. +""" + +from __future__ import annotations + +import statistics +import time +from collections.abc import Callable +from dataclasses import asdict, dataclass +from typing import Any + +import torch + + +def synchronize(device: torch.device | str | None) -> None: + """Block until queued work on ``device`` has actually finished.""" + if device is None: + return + kind = torch.device(device).type + if kind == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize() + elif kind == "mps" and torch.backends.mps.is_available(): + torch.mps.synchronize() + # CPU is synchronous already. + + +@dataclass(frozen=True) +class Timing: + """Wall-clock statistics for a repeated operation, in milliseconds.""" + + runs: int + mean_ms: float + stdev_ms: float + p50_ms: float + p90_ms: float + p95_ms: float + p99_ms: float + min_ms: float + max_ms: float + + @property + def median_ms(self) -> float: + """Alias for ``p50_ms`` โ€” the same statistic under its readable name.""" + return self.p50_ms + + @property + def resolvable_percentile(self) -> float: + """Highest percentile this sample count can distinguish from ``max``. + + With ``n`` samples the smallest resolvable tail probability is ``1/n``, + so 10 runs cannot say anything about p99 that it doesn't also say about + p90 โ€” both interpolate against the same top-most observation. Reported + alongside the percentiles so an undersampled p99 is visibly + undersampled rather than quietly meaningless. + """ + return round(100.0 * (1.0 - 1.0 / self.runs), 2) + + @classmethod + def from_samples(cls, samples_s: list[float]) -> Timing: + if not samples_s: + raise ValueError("need at least one sample") + ms = sorted(s * 1000.0 for s in samples_s) + + # Linear interpolation between order statistics (the standard + # "inclusive" estimator, matching numpy's default). The previous + # `ms[int(q * n)]` indexing collapsed every high percentile onto the + # maximum at small sample counts: with n=10, int(0.90*10) == int(0.99*10) == 9. + cuts = statistics.quantiles(ms, n=100, method="inclusive") if len(ms) > 1 else None + + def pct(q: int) -> float: + return ms[0] if cuts is None else cuts[q - 1] + + return cls( + runs=len(ms), + mean_ms=statistics.fmean(ms), + stdev_ms=statistics.stdev(ms) if len(ms) > 1 else 0.0, + p50_ms=pct(50), + p90_ms=pct(90), + p95_ms=pct(95), + p99_ms=pct(99), + min_ms=ms[0], + max_ms=ms[-1], + ) + + def as_dict(self) -> dict[str, Any]: + return {**asdict(self), "resolvable_percentile": self.resolvable_percentile} + + +def time_it( + fn: Callable[[], Any], + *, + warmup: int = 3, + repeats: int = 10, + device: torch.device | str | None = None, +) -> Timing: + """Run ``fn`` ``repeats`` times after ``warmup`` untimed calls.""" + if repeats < 1: + raise ValueError(f"repeats must be >= 1, got {repeats}") + + for _ in range(warmup): + fn() + synchronize(device) + + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + synchronize(device) # inside the timed region, or this measures nothing + samples.append(time.perf_counter() - start) + + return Timing.from_samples(samples) diff --git a/benchmark/latency.py b/benchmark/latency.py new file mode 100644 index 0000000..59aa6f6 --- /dev/null +++ b/benchmark/latency.py @@ -0,0 +1,134 @@ +"""Inference latency and throughput (TDD ยง6.4, ยง9). + +Generation has two distinct phases and averaging them together hides the thing +you actually care about: + +- **Prefill / first token** โ€” the whole prompt in one parallel pass. Compute-bound, + scales with prompt length. +- **Decode / per token** โ€” one token at a time against the cache. Memory-bandwidth + bound, roughly flat per step. + +A single "tokens/sec" number mixes the two and moves around with prompt length +for reasons that have nothing to do with decode speed, so both are reported. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from benchmark.harness import synchronize, time_it +from inference import generate_ids +from model import CausalLM + + +def benchmark_prefill( + model: CausalLM, + *, + prompt_len: int = 64, + batch_size: int = 1, + warmup: int = 2, + repeats: int = 10, +) -> dict[str, Any]: + """Time to process a prompt and produce the first next-token logits.""" + device = next(model.parameters()).device + prompt = torch.randint(0, model.vocab_size, (batch_size, prompt_len), device=device) + + def run() -> None: + cache = model.new_cache() + with torch.no_grad(): + model(prompt, kv_cache=cache, use_cache=True) + + timing = time_it(run, warmup=warmup, repeats=repeats, device=device) + tokens = batch_size * prompt_len + return { + "phase": "prefill", + "prompt_len": prompt_len, + "batch_size": batch_size, + "timing": timing.as_dict(), + "tokens_per_sec": round(tokens / (timing.median_ms / 1000.0), 2), + } + + +def benchmark_decode( + model: CausalLM, + *, + prompt_len: int = 16, + max_new_tokens: int = 32, + use_cache: bool = True, + warmup: int = 1, + repeats: int = 5, +) -> dict[str, Any]: + """Full generation, reported as per-token cost.""" + device = next(model.parameters()).device + prompt = torch.randint(0, model.vocab_size, (1, prompt_len), device=device) + + def run() -> None: + generate_ids( + model, + prompt, + max_new_tokens=max_new_tokens, + temperature=0.0, + use_cache=use_cache, + ) + + timing = time_it(run, warmup=warmup, repeats=repeats, device=device) + per_token_ms = timing.median_ms / max_new_tokens + return { + "phase": "decode", + "use_cache": use_cache, + "prompt_len": prompt_len, + "max_new_tokens": max_new_tokens, + "timing": timing.as_dict(), + "per_token_ms": round(per_token_ms, 4), + "tokens_per_sec": round(1000.0 / per_token_ms, 2), + } + + +def kv_cache_speedup( + model: CausalLM, *, prompt_len: int = 16, max_new_tokens: int = 32, repeats: int = 3 +) -> dict[str, Any]: + """Cached vs uncached decoding โ€” the measurement TDR-008 is claimed on. + + TDR-012 requires an optimization to demonstrate a measured gain rather than + assert one, so the ratio is reported whichever way it falls. At small model + sizes and short sequences the cache can genuinely lose: the O(Tยฒ) recompute + it removes is cheap, while the per-step Python and allocation overhead it + adds is not. That crossover is a real property worth surfacing, not a bug. + """ + common = {"prompt_len": prompt_len, "max_new_tokens": max_new_tokens, "repeats": repeats} + cached = benchmark_decode(model, use_cache=True, **common) + uncached = benchmark_decode(model, use_cache=False, **common) + ratio = uncached["timing"]["p50_ms"] / cached["timing"]["p50_ms"] + return { + "cached_ms": cached["timing"]["p50_ms"], + "uncached_ms": uncached["timing"]["p50_ms"], + "speedup": round(ratio, 3), + "cache_wins": ratio > 1.0, + "max_new_tokens": max_new_tokens, + } + + +def benchmark_attention_impls( + model_cfg: dict[str, Any], *, seq_len: int = 128, repeats: int = 5 +) -> dict[str, Any]: + """Compare the three attention kernels on one forward pass (TDD ยง4.3).""" + from attention import VALID_IMPLS + + results = {} + for impl in VALID_IMPLS: + cfg = {**model_cfg, "attention": {"impl": impl}} + torch.manual_seed(0) + model = CausalLM.from_config(cfg).eval() + device = next(model.parameters()).device + ids = torch.randint(0, model.vocab_size, (1, seq_len), device=device) + + def run(m: CausalLM = model, x: torch.Tensor = ids) -> None: + with torch.no_grad(): + m(x) + + synchronize(device) + results[impl] = time_it(run, warmup=2, repeats=repeats, device=device).as_dict() + + return {"seq_len": seq_len, "impls": results} diff --git a/benchmark/memory.py b/benchmark/memory.py new file mode 100644 index 0000000..66673c2 --- /dev/null +++ b/benchmark/memory.py @@ -0,0 +1,107 @@ +"""Memory metrics (TDD ยง9). + +Unlike latency, these numbers are **exact and machine-independent** โ€” they are +computed from tensor shapes and element sizes, not measured. That makes them the +only part of the benchmark suite that can be asserted against a recorded +baseline in CI, which is what `test_memory_report_reproduces_the_baseline` does. + +Deliberately *not* using peak-allocator readings here: `torch.cuda.max_memory_allocated` +has no MPS equivalent (TDR-019), so a peak-memory number would be CUDA-only and +not comparable across the devices this project actually runs on. +""" + +from __future__ import annotations + +from typing import Any + +from torch import nn + +from optimization.quantize import model_nbytes, quantize + +MIB = 1024 * 1024 + + +def model_memory_report(model: nn.Module) -> dict[str, Any]: + """Bytes held by the model under each precision, plus compression ratios.""" + fp32 = model_nbytes(model) + variants = { + "fp32": fp32, + "fp16": model_nbytes(quantize(model, scheme="fp16")), + "int8": model_nbytes(quantize(model, bits=8)), + "int4": model_nbytes(quantize(model, bits=4)), + } + return { + "bytes": variants, + "mib": {k: round(v / MIB, 4) for k, v in variants.items()}, + "compression_vs_fp32": {k: round(fp32 / v, 4) for k, v in variants.items()}, + "parameters": sum(p.numel() for p in model.parameters()), + } + + +def kv_cache_memory( + *, + n_layers: int, + n_kv_heads: int, + head_dim: int, + seq_len: int, + batch_size: int = 1, + bytes_per_element: int = 4, +) -> int: + """Bytes a full KV cache occupies: 2 (K and V) x layers x heads x T x dim.""" + return 2 * n_layers * batch_size * n_kv_heads * seq_len * head_dim * bytes_per_element + + +def gqa_saving_report( + *, + n_layers: int, + n_heads: int, + n_kv_heads: int, + head_dim: int, + seq_len: int, + batch_size: int = 1, +) -> dict[str, Any]: + """What grouped-query attention actually saves on the cache (TDR-016). + + Quoted per *config*, not per measurement: the saving is exactly the head + ratio, so reporting it as an empirical result would dress up arithmetic as + an experiment. + """ + common = { + "n_layers": n_layers, + "head_dim": head_dim, + "seq_len": seq_len, + "batch_size": batch_size, + } + mha = kv_cache_memory(n_kv_heads=n_heads, **common) + gqa = kv_cache_memory(n_kv_heads=n_kv_heads, **common) + return { + "mha_bytes": mha, + "gqa_bytes": gqa, + "mha_mib": round(mha / MIB, 4), + "gqa_mib": round(gqa / MIB, 4), + "saving_ratio": round(mha / gqa, 4), + "n_heads": n_heads, + "n_kv_heads": n_kv_heads, + # Echoed back because a cache size is uninterpretable without the length + # and batch it was computed at. + "seq_len": seq_len, + "batch_size": batch_size, + } + + +def activation_memory_estimate( + *, batch_size: int, n_heads: int, seq_len: int, bytes_per_element: int = 4 +) -> dict[str, Any]: + """Score-matrix footprint for standard vs tiled attention (TDR-007). + + Standard attention materialises ``[B, H, T, T]``; the tiled kernel never + holds more than one ``[block_q, block_k]`` tile, so its score memory is + independent of ``T``. This is the O(Tยฒ) โ†’ O(T) claim, stated in bytes. + """ + standard = batch_size * n_heads * seq_len * seq_len * bytes_per_element + return { + "standard_score_bytes": standard, + "standard_score_mib": round(standard / MIB, 4), + "seq_len": seq_len, + "scales_as": "O(T^2)", + } diff --git a/benchmark/report.py b/benchmark/report.py new file mode 100644 index 0000000..94c906a --- /dev/null +++ b/benchmark/report.py @@ -0,0 +1,143 @@ +"""Assemble and emit a benchmark report (TDD ยง9). + + python -m benchmark.report --config configs/default.yaml --out report.json + +The report separates two kinds of number, because they deserve very different +levels of trust: + +- ``deterministic`` โ€” memory footprints, compression ratios, cache sizes. Exact + arithmetic over shapes, identical on every machine. These are the ones a test + can assert against a recorded baseline. +- ``measured`` โ€” latency and throughput. Machine-, thermal- and load-dependent. + Reproducible only in shape, never in value, so nothing asserts them. + +Keeping them in one flat blob would invite exactly the mistake TDR-012 exists to +prevent: quoting a timing from one machine as though it were a property of the +code. +""" + +from __future__ import annotations + +import argparse +import json +import platform +from pathlib import Path +from typing import Any + +import torch + +from benchmark.latency import benchmark_decode, benchmark_prefill, kv_cache_speedup +from benchmark.memory import ( + activation_memory_estimate, + gqa_saving_report, + model_memory_report, +) +from model import CausalLM +from training.device import resolve_device + + +def environment() -> dict[str, Any]: + """Everything needed to interpret a measured number later.""" + device = resolve_device(None) + return { + "torch": torch.__version__, + "python": platform.python_version(), + "platform": platform.platform(), + "machine": platform.machine(), + "device": str(device), + "mps_available": torch.backends.mps.is_available(), + "cuda_available": torch.cuda.is_available(), + } + + +def deterministic_report(cfg: dict[str, Any]) -> dict[str, Any]: + """Exact, machine-independent metrics. Safe to assert against a baseline.""" + m = cfg["model"] + torch.manual_seed(0) + model = CausalLM.from_config(cfg) + head_dim = int(m["d_model"]) // int(m["n_heads"]) + seq_len = int(m.get("max_seq_len", 1024)) + + return { + "config": {k: m[k] for k in sorted(m)}, + "model_memory": model_memory_report(model), + "gqa_kv_cache": gqa_saving_report( + n_layers=int(m["n_layers"]), + n_heads=int(m["n_heads"]), + n_kv_heads=int(m.get("n_kv_heads", m["n_heads"])), + head_dim=head_dim, + seq_len=seq_len, + ), + "attention_scores": activation_memory_estimate( + batch_size=1, n_heads=int(m["n_heads"]), seq_len=seq_len + ), + } + + +def measured_report( + cfg: dict[str, Any], *, prompt_len: int = 32, max_new_tokens: int = 32 +) -> dict[str, Any]: + """Wall-clock metrics. Interpret only alongside ``environment()``. + + Sample counts are set by what each phase costs, not by a single default. + A percentile resolves tail probabilities no finer than ``1/runs``, so the + 5-repeat default would report a p99 that is just the maximum under another + name. Prefill is ~3 ms, so 100 runs is nearly free; decode is ~80 ms, so 20 + runs is the affordable compromise โ€” and ``resolvable_percentile`` in the + output states the ceiling either way. + """ + torch.manual_seed(0) + model = CausalLM.from_config(cfg).eval().to(resolve_device(cfg.get("device"))) + return { + "prefill": benchmark_prefill(model, prompt_len=prompt_len, warmup=5, repeats=100), + "decode_cached": benchmark_decode( + model, + prompt_len=prompt_len, + max_new_tokens=max_new_tokens, + use_cache=True, + warmup=2, + repeats=20, + ), + "kv_cache_speedup": kv_cache_speedup( + model, prompt_len=prompt_len, max_new_tokens=max_new_tokens, repeats=10 + ), + } + + +def build_report(cfg: dict[str, Any], *, measured: bool = True) -> dict[str, Any]: + report: dict[str, Any] = { + "environment": environment(), + "deterministic": deterministic_report(cfg), + } + if measured: + report["measured"] = measured_report(cfg) + return report + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run the benchmark suite") + parser.add_argument("--config", default="configs/default.yaml") + parser.add_argument("--out", default=None, help="write JSON here") + parser.add_argument( + "--deterministic-only", + action="store_true", + help="skip wall-clock measurement (exact metrics only)", + ) + args = parser.parse_args() + + from training.train import load_config + + cfg = load_config(args.config) + report = build_report(cfg, measured=not args.deterministic_only) + + text = json.dumps(report, indent=2, sort_keys=True) + if args.out: + Path(args.out).parent.mkdir(parents=True, exist_ok=True) + Path(args.out).write_text(text + "\n", encoding="utf-8") + print(f"wrote {args.out}") + else: + print(text) + + +if __name__ == "__main__": + main() diff --git a/benchmark/showcase.py b/benchmark/showcase.py new file mode 100644 index 0000000..d2569c3 --- /dev/null +++ b/benchmark/showcase.py @@ -0,0 +1,345 @@ +"""A single screenshot-ready summary of the engine. + + python -m benchmark.showcase --config configs/tinyshakespeare.yaml \ + --checkpoint checkpoints/tinyshakespeare/best.pt \ + --vocab data/tinyshakespeare/vocab.json + +Prints four blocks: the architecture as configured, the exact memory figures, the +measured latency, and live completions from the checkpoint. Everything here is +computed at run time โ€” nothing is a stored string โ€” so a screenshot of the output +is a screenshot of the code actually working. + +Deliberately prints the held-out perplexity next to the completions. A generation +sample on its own invites the reader to judge fluency, which at this scale is +flattering and meaningless; the perplexity is the number that says what the model +is really worth. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from torch import nn + +from benchmark.memory import gqa_saving_report, model_memory_report +from model import CausalLM +from tokenizer import BPE +from training.checkpoint import load_checkpoint +from training.device import resolve_device + +RULE = "โ”€" * 74 + +DEFAULT_PROMPTS = ( + "ROMEO:", + "To be, or not to be", + "My lord, the queen", + "What say you to", +) + + +def header(title: str) -> None: + print(f"\n{RULE}\n {title}\n{RULE}") + + +def show_architecture(cfg: dict, model: CausalLM) -> None: + m = cfg["model"] + header("ARCHITECTURE โ€” decoder-only, built from scratch") + rows = [ + ("parameters", f"{model.num_parameters():,}"), + ("layers x d_model", f"{m['n_layers']} x {m['d_model']}"), + ("attention", f"{m['n_heads']} query heads / {m['n_kv_heads']} KV heads (GQA)"), + ("head dim", str(m["d_model"] // m["n_heads"])), + ("feed-forward", f"SwiGLU, d_ff {m['d_ff']}"), + ("normalisation", "RMSNorm, pre-norm"), + ("positions", f"RoPE, theta {m.get('rope_theta', 10000.0):.0f}"), + ("context", f"{m['max_seq_len']} tokens"), + ("vocabulary", f"{m['vocab_size']} (byte-level BPE, no OOV)"), + ("kernel", f"{cfg.get('attention', {}).get('impl', 'sdpa')} (manual | sdpa | flash)"), + ] + for key, value in rows: + print(f" {key:<20} {value}") + + +def show_memory(model: CausalLM, cfg: dict) -> None: + m = cfg["model"] + header("MEMORY โ€” exact arithmetic, identical on every machine") + report = model_memory_report(model) + print(f" {'precision':<12}{'size':>12}{'compression':>14}") + for name in ("fp32", "fp16", "int8", "int4"): + mib = report["mib"][name] + ratio = report["compression_vs_fp32"][name] + print(f" {name:<12}{mib:>9.2f} MiB{ratio:>13.2f}x") + + gqa = gqa_saving_report( + n_layers=int(m["n_layers"]), + n_heads=int(m["n_heads"]), + n_kv_heads=int(m.get("n_kv_heads", m["n_heads"])), + head_dim=int(m["d_model"]) // int(m["n_heads"]), + seq_len=int(m["max_seq_len"]), + ) + print( + f"\n KV cache @ {gqa['seq_len']} tokens: " + f"MHA {gqa['mha_mib']:.2f} MiB -> GQA {gqa['gqa_mib']:.2f} MiB " + f"({gqa['saving_ratio']:.2f}x smaller)" + ) + print(" int8 is not 4x: the tied embedding stays fp32 and sets the floor.") + + +def show_latency(model: CausalLM) -> None: + from benchmark.latency import benchmark_decode, benchmark_prefill, kv_cache_speedup + + header("LATENCY โ€” wall clock, this machine only") + prefill = benchmark_prefill(model, prompt_len=32, warmup=3, repeats=50) + decode = benchmark_decode(model, prompt_len=32, max_new_tokens=32, warmup=2, repeats=10) + + for label, report in (("prefill", prefill), ("decode", decode)): + t = report["timing"] + print( + f" {label:<8} p50 {t['p50_ms']:7.2f} ms p90 {t['p90_ms']:7.2f} " + f"p95 {t['p95_ms']:7.2f} p99 {t['p99_ms']:7.2f} " + f"({t['runs']} runs, resolvable <= p{t['resolvable_percentile']})" + ) + print(f"\n prefill throughput {prefill['tokens_per_sec']:>10,.0f} tok/s") + print(f" decode throughput {decode['tokens_per_sec']:>10,.0f} tok/s") + + # Reported across several lengths, because a single figure hides the point: + # the cache removes O(T^2) recompute, so its advantage *grows* with the + # generation. Quoting only a short run makes a real asymptotic win look + # like noise. + print("\n KV cache โ€” measured, not assumed (TDR-012):") + print(f" {'new tokens':>11}{'uncached':>11}{'cached':>10}{'speedup':>10}") + for length in (32, 128, 256): + result = kv_cache_speedup(model, prompt_len=16, max_new_tokens=length, repeats=3) + print( + f" {length:>11}{result['uncached_ms']:>8.0f} ms{result['cached_ms']:>7.0f} ms" + f"{result['speedup']:>9.2f}x" + ) + + +def show_kernel_equivalence(cfg: dict) -> None: + """Prove the three attention kernels agree, rather than assert it. + + An optimised kernel with no reference to check against is not trustworthy. + `manual` is the readable oracle; the other two have to match it. + """ + import torch + + from attention import CausalAttention + from attention.kernels import causal_block_mask, flash_attention, manual_attention + + header("KERNEL EQUIVALENCE โ€” three implementations, one answer") + m = cfg["model"] + torch.manual_seed(0) + reference = CausalAttention( + int(m["d_model"]), + int(m["n_heads"]), + int(m["n_kv_heads"]), + max_seq_len=int(m["max_seq_len"]), + impl="manual", + ).eval() + x = torch.randn(2, 64, int(m["d_model"])) + + baseline, _ = reference(x) + print(f" {'kernel':<10}{'max abs diff vs manual':>26}") + for impl in ("manual", "sdpa", "flash"): + torch.manual_seed(0) + other = CausalAttention( + int(m["d_model"]), + int(m["n_heads"]), + int(m["n_kv_heads"]), + max_seq_len=int(m["max_seq_len"]), + impl=impl, + ).eval() + other.load_state_dict(reference.state_dict()) + out, _ = other(x) + print(f" {impl:<10}{(out - baseline).abs().max().item():>26.2e}") + + # Gradient check in float64: float32 noise hides sign and scale errors in a + # hand-written backward. + shape = (2, 3, 37, 16) + q, k, v = (torch.randn(*shape, dtype=torch.float64, requires_grad=True) for _ in range(3)) + qf, kf, vf = (t.detach().clone().requires_grad_(True) for t in (q, k, v)) + ref = manual_attention(q, k, v, mask=causal_block_mask(37, 37, 0, q.device)) + tiled = flash_attention(qf, kf, vf, causal=True, block_q=8, block_k=8) + seed = torch.randn_like(ref) + ref.backward(seed) + tiled.backward(seed) + pairs = ((q, qf), (k, kf), (v, vf)) + assert all(a.grad is not None and b.grad is not None for a, b in pairs) + worst = max( + (a.grad - b.grad).abs().max().item() # type: ignore[operator] + for a, b in pairs + ) + print(f"\n hand-written flash backward vs autograd (float64): {worst:.2e}") + + +def show_flash_memory(cfg: dict) -> None: + """The O(T) memory claim, measured for *training* rather than inference.""" + import torch + + from attention import CausalAttention + + header("TILED ATTENTION MEMORY โ€” bytes autograd retains for backward") + d_model = int(cfg["model"]["d_model"]) + + def retained(impl: str, seq_len: int) -> int: + torch.manual_seed(0) + attn = CausalAttention( + d_model, 4, 2, max_seq_len=4096, impl=impl, block_q=32, block_k=32 + ).eval() + total = 0 + + def pack(t: torch.Tensor) -> torch.Tensor: + nonlocal total + total += t.numel() * t.element_size() + return t + + with torch.enable_grad(), torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t): + attn(torch.randn(1, seq_len, d_model)) + return total + + print( + f" {'T':>6}{'standard':>14}{'tiled':>12}{'ratio':>9}{'std growth':>13}{'tiled growth':>15}" + ) + prev_std = prev_tiled = None + for seq_len in (128, 256, 512, 1024): + std, tiled = retained("manual", seq_len), retained("flash", seq_len) + g_std = f"{std / prev_std:.2f}x" if prev_std else "โ€”" + g_tiled = f"{tiled / prev_tiled:.2f}x" if prev_tiled else "โ€”" + print( + f" {seq_len:>6}{std / 1024:>11.0f} KiB{tiled / 1024:>9.0f} KiB" + f"{tiled / std:>8.2f}x{g_std:>13}{g_tiled:>15}" + ) + prev_std, prev_tiled = std, tiled + print("\n 2.0x per doubling is linear; 4.0x is quadratic.") + print(" Before the custom backward the tiled path retained 1.30x MORE than standard.") + + +def show_quantization(model: CausalLM, cfg: dict, val_path: str | None, batches: int) -> None: + """Compression against its actual cost in held-out perplexity.""" + from optimization import model_nbytes, quantize + + header("QUANTIZATION โ€” compression against measured quality cost") + fp32_bytes = model_nbytes(model) + print(f" {'precision':<12}{'size':>12}{'compression':>14}{'held-out ppl':>16}") + for label, bits in (("fp32", None), ("int8", 8), ("int4", 4)): + target: nn.Module = model if bits is None else quantize(model, bits=bits) + size = model_nbytes(target) + ppl = ( + f"{measure_perplexity(target, cfg, val_path, batches):.1f}" + if val_path + else "not measured" + ) + print(f" {label:<12}{size / (1024 * 1024):>9.2f} MiB{fp32_bytes / size:>13.2f}x{ppl:>16}") + print("\n Weight-only: buys memory, not speed. Real int8 throughput needs fused kernels.") + + +def measure_perplexity(model: nn.Module, cfg: dict, val_path: str, batches: int) -> float: + """Evaluate held-out perplexity here and now, rather than accept a number. + + A figure typed in on the command line is indistinguishable from one that was + invented. Measuring it in the same process that prints it means the whole + output is generated. + """ + from dataset import DataEngine + from training.trainer import Trainer + + eval_cfg = { + **cfg, + "dataset": {**cfg["dataset"], "source": "file", "path": val_path, "shuffle": False}, + "training": {**cfg.get("training", {}), "amp": False}, + } + engine = DataEngine.from_config(eval_cfg, vocab_size=int(cfg["model"]["vocab_size"])) + trainer = Trainer(model, engine, eval_cfg, device=next(model.parameters()).device) + return float(trainer.evaluate(engine, batches=batches)["val_perplexity"]) + + +def show_completions( + model: CausalLM, tokenizer: BPE, prompts: tuple[str, ...], val_ppl: float | None +) -> None: + from inference import generate + + header("AUTOCOMPLETION โ€” greedy-free sampling from the trained checkpoint") + if val_ppl is not None: + print( + f" held-out perplexity {val_ppl:.1f} (measured now) โ€” read the samples against this\n" + ) + + for prompt in prompts: + text = generate( + model, + tokenizer, + prompt, + max_new_tokens=48, + temperature=0.8, + top_k=40, + generator=torch.Generator().manual_seed(0), + ) + continuation = text[len(prompt) :].replace("<|endoftext|>", " โŽ ").strip() + print(f" > {prompt}") + for line in continuation.splitlines()[:4]: + if line.strip(): + print(f" {line.strip()}") + print() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Print a showcase summary") + parser.add_argument("--config", default="configs/tinyshakespeare.yaml") + parser.add_argument("--checkpoint", default=None) + parser.add_argument("--vocab", default=None) + parser.add_argument( + "--val-path", + default=None, + help="held-out .bin; perplexity is measured here rather than passed in", + ) + parser.add_argument("--eval-batches", type=int, default=20) + parser.add_argument("--skip-latency", action="store_true") + parser.add_argument( + "--evidence", + action="store_true", + help="print kernel equivalence, tiled-attention memory and quantization cost", + ) + args = parser.parse_args() + + from training.train import load_config + + cfg = load_config(args.config) + device = resolve_device(cfg.get("device")) + torch.manual_seed(int(cfg.get("seed", 0))) + + model = CausalLM.from_config(cfg) + step = None + if args.checkpoint and Path(args.checkpoint).exists(): + payload = load_checkpoint(args.checkpoint, model=model) + step = payload.get("step") + model = model.to(device).eval() + + print(f"\n Language Model Engine โ€” {device}, torch {torch.__version__}") + if step is not None: + print(f" checkpoint: {args.checkpoint} (step {step})") + + show_architecture(cfg, model) + show_memory(model, cfg) + if not args.skip_latency: + show_latency(model) + if args.evidence: + show_kernel_equivalence(cfg) + show_flash_memory(cfg) + show_quantization(model, cfg, args.val_path, args.eval_batches) + + val_ppl = None + if args.val_path and Path(args.val_path).exists(): + val_ppl = measure_perplexity(model, cfg, args.val_path, args.eval_batches) + + if args.vocab and Path(args.vocab).exists(): + tokenizer = BPE(vocab_size=int(cfg["model"]["vocab_size"])) + tokenizer.load(args.vocab) + show_completions(model, tokenizer, DEFAULT_PROMPTS, val_ppl) + print(RULE) + + +if __name__ == "__main__": + main() diff --git a/benchmark/throughput.py b/benchmark/throughput.py new file mode 100644 index 0000000..467bef7 --- /dev/null +++ b/benchmark/throughput.py @@ -0,0 +1,49 @@ +"""Training throughput (TDD ยง9). + +Reports tokens/sec for a full optimizer step โ€” forward, backward, clip and +update โ€” because that is the number that determines how long a run takes. +Timing the forward pass alone overstates throughput by roughly 3x, since the +backward pass costs about twice the forward. +""" + +from __future__ import annotations + +from typing import Any + +from benchmark.harness import synchronize +from training.trainer import Trainer + + +def benchmark_training_step( + trainer: Trainer, *, warmup: int = 3, repeats: int = 10 +) -> dict[str, Any]: + """Measure steady-state step time on an already-constructed Trainer. + + Warmup steps are real optimizer steps, so this mutates the model โ€” it is a + throughput probe, not something to run against weights you care about. + """ + for _ in range(warmup): + trainer.step() + synchronize(trainer.device) + + samples_ms = [] + tokens = 0 + for _ in range(repeats): + metrics = trainer.step() + synchronize(trainer.device) + # Trainer already measures its own step wall-clock; reuse it rather than + # double-counting with an outer timer. + samples_ms.append(metrics["tokens"] / metrics["tokens_per_sec"] * 1000.0) + tokens += int(metrics["tokens"]) + + samples_ms.sort() + median = samples_ms[len(samples_ms) // 2] + return { + "steps": repeats, + "median_step_ms": round(median, 4), + "tokens_per_step": tokens // repeats, + "tokens_per_sec": round((tokens // repeats) / (median / 1000.0), 2), + "device": str(trainer.device), + "amp_dtype": str(trainer.amp_dtype), + "grad_accum_steps": trainer.grad_accum_steps, + } diff --git a/checkpoints/.gitkeep b/checkpoints/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/configs/.gitkeep b/configs/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/configs/default.yaml b/configs/default.yaml new file mode 100644 index 0000000..11d14e3 --- /dev/null +++ b/configs/default.yaml @@ -0,0 +1,47 @@ +# Default training / model configuration. +# Overridable per-key via the CLI: python -m training.train --config configs/default.yaml training.lr=1e-4 + +seed: 42 + +model: + d_model: 512 + n_layers: 8 + n_heads: 8 + n_kv_heads: 2 # grouped-query attention (n_kv_heads < n_heads) + d_ff: 2048 + vocab_size: 4096 + max_seq_len: 1024 + rope_theta: 10000.0 + dropout: 0.0 + tie_weights: true # reuse the token embedding as the LM head + +attention: + impl: sdpa # sdpa | flash | manual + kv_cache: true + +dataset: + source: synthetic # synthetic | file (mmap) | stream (sequential) + path: null # required for source: file | stream + seq_len: 512 + batch_size: 16 + num_workers: 0 + +training: + steps: 2000 + lr: 3.0e-4 + min_lr: 3.0e-5 + warmup_steps: 200 + weight_decay: 0.1 + grad_clip: 1.0 + amp: true + amp_dtype: bfloat16 + grad_accum_steps: 1 + log_every: 50 + ckpt_every: 500 + out_dir: checkpoints + +inference: + max_new_tokens: 128 + temperature: 1.0 + top_k: 0 + top_p: 1.0 diff --git a/configs/tinyshakespeare.yaml b/configs/tinyshakespeare.yaml new file mode 100644 index 0000000..999900f --- /dev/null +++ b/configs/tinyshakespeare.yaml @@ -0,0 +1,61 @@ +# A ready-to-run config over the TinyShakespeare corpus (public domain). +# +# Prepare the corpus once, then train: +# +# python -m dataset.prepare_corpus --corpus tinyshakespeare --vocab-size 8192 +# python -m training.train --config configs/tinyshakespeare.yaml +# +# Sized from a measured run rather than guessed. The first attempt at 2000 steps +# was 41.6 epochs over a 295k-token corpus at 26 parameters per token: training +# loss fell 5.26 -> 1.86 while validation perplexity bottomed at 98.9 around +# step 500 and then climbed to 304.6. The model was memorising, and because +# checkpoints were periodic the best weights were never saved. +# +# So: fewer steps, evaluate 2.5x more often, save on improvement, and stop early +# when four consecutive evaluations fail to improve. Expect validation to turn +# upward before training does โ€” with a corpus this small that is arithmetic, not +# a bug, and it is exactly what the held-out split is for. + +seed: 42 +device: auto # auto | mps | cuda | cpu + +model: + d_model: 256 + n_layers: 6 + n_heads: 8 + n_kv_heads: 2 + d_ff: 1024 + vocab_size: 8192 # must match the trained tokenizer + max_seq_len: 256 + rope_theta: 10000.0 + dropout: 0.1 # a corpus this small needs the regularisation + tie_weights: true + +attention: + impl: sdpa + kv_cache: true + +dataset: + source: file + path: data/tinyshakespeare/train.bin + val_path: data/tinyshakespeare/val.bin + seq_len: 256 + batch_size: 24 + num_workers: 0 + +training: + steps: 800 + lr: 6.0e-4 + min_lr: 6.0e-5 + warmup_steps: 100 + weight_decay: 0.1 + grad_clip: 1.0 + amp: true + amp_dtype: bfloat16 + grad_accum_steps: 1 + log_every: 100 + eval_every: 100 + eval_batches: 20 + early_stop_patience: 4 + ckpt_every: 0 # best.pt is saved on improvement instead + out_dir: checkpoints/tinyshakespeare diff --git a/dataset/__init__.py b/dataset/__init__.py new file mode 100644 index 0000000..52f6316 --- /dev/null +++ b/dataset/__init__.py @@ -0,0 +1,48 @@ +"""Efficient dataset engine: memory-mapped, streaming, worker-sharded. + +See docs/TDD.md ยง2 and TDR-004. +""" + +from dataset.binfile import ( + TOKEN_DTYPE, + count_tokens_bin, + encode_texts_to_bin, + read_tokens_bin, + write_tokens_bin, +) +from dataset.corpora import ( + CORPORA, + Corpus, + download, + prepare, + prepare_documents, + read_documents, + split_documents, + use_system_trust_store, +) +from dataset.engine import DataEngine +from dataset.token_dataset import ( + MemmapTokenDataset, + StreamingTokenDataset, + SyntheticTokenDataset, +) + +__all__ = [ + "CORPORA", + "TOKEN_DTYPE", + "Corpus", + "DataEngine", + "MemmapTokenDataset", + "StreamingTokenDataset", + "SyntheticTokenDataset", + "count_tokens_bin", + "download", + "encode_texts_to_bin", + "prepare", + "prepare_documents", + "read_documents", + "read_tokens_bin", + "split_documents", + "use_system_trust_store", + "write_tokens_bin", +] diff --git a/dataset/binfile.py b/dataset/binfile.py new file mode 100644 index 0000000..a98748d --- /dev/null +++ b/dataset/binfile.py @@ -0,0 +1,80 @@ +"""On-disk token format: a bare flat array of little-endian ``uint16``. + +No header, no index, no framing. That is deliberate -- the file is exactly +``n_tokens * 2`` bytes, so the window count is derivable from ``stat()`` alone +(no need to open or parse anything), and any tool can read it: ``np.memmap``, +``torch.from_file``, or ``xxd``. Metadata that would otherwise live in a header +(vocab size, tokenizer version) belongs with the tokenizer's ``vocab.json``, +which must be kept alongside the ``.bin`` -- token ids are meaningless without +the vocabulary that produced them. + +``uint16`` caps the vocabulary at 65,536, which is comfortably above the 4,096 +we train with and matches GPT-2's 50,257. Going past that means ``uint32`` and +double the disk/page-cache footprint, so we validate rather than widen. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Iterator +from pathlib import Path + +import numpy as np + +# The single source of truth for the on-disk dtype. Import this rather than +# writing np.uint16 at call sites, so a future widening is one edit. +TOKEN_DTYPE = np.uint16 + +_MAX_TOKEN_ID = np.iinfo(TOKEN_DTYPE).max + + +def write_tokens_bin(ids: Iterable[int], path: str | os.PathLike[str]) -> int: + """Write token ids to ``path`` as a flat ``uint16`` array. Returns the count. + + Raises ``ValueError`` on any id outside ``[0, 65535]``. We check explicitly + because numpy would otherwise wrap silently (``65536 -> 0``), corrupting the + corpus in a way that surfaces much later as inexplicable training loss. + """ + arr = np.fromiter(ids, dtype=np.int64) + if arr.size and (arr.min() < 0 or arr.max() > _MAX_TOKEN_ID): + raise ValueError( + f"token ids must fit uint16 [0, {_MAX_TOKEN_ID}]; got range [{arr.min()}, {arr.max()}]" + ) + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + arr.astype(TOKEN_DTYPE).tofile(path) + return int(arr.size) + + +def read_tokens_bin(path: str | os.PathLike[str]) -> np.memmap: + """Open a token file read-only as a memory map (no bulk copy into RAM).""" + return np.memmap(path, dtype=TOKEN_DTYPE, mode="r") + + +def count_tokens_bin(path: str | os.PathLike[str]) -> int: + """Number of tokens in a token file, from its size alone -- no open, no read.""" + return Path(path).stat().st_size // np.dtype(TOKEN_DTYPE).itemsize + + +def encode_texts_to_bin( + tokenizer: object, + texts: Iterable[str], + path: str | os.PathLike[str], + *, + eot_token: str | None = None, +) -> int: + """Tokenize ``texts`` into a single ``.bin`` corpus. Returns the token count. + + ``eot_token`` (e.g. ``"<|endoftext|>"``) is appended after each text so the + model learns a document boundary; without it, packed windows silently splice + the end of one document onto the start of the next. + """ + encode = tokenizer.encode # type: ignore[attr-defined] + eot_ids: list[int] = encode(eot_token) if eot_token else [] + + def stream() -> Iterator[int]: + for text in texts: + yield from encode(text) + yield from eot_ids + + return write_tokens_bin(stream(), path) diff --git a/dataset/corpora.py b/dataset/corpora.py new file mode 100644 index 0000000..efe6556 --- /dev/null +++ b/dataset/corpora.py @@ -0,0 +1,210 @@ +"""Corpus acquisition and document-level splitting (TDD ยง2). + +Corpora are downloaded on demand and never committed โ€” they are large and +licence-encumbered, and `.gitignore` already excludes `data/`. Each entry names +its licence, because "where did the training data come from" is a question that +gets asked later and is unpleasant to answer retroactively. + +**The split is by document, never by token.** Windows are contiguous slices of +one flat token array, so splitting a *tokenised* stream puts validation tokens +inside training windows: the model trains on text it is then evaluated on, and +the reported perplexity is quietly too good. Splitting the document list first +makes that impossible by construction. +""" + +from __future__ import annotations + +import hashlib +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +DATA_DIR = Path("data") + + +@dataclass(frozen=True) +class Corpus: + """A downloadable text corpus and the terms it comes under.""" + + name: str + url: str + licence: str + description: str + doc_separator: str = "\n\n" + + @property + def raw_path(self) -> Path: + return DATA_DIR / f"{self.name}.txt" + + +# Small, permissively licensed, and plain-text โ€” the constraints that matter for +# a from-scratch project with no data-engineering budget. +CORPORA: dict[str, Corpus] = { + "tinyshakespeare": Corpus( + name="tinyshakespeare", + url=( + "https://raw.githubusercontent.com/karpathy/char-rnn/" + "master/data/tinyshakespeare/input.txt" + ), + licence="public domain (Shakespeare)", + description="~1.1 MB of Shakespeare. Small enough to tokenise in seconds.", + ), +} + + +def use_system_trust_store() -> bool: + """Route TLS verification through the OS trust store. True if it took effect. + + Needed on any machine behind a TLS-inspecting corporate proxy (Zscaler, + Netskope, Blue Coat). The proxy re-signs every certificate with a corporate + root CA that lives in the OS keychain, so `curl` succeeds โ€” but `requests` + verifies against certifi's bundle, which has never heard of it, and fails + with `CERTIFICATE_VERIFY_FAILED: unable to get local issuer certificate`. + + Never disable verification to work around this: that turns a working + corporate MITM into an accepted arbitrary one. Point at the real CA instead, + either via the OS store here or `REQUESTS_CA_BUNDLE` (which requests honours + natively). + """ + try: + import truststore + except ImportError: + return False + truststore.inject_into_ssl() + return True + + +def download(corpus: Corpus | str, *, force: bool = False, timeout: int = 60) -> Path: + """Fetch a corpus to ``data/``. Returns the local path; skips if present.""" + entry = CORPORA[corpus] if isinstance(corpus, str) else corpus + path = entry.raw_path + if path.exists() and not force: + return path + + import requests # imported lazily so the package works offline + + used_system_trust = use_system_trust_store() + path.parent.mkdir(parents=True, exist_ok=True) + try: + response = requests.get(entry.url, timeout=timeout) + response.raise_for_status() + except requests.exceptions.SSLError as exc: + raise RuntimeError( + f"TLS verification failed fetching {entry.url}.\n" + f"System trust store in use: {used_system_trust}.\n" + "On a TLS-inspecting corporate network the certificate is re-signed " + "by a corporate root CA that certifi does not carry. Fix with either:\n" + " pip install truststore # use the OS keychain\n" + " export REQUESTS_CA_BUNDLE=/path/to/corporate-ca.pem\n" + "Do not disable verification." + ) from exc + + path.write_text(response.text, encoding="utf-8") + return path + + +def read_documents(path: str | os.PathLike[str], separator: str = "\n\n") -> list[str]: + """Split a raw text file into documents, dropping blanks.""" + text = Path(path).read_text(encoding="utf-8") + return [doc.strip() for doc in text.split(separator) if doc.strip()] + + +def split_documents( + documents: list[str], *, val_fraction: float = 0.05, seed: int = 0 +) -> tuple[list[str], list[str]]: + """Partition documents into (train, val) with a stable, content-derived hash. + + Hashing the document text rather than shuffling an index means the split is + reproducible without carrying a random seed alongside the data, and a + document keeps its side even if the corpus is re-ordered or extended. Adding + new documents therefore never reshuffles the existing ones into the other + split, which would leak across a resumed run. + """ + if not 0.0 < val_fraction < 1.0: + raise ValueError(f"val_fraction must be in (0, 1), got {val_fraction}") + + threshold = int(val_fraction * (1 << 32)) + train: list[str] = [] + val: list[str] = [] + for doc in documents: + digest = hashlib.blake2b(f"{seed}:{doc}".encode(), digest_size=4).digest() + (val if int.from_bytes(digest, "big") < threshold else train).append(doc) + + if not train or not val: + raise ValueError( + f"split produced an empty side ({len(train)} train, {len(val)} val); " + f"corpus has only {len(documents)} documents" + ) + return train, val + + +def prepare_documents( + documents: list[str], + tokenizer: object, + *, + name: str = "custom", + licence: str = "unspecified", + val_fraction: float = 0.05, + seed: int = 0, + eot_token: str | None = "<|endoftext|>", + out_dir: str | os.PathLike[str] | None = None, +) -> dict[str, Any]: + """Split โ†’ encode โ†’ write ``train.bin`` / ``val.bin`` from documents in memory. + + The entry point for any corpus that is not a plain URL download โ€” a local + export, a scrape, a mailbox. + + **On personal data.** A language model memorises: this project measured a + 32.5M-parameter model reaching loss 0.017 on a 10,677-token corpus, i.e. + reproducing it. Any corpus fed here ends up recoverable from the weights, so + a checkpoint trained on mail, chat logs or customer records *is* a copy of + that data and inherits its handling obligations. Scrub before this call, not + after โ€” and note ``licence`` so the provenance travels with the artefacts. + """ + from dataset.binfile import encode_texts_to_bin + + train_docs, val_docs = split_documents(documents, val_fraction=val_fraction, seed=seed) + + # Coerce: the natural call passes a str, and `str / str` is a TypeError. + destination = Path(out_dir) if out_dir is not None else DATA_DIR / name + train_bin = destination / "train.bin" + val_bin = destination / "val.bin" + n_train = encode_texts_to_bin(tokenizer, train_docs, train_bin, eot_token=eot_token) + n_val = encode_texts_to_bin(tokenizer, val_docs, val_bin, eot_token=eot_token) + + return { + "corpus": name, + "licence": licence, + "documents": len(documents), + "train_documents": len(train_docs), + "val_documents": len(val_docs), + "train_tokens": n_train, + "val_tokens": n_val, + "train_bin": str(train_bin), + "val_bin": str(val_bin), + } + + +def prepare( + corpus: Corpus | str, + tokenizer: object, + *, + val_fraction: float = 0.05, + seed: int = 0, + eot_token: str | None = "<|endoftext|>", + out_dir: str | os.PathLike[str] | None = None, +) -> dict[str, Any]: + """Download โ†’ split by document โ†’ encode โ†’ write ``train.bin`` / ``val.bin``.""" + entry = CORPORA[corpus] if isinstance(corpus, str) else corpus + documents = read_documents(download(entry), entry.doc_separator) + return prepare_documents( + documents, + tokenizer, + name=entry.name, + licence=entry.licence, + val_fraction=val_fraction, + seed=seed, + eot_token=eot_token, + out_dir=out_dir, + ) diff --git a/dataset/engine.py b/dataset/engine.py new file mode 100644 index 0000000..bba4d11 --- /dev/null +++ b/dataset/engine.py @@ -0,0 +1,148 @@ +"""``DataEngine`` -- config in, endless device-ready batches out. + +The training loop counts *steps*, not epochs, so ``next_batch()`` never raises +``StopIteration``: a finite dataset simply wraps around. Everything else here is +DataLoader plumbing whose defaults are chosen to avoid three specific traps: + +- ``prefetch_factor`` is illegal when ``num_workers == 0`` (torch raises), so it + is only passed when workers exist. +- ``pin_memory`` allocates CUDA page-locked host memory. With no CUDA device it + is at best a no-op and at worst a warning on every batch, so it defaults to + ``torch.cuda.is_available()`` rather than ``True``. Apple MPS reads from + unified memory and gains nothing from pinning (TDR-019). +- ``shuffle`` is rejected outright by DataLoader for an ``IterableDataset``; + streaming sources shuffle internally instead. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +import torch +from torch.utils.data import DataLoader, Dataset, IterableDataset + +from dataset.token_dataset import ( + Item, + MemmapTokenDataset, + StreamingTokenDataset, + SyntheticTokenDataset, +) + +Batch = dict[str, torch.Tensor] + +VALID_SOURCES = ("synthetic", "file", "stream") + + +class DataEngine: + """Wraps a dataset in a DataLoader and serves an endless stream of batches.""" + + def __init__( + self, + dataset: Dataset[Item], + *, + batch_size: int = 8, + num_workers: int = 0, + prefetch_factor: int | None = None, + pin_memory: bool | None = None, + shuffle: bool = True, + drop_last: bool = True, + ) -> None: + self.dataset = dataset + self.batch_size = batch_size + self.num_workers = num_workers + # An IterableDataset owns its own ordering; DataLoader raises if asked + # to shuffle one. + self.shuffle = shuffle and not isinstance(dataset, IterableDataset) + self.drop_last = drop_last + self.pin_memory = torch.cuda.is_available() if pin_memory is None else pin_memory + self.prefetch_factor = prefetch_factor + + self._loader = self._build_loader() + self._iter: Iterator[Batch] = iter(self._loader) + + # ------------------------------------------------------------------ + # Construction + # ------------------------------------------------------------------ + @classmethod + def from_config(cls, cfg: dict[str, Any], vocab_size: int) -> DataEngine: + """Build from the ``dataset`` block of a loaded YAML config.""" + dcfg = cfg.get("dataset", {}) + source = dcfg.get("source", "synthetic") + seq_len = int(dcfg.get("seq_len", 512)) + seed = int(cfg.get("seed", 0)) + + dataset: Dataset[Item] + if source == "synthetic": + dataset = SyntheticTokenDataset(vocab_size, seq_len, seed=seed) + elif source in ("file", "stream"): + path = dcfg.get("path") + if not path: + raise ValueError(f"dataset.source={source!r} requires dataset.path") + if source == "file": + dataset = MemmapTokenDataset(path, seq_len) + else: + dataset = StreamingTokenDataset( + path, seq_len, shuffle=dcfg.get("shuffle", True), seed=seed + ) + else: + raise ValueError(f"unknown dataset.source {source!r}; expected one of {VALID_SOURCES}") + + return cls( + dataset, + batch_size=int(dcfg.get("batch_size", 8)), + num_workers=int(dcfg.get("num_workers", 0)), + prefetch_factor=dcfg.get("prefetch_factor"), + pin_memory=dcfg.get("pin_memory"), + shuffle=bool(dcfg.get("shuffle", True)), + drop_last=bool(dcfg.get("drop_last", True)), + ) + + def _build_loader(self) -> DataLoader[Item]: + kwargs: dict[str, Any] = { + "batch_size": self.batch_size, + "num_workers": self.num_workers, + "pin_memory": self.pin_memory, + "drop_last": self.drop_last, + } + if not isinstance(self.dataset, IterableDataset): + kwargs["shuffle"] = self.shuffle + if self.num_workers > 0: + # Both of these are invalid with num_workers=0. + if self.prefetch_factor is not None: + kwargs["prefetch_factor"] = int(self.prefetch_factor) + kwargs["persistent_workers"] = True + return DataLoader(self.dataset, **kwargs) + + # ------------------------------------------------------------------ + # Batch production + # ------------------------------------------------------------------ + def next_batch(self, device: torch.device | str | None = None) -> Batch: + """Next batch, wrapping around at the end of a finite dataset.""" + try: + batch = next(self._iter) + except StopIteration: + self._iter = iter(self._loader) + try: + batch = next(self._iter) + except StopIteration as exc: + # An empty loader means the dataset cannot fill even one batch, + # which `drop_last=True` turns into zero batches rather than a + # short one. Raised here with the arithmetic spelled out, because + # the bare StopIteration surfaces from inside torch's sampler and + # says nothing about why. Hit most often by a small val split. + windows = len(self.dataset) if hasattr(self.dataset, "__len__") else "unknown" + raise RuntimeError( + f"dataset yielded no batches: {windows} windows available but " + f"batch_size={self.batch_size} with drop_last={self.drop_last}. " + f"Lower batch_size or seq_len, or set drop_last=False." + ) from exc + if device is None: + return batch + # non_blocking pairs with pinned memory to overlap H2D with compute; it + # degrades to a synchronous copy elsewhere. + return {k: v.to(device, non_blocking=self.pin_memory) for k, v in batch.items()} + + def __iter__(self) -> Iterator[Batch]: + while True: + yield self.next_batch() diff --git a/dataset/messages.py b/dataset/messages.py new file mode 100644 index 0000000..f4af648 --- /dev/null +++ b/dataset/messages.py @@ -0,0 +1,173 @@ +"""Preparing email and chat data as training documents (TDD ยง2). + +This module builds documents from mail/message records. It deliberately ships +with **no connector to any mailbox or chat workspace** โ€” the caller supplies the +records, so pointing it at real data is an explicit act rather than a default. + +## Read this before using it on real mail + +**A language model memorises its training data.** This project measured a +32.5M-parameter model reaching loss 0.017 on a 10,677-token corpus โ€” i.e. +reproducing it. Anything fed through here is recoverable from the resulting +weights, so **a checkpoint trained on mail is a copy of that mail** and inherits +every handling obligation the mail had: retention, access control, deletion +requests, cross-border transfer. + +**`redact()` is best-effort, not a compliance control.** It removes shapes it +recognises (addresses, phone numbers, long digit runs). It cannot recognise a +name in prose, a deal codename, a customer referred to obliquely, or an amount +that is sensitive only in context. Treat it as noise reduction on top of a +decision that the corpus is already acceptable to train on โ€” never as the thing +that makes it acceptable. + +**Scale makes it a fine-tuning corpus, not a pretraining one.** A mailbox is +~1-10M tokens against the ~650M a 32.5M-parameter model wants. Trained alone it +memorises rather than generalises; the sound route is pretrain on a large public +corpus, then fine-tune. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field + +# Structure markers, so the model can learn where a field begins rather than +# inferring it from punctuation. Register these as tokenizer special tokens. +FROM_TOKEN = "<|from|>" +SUBJECT_TOKEN = "<|subject|>" +BODY_TOKEN = "<|body|>" +MESSAGE_TOKEN = "<|message|>" +STRUCTURE_TOKENS = (FROM_TOKEN, SUBJECT_TOKEN, BODY_TOKEN, MESSAGE_TOKEN) + +# Quoted-reply openers. Without stripping these, a thread of N replies contributes +# the same text N times and the model over-weights whatever gets quoted most. +_QUOTE_MARKERS = ( + re.compile(r"^\s*>.*$", re.MULTILINE), + re.compile(r"^\s*On .{0,80}wrote:\s*$.*", re.MULTILINE | re.DOTALL), + re.compile(r"^\s*-{2,}\s*Original Message\s*-{2,}.*", re.MULTILINE | re.DOTALL), + re.compile(r"^\s*From:.*$.*", re.MULTILINE | re.DOTALL), +) + +# `-- ` on its own line is the RFC 3676 signature delimiter. +_SIGNATURE = re.compile(r"^--\s*$.*", re.MULTILINE | re.DOTALL) + +# Boilerplate footers dominate corporate mail by volume. Left in, the model +# learns to emit disclaimers, because that is genuinely the most predictable text +# in the corpus. +_DISCLAIMER = re.compile( + # "is confidential", "are confidential and intended", "is intended solely" -- + # the copula varies with whether attachments are mentioned, and clauses sit + # between it and the keyword. Bounded to one sentence ([^.\n]) so ordinary + # prose that merely uses the word "confidential" is not swallowed. + r"(this (e-?mail|message)\b[^.\n]{0,80}\b(is|are)\b[^.\n]{0,40}(confidential|intended)" + r"|if you (are not|have received) th(is|e intended))", + re.IGNORECASE, +) + +# Order matters: the phone pattern accepts separators and would otherwise claim a +# bare 16-digit card or account number. Digit runs with no separators are matched +# first and labelled <|number|>; only separated runs reach <|phone|>. +_REDACTIONS: tuple[tuple[re.Pattern[str], str], ...] = ( + (re.compile(r"[\w.+-]+@[\w-]+\.[\w.-]+"), "<|email|>"), + (re.compile(r"https?://\S+"), "<|url|>"), + (re.compile(r"\b\d{6,}\b"), "<|number|>"), + (re.compile(r"\+?\d[\d\s().-]{7,}\d"), "<|phone|>"), +) + + +@dataclass +class Message: + """One mail or chat record. ``sender`` may already be a pseudonym.""" + + body: str + sender: str | None = None + subject: str | None = None + metadata: dict[str, str] = field(default_factory=dict) + + +def strip_quoted_replies(text: str) -> str: + """Remove quoted history so a thread is not counted once per reply.""" + for pattern in _QUOTE_MARKERS: + text = pattern.sub("", text) + return text + + +def strip_signature(text: str) -> str: + return _SIGNATURE.sub("", text) + + +def looks_like_disclaimer(line: str) -> bool: + return bool(_DISCLAIMER.search(line)) + + +def strip_disclaimers(text: str) -> str: + """Drop everything from the first confidentiality-notice line onward.""" + lines = text.splitlines() + for i, line in enumerate(lines): + if looks_like_disclaimer(line): + return "\n".join(lines[:i]) + return text + + +def redact(text: str) -> str: + """Mask recognisable identifier *shapes*. Best-effort โ€” see the module note.""" + for pattern, placeholder in _REDACTIONS: + text = pattern.sub(placeholder, text) + return text + + +def clean_body(text: str, *, apply_redaction: bool = True) -> str: + """Full cleaning pipeline for one message body.""" + text = strip_quoted_replies(text) + text = strip_disclaimers(text) + text = strip_signature(text) + if apply_redaction: + text = redact(text) + # Collapse the runs of blank lines that stripping leaves behind. + return re.sub(r"\n{3,}", "\n\n", text).strip() + + +def to_document( + message: Message, *, include_structure: bool = True, apply_redaction: bool = True +) -> str: + """Render one message as a training document.""" + body = clean_body(message.body, apply_redaction=apply_redaction) + if not include_structure: + return body + + parts = [MESSAGE_TOKEN] + if message.sender: + parts += [FROM_TOKEN, redact(message.sender) if apply_redaction else message.sender] + if message.subject: + parts += [SUBJECT_TOKEN, redact(message.subject) if apply_redaction else message.subject] + parts += [BODY_TOKEN, body] + return "\n".join(parts) + + +def build_documents( + messages: list[Message], + *, + include_structure: bool = True, + apply_redaction: bool = True, + min_chars: int = 40, +) -> list[str]: + """Clean and render messages, dropping ones too short to carry signal. + + Very short messages ("thanks", "ok, done") are the bulk of a real mailbox by + count and near-worthless by content; keeping them teaches the model to emit + acknowledgements. ``min_chars`` applies to the *cleaned* body, so a long + quoted reply with two new words is correctly dropped. + """ + documents = [] + for message in messages: + body = clean_body(message.body, apply_redaction=apply_redaction) + if len(body) < min_chars: + continue + documents.append( + to_document( + Message(body=body, sender=message.sender, subject=message.subject), + include_structure=include_structure, + apply_redaction=apply_redaction, + ) + ) + return documents diff --git a/dataset/prepare_corpus.py b/dataset/prepare_corpus.py new file mode 100644 index 0000000..3c45b99 --- /dev/null +++ b/dataset/prepare_corpus.py @@ -0,0 +1,73 @@ +"""Corpus preparation CLI. + + python -m dataset.prepare_corpus --corpus tinyshakespeare --vocab-size 8192 + +Downloads, trains a tokenizer, splits by document, and writes +``train.bin`` / ``val.bin`` plus the ``vocab.json`` they were encoded with. + +The vocabulary is written **beside the token files on purpose**. Token ids are +meaningless without the vocabulary that produced them, and a `.bin` paired with +the wrong `vocab.json` decodes to plausible-looking nonsense rather than failing, +so the two travel together. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from dataset.corpora import CORPORA, prepare, read_documents +from tokenizer import BPE + + +def main() -> None: + parser = argparse.ArgumentParser(description="Download and prepare a corpus") + parser.add_argument("--corpus", default="tinyshakespeare", choices=sorted(CORPORA)) + parser.add_argument("--vocab-size", type=int, default=8192) + parser.add_argument("--val-fraction", type=float, default=0.05) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--tokenizer-sample-mb", + type=float, + default=64.0, + help="cap the text used to fit the vocabulary (TDD ยง1.5)", + ) + args = parser.parse_args() + + entry = CORPORA[args.corpus] + from dataset.corpora import download + + path = download(entry) + documents = read_documents(path, entry.doc_separator) + print(f"{entry.name}: {len(documents)} documents from {path} ({entry.licence})") + + # Fit the vocabulary on a bounded sample rather than the whole corpus: merge + # frequencies converge long before the corpus is exhausted, and training cost + # scales with the text fed in. + budget = int(args.tokenizer_sample_mb * 1024 * 1024) + sample: list[str] = [] + used = 0 + for doc in documents: + if used >= budget: + break + sample.append(doc) + used += len(doc) + print(f"fitting vocabulary on {used / 1024 / 1024:.2f} MB ({len(sample)} documents)") + + tokenizer = BPE(vocab_size=args.vocab_size, special_tokens=["<|endoftext|>"]) + tokenizer.train(sample) + print(f"vocabulary: {tokenizer.vocab_size} tokens, {len(tokenizer.merges)} merges") + + info = prepare(entry, tokenizer, val_fraction=args.val_fraction, seed=args.seed) + vocab_path = Path(str(info["train_bin"])).parent / "vocab.json" + tokenizer.save(str(vocab_path)) + info["vocab"] = str(vocab_path) + + total = int(info["train_tokens"]) + int(info["val_tokens"]) + info["chars_per_token"] = round(sum(len(d) for d in documents) / max(1, total), 3) + print(json.dumps(info, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/dataset/token_dataset.py b/dataset/token_dataset.py new file mode 100644 index 0000000..8cbf7bc --- /dev/null +++ b/dataset/token_dataset.py @@ -0,0 +1,168 @@ +"""Window readers over a ``uint16`` token corpus. + +Two datasets, one item contract. Both yield +``{"input_ids": LongTensor[T], "labels": LongTensor[T]}`` where ``labels`` is +``input_ids`` shifted one position left. + +The shift is done at *window* level: a window reads ``seq_len + 1`` ids and +splits them, rather than reading ``seq_len`` and shifting inside. That costs one +extra token of overlap per window and buys a real label for the final position +instead of padding -- with a 512-token window, shifting inside would silently +discard 1/512 of the training signal and leave one position learning from a fake +target. + +Which class to use: + +- ``MemmapTokenDataset`` (map-style) -- a fixed corpus on disk. Supports + ``len()`` and random access, so ``DataLoader(shuffle=True)`` can shuffle + window order for free. The default (TDR-004). +- ``StreamingTokenDataset`` (iterable) -- sequential passes and worker-sharded + reads, for corpora where random access is undesirable. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.utils.data import Dataset, IterableDataset, get_worker_info + +from dataset.binfile import TOKEN_DTYPE, count_tokens_bin + +Item = dict[str, torch.Tensor] + + +def _n_windows(n_tokens: int, seq_len: int) -> int: + """Whole windows available. Each consumes seq_len+1 ids (T inputs + 1 label).""" + return max(0, (n_tokens - 1) // seq_len) + + +class _MemmapReader: + """Shared lazy-memmap machinery for both dataset classes. + + The memmap handle is opened on first use and *dropped when pickled*. That + matters: ``DataLoader(num_workers>0)`` uses ``spawn`` on macOS, which pickles + the dataset into each worker. A live ``np.memmap`` pickles as a plain + ``ndarray`` -- i.e. it would serialise the entire corpus through a pipe into + every worker, which is precisely the bulk copy mmap exists to avoid. Opening + lazily per process gives each worker its own independent handle, as intended + by TDR-004. + """ + + def __init__(self, path: str | os.PathLike[str], seq_len: int) -> None: + if seq_len < 1: + raise ValueError(f"seq_len must be >= 1, got {seq_len}") + self.path = Path(path) + self.seq_len = int(seq_len) + self.n_tokens = count_tokens_bin(self.path) + self.n_windows = _n_windows(self.n_tokens, self.seq_len) + if self.n_windows < 1: + raise ValueError( + f"corpus too small: {self.path} holds {self.n_tokens} tokens but a " + f"window of seq_len={self.seq_len} needs {self.seq_len + 1}" + ) + self._tokens: np.memmap | None = None + + @property + def tokens(self) -> np.memmap: + if self._tokens is None: + self._tokens = np.memmap(self.path, dtype=TOKEN_DTYPE, mode="r") + return self._tokens + + def window(self, index: int) -> Item: + """Read window ``index`` and split it into inputs and shifted labels.""" + start = index * self.seq_len + # Copy out of the mmap: the returned tensors outlive this call and must + # not alias page-cache memory that a later window may remap. + chunk = np.asarray(self.tokens[start : start + self.seq_len + 1], dtype=np.int64) + return { + "input_ids": torch.from_numpy(chunk[:-1].copy()), + "labels": torch.from_numpy(chunk[1:].copy()), + } + + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__.copy() + state["_tokens"] = None # never ship an open handle across a process boundary + return state + + +class MemmapTokenDataset(_MemmapReader, Dataset[Item]): + """Map-style, memory-mapped view over a token corpus.""" + + def __len__(self) -> int: + return self.n_windows + + def __getitem__(self, index: int) -> Item: + if index < 0: + index += self.n_windows + if not 0 <= index < self.n_windows: + raise IndexError(f"window index out of range: {index} (have {self.n_windows})") + return self.window(index) + + +class StreamingTokenDataset(_MemmapReader, IterableDataset[Item]): + """Iterable, worker-sharded view over a token corpus. + + Sharding is strided (``order[worker_id::num_workers]``) so every window is + emitted by exactly one worker -- duplicated windows would quietly inflate the + effective epoch and bias training toward whatever the overlap contains. + """ + + def __init__( + self, + path: str | os.PathLike[str], + seq_len: int, + *, + shuffle: bool = False, + seed: int = 0, + ) -> None: + super().__init__(path, seq_len) + self.shuffle = shuffle + self.seed = seed + self.epoch = 0 + + def __iter__(self) -> Iterator[Item]: + order = np.arange(self.n_windows) + if self.shuffle: + # Reseeded per epoch so successive passes differ, but reproducibly. + np.random.default_rng([self.seed, self.epoch]).shuffle(order) + self.epoch += 1 + + info = get_worker_info() + if info is not None: + order = order[info.id :: info.num_workers] + + for index in order: + yield self.window(int(index)) + + +class SyntheticTokenDataset(IterableDataset[Item]): + """Endless uniform-random token stream -- a corpus-free smoke-test source. + + There is no signal to learn here, so loss should sit near ``ln(vocab_size)`` + and stay there. That makes it a useful control: a model that "improves" on + synthetic data has a leak (labels in the inputs, or a shifted mask). + """ + + def __init__(self, vocab_size: int, seq_len: int, *, seed: int = 0) -> None: + if vocab_size < 1: + raise ValueError(f"vocab_size must be >= 1, got {vocab_size}") + if seq_len < 1: + raise ValueError(f"seq_len must be >= 1, got {seq_len}") + self.vocab_size = vocab_size + self.seq_len = seq_len + self.seed = seed + + def __iter__(self) -> Iterator[Item]: + info = get_worker_info() + # Offset the seed per worker, or every worker emits the identical stream. + generator = torch.Generator().manual_seed(self.seed + (info.id if info else 0)) + while True: + chunk = torch.randint( + 0, self.vocab_size, (self.seq_len + 1,), generator=generator, dtype=torch.long + ) + yield {"input_ids": chunk[:-1], "labels": chunk[1:]} diff --git a/docs/BUILD_ORDER.md b/docs/BUILD_ORDER.md new file mode 100644 index 0000000..4113fc4 --- /dev/null +++ b/docs/BUILD_ORDER.md @@ -0,0 +1,109 @@ +# Build Order & Module Interfaces + +Dependency-aware implementation plan for the remaining modules. Each stage is +independently testable and unblocks the next. The architecture is specified in +[`TDD.md`](TDD.md); this document is the *sequencing* and the *contracts* that +keep the modules decoupled. + +## Dependency graph + +``` +tokenizer โ”€โ”€โ–บ dataset โ”€โ”€โ–บ training โ”€โ”€โ–บ checkpoints โ”€โ”€โ–บ inference + โ–ฒ โ–ฒ โ–ฒ + โ”‚ โ”‚ โ”‚ + (token ids) model โ—„โ”€โ”€ attention optimization + โ–ฒ โ–ฒ + โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ benchmark โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ +``` + +- `attention` and `dataset` only depend on `tokenizer`/config โ†’ build them early, in parallel. +- `model` composes `attention`. +- `training` needs `dataset` + `model`. +- `inference` needs a trained `model` checkpoint. +- `optimization` transforms a trained model โ†’ feeds back into `inference`. +- `benchmark` instruments everything; build a thin harness early, grow it last. + +## Stages (in order) + +| # | Module | Depends on | Deliverable | Exit test | +|---|--------|-----------|-------------|-----------| +| 0 | **tokenizer** โœ… | โ€” | `BPE` + GPT-2 pre-tokenization, piece cache | 54 unit tests | +| 1 | **dataset** โœ… | tokenizer | mmap/streaming datasets, corpus download, document split | 76 unit tests | +| 2 | **attention** โœ… | config | causal GQA + RoPE + KV cache, 3 kernels + recomputing backward | 43 unit tests | +| 3 | **model** โœ… | attention | `CausalLM` (embed โ†’ blocks โ†’ RMSNorm โ†’ LM head) | 25 unit tests | +| 4 | **training** โœ… | dataset + model | AdamW + cosine + AMP loop | 58 unit tests (incl. overfit-one-batch) | +| 5 | **inference** โœ… | model ckpt | greedy / top-k / top-p, cached + batched generation | 29 unit tests | +| 6 | **optimization** โœ… | model ckpt | int8/int4 per-channel quant + preallocated KV arena | 38 unit tests (incl. perplexity-within-tolerance) | +| 7 | **benchmark** โœ… | all | latency / throughput / memory harness | 18 unit tests (incl. baseline reproduction) | + +> All 8 stages are implemented (371 tests). + +## Contracts (stable interfaces โ€” change only via ADR) + +```python +# dataset +class TokenDataset(IterableDataset): + def __iter__( + self, + ) -> Iterator[dict]: ... # {"input_ids": LongTensor[T], "labels": LongTensor[T]} + + +# attention +class CausalAttention(nn.Module): + def forward( + self, x, *, kv_cache=None, use_cache=False, key_padding_mask=None + ): ... # -> (out, new_kv_cache) + + +# model +class CausalLM(nn.Module): + def forward( + self, input_ids, *, kv_cache=None, use_cache=False, key_padding_mask=None + ): ... # -> logits [B, T, V] + def new_cache(self) -> list[KVCache]: ... # one entry per layer + @classmethod + def from_config(cls, cfg: dict) -> "CausalLM": ... + + +# forward returns logits only -- cross-entropy lives in the trainer, so the +# model stays pure for inference. Because of that there is nowhere to return a +# freshly built cache, so the cache is CALLER-OWNED: allocate once with +# new_cache() and pass the same list every step; entries mutate in place. +# use_cache=True without a cache raises rather than silently discarding it. + + +# inference +def sample(logits, *, temperature=1.0, top_k=0, top_p=1.0) -> int: ... +def generate(model, tokenizer, prompt: str, *, max_new_tokens: int, **sampling) -> str: ... +def generate_batch( + model, prompts: list[list[int]], *, max_new_tokens: int, **sampling +) -> list[list[int]]: ... + + +# Batched generation left-pads and masks the padded keys, so a ragged batch +# yields the same tokens as decoding each prompt alone. Requires impl sdpa or +# manual; the tiled kernel rejects a per-batch padding mask rather than +# silently ignoring it. + + +# optimization +def quantize(model, *, bits=8, scheme="int_weight", skip=("lm_head",)) -> nn.Module: ... +def preallocated_cache(model, *, batch_size=1, max_seq_len=None) -> list[PreallocatedKVCache]: ... + + +# `skip` defaults to the LM head because it is normally tied to the token +# embedding; quantising it would sever the tie and save nothing. + + +# benchmark -- metrics split by how much they can be trusted +def deterministic_report(cfg) -> dict: ... # exact; asserted against a baseline +def measured_report(cfg) -> dict: ... # wall-clock; reproducible in shape only +``` + +## Conventions + +- **Config-driven**: every module reads from the loaded YAML dict (see `configs/default.yaml`); no magic constants. +- **tokenizer representation** is GPT-2 style rendered-unicode (TDR-002). Raw bytes only at the UTF-8 boundary via `token_bytes()`. +- **KV cache** is the canonical inference path; training never uses it. +- **TDD**: write the exit test in `tests/unit/` *before* implementing each stage. +- One module = one public API (exported in its `__init__.py`). diff --git a/docs/TDD.md b/docs/TDD.md new file mode 100644 index 0000000..0f79b94 --- /dev/null +++ b/docs/TDD.md @@ -0,0 +1,1006 @@ +# Optimized GPT-style Language Model Engine Built From Scratch + +**Technical Design Document (TDD)** +**Status:** Draft v1.2 โ€” modern decoder stack (RoPE/RMSNorm/GQA/SwiGLU) adopted; +TDR-006 superseded; module contracts unified with `BUILD_ORDER.md`; vault +relocated (TDR-015 โ€ฆ TDR-019) +**Audience:** Senior ML Engineers, Systems Engineers, Infrastructure Engineers + +> This document specifies the architecture, interfaces, algorithms, and engineering +> strategy for a complete decoder-only transformer language model engine built from +> scratch. It prioritizes performance, scalability, memory efficiency, clean +> architecture, and explicit engineering tradeoffs. It is a design/reference +> document, not a tutorial. + +--- + +## 0. Executive Summary + +We design a **from-scratch**, decoder-only GPT-style transformer language model +system with a production-grade engineering posture. No pretrained weights. No +HuggingFace abstractions. All core components are implemented manually on PyTorch. + +The system is decomposed into seven cooperating modules plus an overarching +benchmarking and testing strategy: + +1. **Custom Tokenizer** โ€” byte-level BPE. +2. **Efficient Dataset Engine** โ€” memory-mapped, streaming, multiparallel feeder. +3. **Transformer Architecture** โ€” decoder-only blocks (embedding, causal + grouped-query attention with RoPE, RMSNorm, SwiGLU FFN, output head). +4. **Attention Optimization** โ€” standard โ†’ block/flash-inspired tiled attention. +5. **Training Engine** โ€” AdamW, LR schedule, mixed precision, checkpointing. +6. **Inference Engine** โ€” autoregressive generation + KV cache + batching + sampling. +7. **Model Optimization Layer** โ€” quantization, memory pooling/allocation. + +Cutting across all modules: a **benchmarking system** and a **TDD test strategy**. + +The project follows a four-phase roadmap: Foundation โ†’ Functional Model โ†’ +Optimization โ†’ Production Engine. Every major decision is recorded as a Technical +Decision Record (TDR) in Section 13. + +--- + +## 1. Custom Tokenizer + +### 1.1 Scope + +A subword tokenizer using **Byte Pair Encoding (BPE)** at the **byte level** +(GPT-2 style). It must convert arbitrary UTF-8 text to a compact sequence of +integer token ids and back losslessly. + +### 1.2 Requirements + +- Vocabulary creation / learning +- Token merging algorithm +- Encoding pipeline (text โ†’ ids) +- Decoding pipeline (ids โ†’ text) +- Special token handling +- Vocabulary serialization (save/load) + +### 1.3 Architecture + +``` +Input text + โ”‚ + โ–ผ encode utf-8 โ†’ byte sequence +Base byte vocabulary (256 tokens: byte 0..255) + โ”‚ apply learned merges greedily + โ–ผ +BPE merges (ranked) โ†’ token ids +``` + +**Internal state:** +- `encoder: Dict[str, int]` โ€” rendered-token โ†’ id. +- `decoder: Dict[int, str]` โ€” id โ†’ rendered-token. +- `merges: List[tuple]` โ€” rank-ordered merge rules. +- `special_tokens: List[str]` โ€” atomic tokens, reserved ids 0..k-1. + +**Rendering table:** A GPT-2-style `bytes_to_unicode` map renders the 256 byte +values to one-to-one printable unicode characters so that every byte sequence is a +valid string key (avoids collisions with special-token text). + +### 1.4 Algorithm + +**Training / merge learning (greedy BPE):** + +1. Initialize each token sequence as its constituent UTF-8 bytes (id = `k + byte`, + where `k` = number of special tokens). +2. Count frequencies of all adjacent pairs. +3. Merge the most frequent pair into a new token id, record its rank. +4. Replace all occurrences of that pair; repeat until the target `vocab_size` is + reached or no pairs remain. + +**Encoding (greedy):** + +1. Detect and reserve special-token substrings (atomic; never merged across). +2. Map remaining bytes to base token ids. +3. Repeatedly merge the adjacent pair with the **lowest learned merge rank** until + no merge applies. + +**Decoding:** + +1. For each id, emit either the special token's literal text or the original bytes + recovered by inverting the rendering table. +2. Concatenate bytes and decode UTF-8 (loss-tolerant with `errors="replace"`). + +### 1.5 Memory Considerations + +- `encoder`/`decoder` dicts are the dominant memory cost; both grow with + `vocab_size`. Each entry is `O(1)`; total is `O(vocab_size)`. +- The byte-decomposition cache (`token_id โ†’ bytes`) lets us render merged tokens + without re-deriving byte sequences; its cost is `O(vocab_size ร— avg_bytes)`. +- Training counts pairs over the corpus each iteration; for very large corpora + this must be streaming / sampled rather than fully resident. + +### 1.6 Performance Optimizations + +Implemented (TDR-020), with measured effect: + +- **Pre-tokenization.** Merges run within pieces of ~1-20 bytes, not over whole + documents, so the quadratic merge term becomes a small constant. +- **Piece cache.** Natural text repeats pieces heavily, so after warmup most of + a corpus costs one dict lookup. Measured: `encode()` 2.9 KB/s โ†’ 67 MB/s. +- **All occurrences per pass.** Each merge pass rewrites every occurrence of the + winning pair, rather than one occurrence followed by a full rescan. +- **Merged id stored in the merge table.** The hot loop never re-renders a token + key or hashes a long string to find the id it just decided on. +- **Incremental training index.** A pair โ†’ containing-pieces index means each + merge revisits only the pieces it affects: O(merges x affected pieces) rather + than O(merges x corpus). Measured: 4096-token vocabulary over 390 KB in 0.4 s. +- **Special tokens split off first**, by regex alternation ordered longest-first, + so the merge loop never branches on them and a special token can never be + shadowed by its own prefix. + +### 1.7 Design Tradeoffs + +| Approach | Pros | Cons | +|----------|------|------| +| Byte-level BPE | No OOV, handles all Unicode/emoji, robust subwords | Slightly more tokens for ASCII than word-level | +| Word-level BPE / SentencePiece-style | Cleaner for Latin scripts | OOV handling and Unicode complexity | + +--- + +## 2. Efficient Dataset Engine + +### 2.1 Scope + +A high-performance data pipeline that feeds training batches to the model without +loading the entire corpus into RAM and without stalling the compute device. + +### 2.2 Requirements + +- Memory-mapped dataset loading +- Streaming dataset support +- Multiprocessing workers +- Asynchronous prefetching +- Pinned-memory transfer (GPU) +- Efficient batching +- Configurable shuffling + +### 2.3 Data Flow Architecture + +``` +Dataset Storage (binary token files) + โ”‚ open (os page cache) + โ–ผ +Memory Mapping (mmap) โ€” virtual address space, no bulk RAM copy + โ”‚ __getitem__(idx) + โ–ผ +Worker Processes (num_workers=N, fork/spawn) + โ”‚ decode + batch construction + โ–ผ (async) +Prefetch Queue (DataLoader `prefetch_factor`) + โ”‚ .to(device, non_blocking=True) on pinned tensors + โ–ผ +GPU Training Batch (x, y) +``` + +### 2.4 Components + +- **Memory-Mapped Reader:** token ids stored as a flat array of `uint16` + (`numpy.memmap` or `torch.from_file`). Only pages actually read are brought into + memory by the OS. +- **TokenDataset:** `__len__` = number of usable `seq_len` windows; `__getitem__` + reads a contiguous slice of `seq_len + 1` ids and splits it into + `{"input_ids": s[:-1], "labels": s[1:]}` โ€” the next-token shift happens at the + window level, so no `collate` fn is needed and the default collate batches the + dict straight to `[B, T]`. Reading `T+1` ids (rather than `T` and shifting + within the window) means the label for the final position is a real token + instead of padding. +- **Multiprocessing workers:** each worker owns an independent memmap handle, + avoiding GIL contention. +- **Prefetching:** `DataLoader(num_workers, prefetch_factor, pin_memory, + non_blocking=True)` keeps the GPU fed. +- **Shuffling:** window-index shuffling (shuffle the starting offsets) rather than + shuffling live tensors; for streaming, reshuffle epoch boundaries. + +### 2.5 CPU โ†” GPU Interaction & Memory Management + +- Pinned host buffers allow `H2D` copies without intermediate staging. +- Prefetching overlaps I/O, tokenization, and device transfer with compute. +- The memmap uses OS page cache, so frequently accessed regions benefit from + kernel caching (re-warm across epochs cheaply). + +### 2.6 Bottlenecks & Optimization Opportunities + +| Bottleneck | Mitigation | +|------------|-----------| +| Disk latency | mmap (page cache), OS-level prefetch, SSDs | +| Python decode/encode | vectorized byte mapping, C-accelerated ops | +| GIL in workers | multiprocessing over threads | +| Device idle | prefetch + async pinned H2D | +| Memory spikes | mmap (no bulk load), fixed `block_size` | + +--- + +## 3. Transformer Architecture + +A standard decoder-only transformer. All tensors are `(B, T, C)` for batch, +sequence, and channel dims. + +### 3.1 High-Level Stack + +``` +tok ids (B, T) + โ”‚ +Token Embedding (V โ†’ C) โ† no positional embedding table; position + โ”‚ enters inside attention via RoPE +N ร— TransformerBlock: + โ”œโ”€ RMSNorm โ†’ CausalAttention (GQA + RoPE) โ†’ Residual + โ”œโ”€ RMSNorm โ†’ SwiGLU MLP (Cโ†’d_ffโ†’C) โ†’ Residual + โ”‚ +Final RMSNorm + โ”‚ +LM Head (C โ†’ V) โ†’ logits (B, T, V) +``` + +This is the modern Llama-style decoder stack (RoPE + RMSNorm + GQA + SwiGLU) +rather than the original GPT-2 stack. See TDR-015/016/017 for the rationale; +TDR-006 (learnable absolute positions) is superseded. + +### 3.2 Embedding Layer + +- **Token embeddings:** `nn.Embedding(vocab_size, d_model)`. +- **No positional embedding table.** Position is injected inside attention by + **Rotary Position Embeddings (RoPE)**, which rotate Q and K by a + position-dependent angle so that attention scores depend only on *relative* + offset. This removes a `max_seq_len ร— d_model` parameter table and permits + context extension by rescaling `rope_theta` (TDR-015). +- **Weight tying:** the token embedding matrix is optionally reused as the LM + head, saving `vocab_size ร— d_model` parameters. + +### 3.3 Self-Attention (Grouped-Query + RoPE) + +- **Q/K/V projections:** `q_proj: Linear(C, n_headsยทd_head)`, + `k_proj`/`v_proj: Linear(C, n_kv_headsยทd_head)`. Q and KV have *different* + output widths under GQA, so the fused `Linear(C, 3C)` of the GPT-2 design does + not apply; the three projections stay separate. +- **Grouped-Query Attention:** `n_kv_heads < n_heads`; each KV head is shared by + `n_heads / n_kv_heads` query heads (must divide evenly). Shrinks the KV cache + by that same factor โ€” the dominant memory term during decoding (TDR-016). +- **RoPE:** applied to Q and K (never V) after projection and head-splitting, + using the absolute position of each token. During cached decoding the position + offset comes from the current cache length, not from the slice index. +- **Causal masking:** future positions masked to `-inf`. The `sdpa` path uses + `is_causal=True`; the `manual` path builds an explicit upper-triangular mask. + During single-token cached decode the query attends over the whole cache and + no mask is needed. +- **Scaling:** `scores /= sqrt(d_head)`. +- **Output projection:** `nn.Linear(n_headsยทd_head, C)`. + +### 3.4 Feed-Forward Network (SwiGLU) + +- `SwiGLU(x) = down( silu(gate(x)) * up(x) )` โ€” three projections + (`gate`, `up`: `C โ†’ d_ff`; `down`: `d_ff โ†’ C`) rather than the two of a GELU MLP. +- Because SwiGLU uses 3 matrices instead of 2, `d_ff` is conventionally set to + ~`8/3 ยท C` to hold the parameter count level with a `4C` GELU MLP. The default + config uses `d_model=512, d_ff=2048` (a `4ร—` ratio), trading ~1.5ร— the FFN + parameters for the gating quality win at this small scale (TDR-017). + +### 3.5 Transformer Block + +- **Pre-normalisation** with **RMSNorm**: `x = x + attn(norm(x))`, then + `x = x + mlp(norm(x))`. Pre-norm improves deep-stack trainability and reduces + warmup sensitivity. +- **RMSNorm** drops the mean-centring and bias of LayerNorm, keeping only the + scale: `x ยท rsqrt(mean(xยฒ) + eps) ยท weight`. Fewer ops, one fewer reduction + pass, no measured quality loss at this scale (TDR-017). +- **Residual connections** around both sublayers. + +### 3.6 Output Layer + +- `nn.Linear(d_model, vocab_size)` produces logits (optionally weight-tied to the + token embedding). +- **Training:** cross-entropy over vocab at each position. +- **Inference:** softmax over the last position โ†’ sampling distribution. + +### 3.7 Configuration + +Config is a plain nested dict loaded from `configs/*.yaml` (no dataclass; every +module reads the dict, per the BUILD_ORDER "config-driven" convention). The +`model` block is authoritative: + +`d_model`, `n_layers`, `n_heads`, `n_kv_heads`, `d_ff`, `vocab_size`, +`max_seq_len`, `rope_theta`, `dropout`. + +The `attention` block selects the kernel path: `impl` (`sdpa | flash | manual`) +and `kv_cache`. + +Sizes target ~10โ€“40M parameters so early iterations run comfortably on CPU and +on Apple Silicon MPS (TDR-019). + +--- + +## 4. Attention Optimization + +### 4.1 Standard Attention + +``` +Attention(Q,K,V) = softmax( Q Kแต€ / sqrt(d) ) V +``` + +- **Compute complexity:** `O(Tยฒ ยท d)` per head, summed over heads. +- **Memory complexity:** `O(Tยฒ ยท n_head)` for the full score tensor. +- **Limitation:** at long contexts, the `Tยฒ` score matrix dominates memory and + bandwidth, becoming the primary scalability bottleneck. + +### 4.2 Optimized / Flash-Attention-Inspired Attention (Tiled) + +Algorithm: process the sequence in fixed-size **blocks**; maintain online softmax +running max `m` and running sum `l`; rescale the accumulated output when a new max +is found; never materialize the full `Tยฒ` matrix. + +- **Memory:** `O(T)` (only output + per-block buffers), regardless of sequence + length; only the current `(block_b, block_k)` tile is materialized. +- **Performance:** improved cache locality (IO-aware), lower bandwidth pressure. +- **Numerical:** online softmax with running max keeps values stable. + +**Implementation challenges:** +- Online rescaling adds complexity (must carry `m`, `l` accumulators). +- Backward pass requires recomputing scores or caching `P = softmax(QKแต€)` per + block. **We recompute** (TDR-021): the forward saves only the per-row + log-sum-exp, and backward rebuilds each score tile from it. Relying on + autograd instead silently restores O(Tยฒ), because every tile's intermediates + are retained. +- Block granule tuning (`BLOCK_M`, `BLOCK_N`) and boundary padding. + +**Tradeoff matrix:** + +| Scheme | Memory | Complexity | Notes | +|--------|--------|-----------|-------| +| Standard | O(Tยฒ) | naive | simple | +| Block/tiled (flash-inspired) | O(T) | tiled, online softmax + recomputing backward | exact, memory-scalable | + +Measured retained-activation bytes for a 4-head, 64-dim layer (the figure that +matters for training, not just inference): + +| T | standard | tiled | ratio | +|---|---|---|---| +| 128 | 800 KiB | 306 KiB | 0.38ร— | +| 256 | 2,640 KiB | 596 KiB | 0.23ร— | +| 512 | 9,488 KiB | 1,176 KiB | 0.12ร— | +| 1024 | 35,856 KiB | 2,336 KiB | **0.07ร—** | + +Growth per doubling: **1.97โ€“1.99ร—** tiled (linear) versus **3.3โ€“3.8ร—** standard +(quadratic). + +### 4.3 Selectable Kernel Paths (`attention.impl`) + +The config exposes three interchangeable implementations behind one module +interface, so the same model can be run on whichever path the hardware favours +and the benchmark layer can compare them directly (TDR-012): + +| `impl` | Path | Use | +|--------|------|-----| +| `manual` | explicit `QKแต€`, mask, softmax, `PV` in PyTorch ops | the readable reference; the correctness oracle every other path is tested against | +| `sdpa` | `torch.nn.functional.scaled_dot_product_attention` | **default.** Dispatches to the best fused backend available (including on Apple MPS); `is_causal=True` avoids materialising the mask | +| `flash` | our own tiled/online-softmax implementation (ยง4.2) | the from-scratch demonstration that the `Tยฒ` matrix need never be materialised | + +All three must produce numerically equivalent output (within fp tolerance) on +the same inputs โ€” this is a regression test, not an aspiration (ยง10.4). + +--- + +## 5. Training Engine + +### 5.1 Training Loop + +``` +for step: + x, y = next(batch) # dataset + collate + with autocast(dtype): logits, loss = model(x, y) + loss_scaler.scale(loss).backward() + grad_norm = clip_grad_norm_(max_norm) + loss_scaler.step(optimizer); update; zero_grad() +``` + +### 5.2 Optimization + +- **Custom AdamW:** decoupled weight decay, bias correction, `eps`. +- **LR schedule:** cosine decay with linear warmup. +- **Gradient clipping:** global norm clipping for stability. +- **Warmup:** prevents early instability from large initial steps. + +### 5.3 Training Improvements + +- **Mixed precision:** `torch.autocast` + `GradScaler` (fp16/bf16). CPU reduces to + fp32 (no-op path preserved so GPU code is portable). +- **Gradient accumulation:** accumulate over `K` micro-batches to emulate a larger + effective batch on limited memory. +- **Checkpointing:** periodic save of model, optimizer, scaler, step, rng state. +- **Experiment tracking:** metrics reported to console/CSV/optional tracker. + +### 5.4 Stability, Performance, Scalability + +- Pre-LN + warmup + gradient clipping together stabilize early training. +- Mixed precision + pinned prefetch both raise throughput. +- Gradient accumulation + batching let batch size scale independent of device memory. +- Target: profile each stage to confirm where time is spent (forward/backward, + dataloader, optimizer, H2D transfer); set throughput targets relative to a + measured CPU baseline before claiming GPU/AMP gains. + +--- + +## 6. Inference Engine + +### 6.1 Text Generation + +Autoregressive loop: encode prompt โ†’ loop `max_new_tokens` โ†’ sample next token โ†’ +append โ†’ decode. + +**Sampling strategies:** +- **Greedy:** `argmax`. +- **Temperature:** `logits / temperature` before softmax. +- **Top-k:** restrict to top-k logits. +- **Top-p (nucleus):** smallest set with cumulative prob โ‰ฅ `p`. + +### 6.2 KV Cache + +``` +Without cache: every generated token recomputes attention over all prior tokens โ†’ O(Tยฒ) cumulative +With cache: keys/values from prior tokens are reused โ†’ O(T) total +``` + +- **Prefill:** encode prompt, compute & store all K/V. +- **Decode:** per new token compute only `q` and new `k/v`, append to cache, + attend over full cached `K`, `V`. + +**Memory tradeoff:** cache grows `2 * n_layer * n_head * d_head * T` floats per +request (linear in context length); mitigates with cache clearing/paging (future). + +**Latency improvement:** transforms decoding from compute-bound into +memory/cache-bandwidth-bound, a large constant-factor speedup. + +### 6.3 Batching + +- Stack multiple requests; each request maintains its own KV cache in the batch dim. +- Improves device utilization and aggregate throughput. + +### 6.4 Measured Metrics + +First-token latency, per-token latency, tokens/sec, throughput, memory footprint. + +--- + +## 7. Model Optimization Layer + +### 7.1 Quantization + +Compare **FP32** (baseline), **FP16**, and **INT8**. + +| Precision | Memory | Speed (mixed HW) | Accuracy | +|-----------|--------|-------------------|----------| +| FP32 | 4 bytes/param | baseline | reference | +| FP16/BF16 | 2 bytes/param | faster w/ mixed HW | near-lossless | +| INT8 | 1 byte/param | fastest | some loss, needs calibration | + +> **Hardware caveat:** On CPU-only validation (our primary portable path), fp16/int8 +> may **not** be faster than fp32. These entries assume mixed-precision-capable +> hardware (e.g., GPU with Tensor Cores). The benchmarking layer measures both to +> validate claims per-device rather than assume a global winner. + +**Approach:** FP16 via `.half()`/autocast; INT8 via post-training quantization with +per-tensor/per-channel scales and calibration on a reference split. + +### 7.2 Memory Optimization + +- **Tensor reuse:** avoid temporaries (in-place where safe), fused projections. +- **Memory pooling:** preallocated key/value caches, buffer reuse across blocks. +- **Allocation strategy:** modeled on CUDA/PyTorch caching allocator ideas โ€” + arena-based pooling so repeated same-shape allocations reuse blocks instead of + hitting `malloc`/`cudaMalloc`. + +--- + +## 8. Low-Level System Design + +### 8.1 High-Level Architecture Diagram + +``` +Tokenizer + โ”‚ (token ids) + โ–ผ +Dataset Engine โ”€โ”€ training batches โ”€โ”€โ–บ Training Engine + โ”‚ โ”‚ (weights) + โ–ผ โ–ผ +Transformer Model โ—„โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ Optimization Layer + โ”‚ โ”‚ + โ–ผ โ–ผ +Inference Engine โ—„โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ (optimized weights) + โ”‚ + โ–ผ +API Layer (programmatic; optional serving interface later) +``` + +### 8.2 Module Responsibilities + +| Module | Responsibility | Inputs | Outputs | Depends on | +|--------|----------------|--------|---------|------------| +| tokenizer | text โ†” token ids | str / list[int] | list[int] / str | config | +| dataset | training batch production | binary token file path | (x, y) batch tensors | tokenizer | +| model | forward, logits, loss | (B,T) token ids | (B,T,V) logits / loss | attention, config | +| attention | causal attention compute | q, k, v | attended output | โ€” | +| training | weight updates | batches, model | weights, metrics | model, optimizer | +| inference | generation | prompt, params | generated text | model, KV cache | +| optimization | reduce footprint/latency | model | quantized/pooled model | model | +| benchmark | measure performance | model, data | metric report | all | + +### 8.3 Class Design + +``` +tokenizer.BPE +Dataset.TokenDataset, Dataset.PackedTokenDataset, Dataset.DataEngine +Attention.RotaryEmbedding, Attention.CausalAttention, Attention.KVCache +Model.TransformerBlock, Model.RMSNorm, Model.SwiGLU, Model.CausalLM +Training.AdamW, Training.CosineScheduler, Training.Trainer, Training.Checkpoint +Inference.Sampler, Inference.Generator +Optimization.quantize +Benchmark.TrainingBenchmark, Benchmark.InferenceBenchmark +``` + +### 8.4 Internal API Design + +These are the canonical contracts, identical to those in +[`BUILD_ORDER.md`](BUILD_ORDER.md). Change them only via a new TDR. + +```python +tokenizer.BPE.encode(text) -> list[int] +tokenizer.BPE.decode(ids) -> str + +# dataset โ€” yields dicts, not tuples, so fields can be added without breaking callers +dataset.TokenDataset.__iter__() -> Iterator[dict] # {"input_ids": LongTensor[T], + # "labels": LongTensor[T]} +dataset.DataEngine.next_batch(device) -> dict # same keys, batched [B, T] + +# attention +attention.CausalAttention.forward(x, *, kv_cache=None, use_cache=False) + -> (out, new_kv_cache) + +# model +model.CausalLM.forward(input_ids, *, kv_cache=None, use_cache=False) -> logits [B,T,V] +model.CausalLM.from_config(cfg: dict) -> CausalLM + +# training / inference / optimization +training.Trainer.step() -> dict # metrics +inference.sample(logits, *, temperature=1.0, top_k=0, top_p=1.0) -> int +inference.generate(model, tokenizer, prompt, *, max_new_tokens, **sampling) -> str +optimization.quantize(model, *, bits=8, scheme="int8_weight") -> nn.Module +``` + +> **Loss placement.** `CausalLM.forward` returns logits only; cross-entropy lives +> in the training loop. This keeps the model pure for inference (where computing +> a loss is wasted work) and is why the signature differs from the +> `GPT.forward(x, y) -> (logits, loss)` sketch in earlier revisions of this +> document. + +--- + +## 9. Benchmarking System + +Structure under `benchmark/`: + +``` +benchmark/ +โ”œโ”€โ”€ training_benchmark.py # tokens/sec, throughput, GPU/CPU util, memory +โ”œโ”€โ”€ inference_benchmark.py # first-token latency, gen latency, tokens/sec, mem +โ”œโ”€โ”€ memory_benchmark.py # footprint FP32 vs FP16 vs INT8, KV cache size +โ””โ”€โ”€ latency_test.py # per-token + end-to-end latency distributions +``` + +Metrics (training): tokens/sec, throughput, GPU utilization, CPU utilization, +peak memory (via `torch.cuda.max_memory_allocated` or `tracemalloc`). + +Metrics (inference): first-token latency, generation latency, tokens/sec, +aggregate throughput, memory footprint (model + KV cache). + +--- + +## 10. Test Driven Development Strategy + +### 10.1 Unit Tests + +- **Tokenizer:** encoding correctness, decoding correctness, round-trip + consistency, vocabulary consistency (encoder/decoder inverse), special tokens, + serialization round-trip. +- **Attention:** tensor shape `(B, H, T, d_head)`, causal-mask correctness, + gradient flow (finite-difference check or `loss.backward()` presence). +- **Dataset:** batch shapes, window slicing, memory loading, offset consistency. +- **AdamW/optimizer:** step reduces loss, clip bounds norm. + +### 10.2 Integration Tests + +- tokenizer + dataset (tokens flow correctly) +- dataset + model (batch shapes through forward) +- model + training loop (loss decreases on a tiny dataset) +- inference pipeline (generate reproduces a deterministic greedy prefix) + +### 10.3 Performance Tests + +- Validate throughput/memory/latency meet targets and that optimizations improve + baselines (gated thresholds). + +### 10.4 Regression Tests + +- Optimization (flash attention, quantization, KV cache) does **not** degrade + final loss/accuracy beyond a tolerance vs the reference implementation. +- Map to module first, then entire suite re-run. + +**Workflow:** every module: write failing test โ†’ implement โ†’ refactor. See +`~/.claude/vault/language-model/roadmap/` for per-phase checklists. + +--- + +## 11. Repository Structure + +``` +language-model-engine/ +โ”œโ”€โ”€ tokenizer/ # custom BPE tokenizer +โ”œโ”€โ”€ dataset/ # memory-mapped, streaming data engine +โ”œโ”€โ”€ model/ # transformer architecture (embedding, blocks, head) +โ”œโ”€โ”€ attention/ # standard + flash-inspired attention +โ”œโ”€โ”€ training/ # trainer, AdamW, schedulers, checkpointing +โ”œโ”€โ”€ inference/ # generator, KV cache, sampling, batching +โ”œโ”€โ”€ optimization/ # quantization, memory pooling +โ”œโ”€โ”€ benchmark/ # training/inference/memory/latency benchmarks +โ”œโ”€โ”€ tests/ # pytest suite (unit + integration) +โ”œโ”€โ”€ configs/ # runnable JSON/configuration files +โ”œโ”€โ”€ checkpoints/ # serialized model/optimizer state +โ”œโ”€โ”€ docs/ # TDD, README, module docs +โ””โ”€โ”€ .github/ # GitHub Actions CI workflows (lint, type, test, smoke) +``` + +Purpose of each directory summarized in the tree above; `tests/` mirrors the +module tree for maintainability. `.github/workflows/` holds the CI definitions +(see ยง16.3). + +The command-line entry point is `training/train.py`, invoked as a module +(`python -m training.train`) rather than a root-level `train.py`, so it resolves +imports through the same package paths as every other module. + +--- + +## 12. Development Roadmap + +### Phase 1 โ€” Foundation +Tokenizer (train/encode/decode/serialize) โ†’ dataset engine (mmap, workers, +prefetch, batching) โ†’ simple transformer forward pass (embedding, causal attention, +MLP, head). *Exit:* all unit tests passing; forward shapes correct. + +### Phase 2 โ€” Functional Model +Training engine (loop, AdamW, scheduler, warmup, grad-clip, AMP, checkpointing) โ†’ +basic autoregressive generation. *Exit:* loss converges on a smoke dataset; +generation produces coherent continuations. + +### Phase 3 โ€” Optimization +KV cache โ†’ quantization (fp16/int8) โ†’ flash-inspired attention โ†’ memory +optimization โ†’ inference batching. *Exit:* quantitative speedup/memory reductions +measured vs Phase 2 baseline. + +### Phase 4 โ€” Production Engine +Benchmarking framework โ†’ API layer โ†’ documentation โ†’ full integration + regression +suite. *Exit:* reproducible benchmark report; documentation complete. + +--- + +## 13. Technical Decision Records (TDRs) + +### TDR-001: Decoder-Only Architecture +- **Decision:** Decoder-only (GPT-style) causal transformer. +- **Reason:** Native fit for autoregressive next-token generation; modern SOTA + architecture class. +- **Alternatives:** encoder-only (BERT), encoder-decoder (T5). +- **Tradeoffs:** Simpler unified stack vs no bidirectional context; sequential + decoding (mitigated by KV cache). + +### TDR-002: Byte-Level BPE Tokenizer +- **Decision:** GPT-2-style byte-level BPE with a `bytes_to_unicode` rendering table. +- **Reason:** Zero OOV for arbitrary UTF-8; robust subwords; well-understood. +- **Alternatives:** word-level BPE, SentencePiece unigram, WordPiece. +- **Tradeoffs:** Slightly larger token counts on Latin text vs OOV/Unicode ease. + +### TDR-003: PyTorch Manual Implementation +- **Decision:** Implement all layers/training manually on PyTorch; no HF/pretrained. +- **Reason:** Autograd and device management provided while we own all ML logic. +- **Alternatives:** NumPy (hard autograd), JAX. +- **Tradeoffs:** Development speed vs authenticity/control; explicitly mandated. + +### TDR-004: Memory-Mapped Dataset Pipeline +- **Decision:** mmap + workers + prefetch for training data. +- **Reason:** Avoids RAM blowup on large corpora; overlaps I/O with compute. +- **Alternatives:** full in-memory load, TFRecord-style shards. +- **Tradeoffs:** OS page cache dependency vs low memory ceiling and streaming. + +### TDR-005: Pre-LayerNorm Transformer +- **Decision:** Pre-LN (normalize-before-sublayer). +- **Reason:** More stable deep training, less warmup sensitivity. +- **Alternatives:** Post-LN, original transformer. +- **Tradeoffs:** Pre-LN may require weighting for very deep stacks; simpler anyway. + +### TDR-006: Learnable Absolute Positional Embeddings โ€” โš ๏ธ SUPERSEDED by TDR-015 +- **Decision:** Learnable absolute position table for Phase 1. +- **Reason:** Simplest correct baseline; acceptable within the fixed `block_size`. +- **Alternatives:** sinusoidal, Rotary, ALiBi. +- **Tradeoffs:** No extrapolation beyond `block_size`. +- **Superseded:** never implemented. RoPE was adopted before any positional code + was written; see TDR-015. Retained for history. + +### TDR-007: Flash-Attention-Inspired Tiled Attention +- **Decision:** Block-tiled attention with online softmax. +- **Reason:** Breaks the O(Tยฒ) memory wall for long contexts; IO-aware. +- **Alternatives:** Standard attention, true FlashAttention kernel. +- **Tradeoffs:** More complex bookkeeping vs large, exact memory savings. + +### TDR-008: KV Cache for Inference +- **Decision:** Per-layer K/V cache with prefill-then-decode. +- **Reason:** Eliminates O(Tยฒ) recomputation; large decode speedup. +- **Alternatives:** full recompute, custom kernels. +- **Tradeoffs:** Linear memory growth vs large latency/throughput win. + +### TDR-009: Custom AdamW + Cosine Schedule with Warmup +- **Decision:** AdamW with decoupled weight decay; cosine decay + linear warmup; + gradient clipping. +- **Reason:** Standard, robust, well-calibrated optimizer recipe. +- **Alternatives:** SGD + momentum, LAMB/LARS, fixed LR. +- **Tradeoffs:** Slightly more hyperparams vs strong convergence/stability defaults. + +### TDR-010: LayerNorm โ†’ FP16/BF16 Mixed Precision with Loss Scaling +- **Decision:** autocast + `GradScaler` for fp16/bf16; CPU degrades to fp32. +- **Reason:** Memory/speed gains on mixed-precision HW with near-lossless fp16. +- **Alternatives:** pure FP32, full bf16. +- **Tradeoffs:** FP16 needs scaled loss (grad underflow) and master weights on some + setups. + +### TDR-011: Global Knowledge Vault +- **Decision:** Maintain a global project knowledge vault at `~/.claude/vault/language-model/` (outside + the repo), as single source of truth for decisions, research, benchmarks. +- **Reason:** Long-lived engineering memory independent of a single codebase; + reusable across projects. +- **Alternatives:** in-repo `vault/`, wiki, issue docs. +- **Tradeoffs:** Slight indirection (not in repo history) vs global, durable, + cross-project knowledge. + +### TDR-012: Benchmarking Gated by Baselines +- **Decision:** Every optimization must demonstrate measurable improvement vs a + stated baseline before acceptance. +- **Reason:** Prevents regressions and validates engineering claims. +- **Alternatives:** ad-hoc manual measurement. +- **Tradeoffs:** Slight overhead vs accountable, quantifiable engineering. + +### TDR-013: Locked `main` + PR-Gated Git Workflow +- **Decision:** `main` is locked; all work on feature branches, merged only via + reviewed PRs. Issues tracked in GitHub Issues. +- **Reason:** Keeps `main` production-stable, enables review, auditable history. +- **Alternatives:** direct-to-main, full git-flow. +- **Tradeoffs:** Slight overhead (branches/PRs) vs strong safety and review gate. + See ยง16 and `~/.claude/vault/language-model/decisions/architecture_decisions.md` (ADR-003). + +### TDR-014: Free-Tier CI (Lint + Type + Tests) Gate +- **Decision:** GitHub Actions free-tier CI enforcing lint, type check, tests, and + import smoke on every PR and push to `main`; CI green is a merge requirement. +- **Reason:** Automated, enforced quality gate at zero cost on a public repo. +- **Alternatives:** local-only checks, paid/self-hosted CI, pre-commit alone. +- **Tradeoffs:** Runner resource limits vs encrypted, reproducible enforcement. + See ยง16.3 and `~/.claude/vault/language-model/decisions/architecture_decisions.md` (ADR-004). + +### TDR-015: Rotary Position Embeddings (supersedes TDR-006) +- **Decision:** Inject position via RoPE applied to Q and K inside attention; no + positional embedding table. +- **Reason:** Attention scores become a function of *relative* offset, which + generalises better than absolute indices; removes a `max_seq_len ร— d_model` + parameter table; context can be extended post-hoc by rescaling `rope_theta` + rather than retraining a table. It is also the current standard (Llama, Qwen, + Mistral), so the implementation is a well-understood reference point. +- **Alternatives:** learnable absolute (TDR-006), sinusoidal, ALiBi. +- **Tradeoffs:** Position handling moves into the attention hot path and must be + correct under KV caching (the offset comes from cache length, not slice index) + โ€” an easy bug, hence a dedicated test. Cos/sin tables are cached as + non-persistent buffers. + +### TDR-016: Grouped-Query Attention +- **Decision:** `n_kv_heads < n_heads` (default 2 KV heads for 8 query heads); + each KV head is shared across a group of query heads. +- **Reason:** The KV cache โ€” not the weights โ€” dominates decode-time memory, and + it scales with `n_kv_heads`. A 4:1 group ratio cuts cache footprint 4ร— for a + quality cost that is negligible at this scale. This matters directly on the + target hardware, where unified memory is shared with the rest of the system + (TDR-019). +- **Alternatives:** full MHA (`n_kv_heads == n_heads`), MQA (`n_kv_heads == 1`). +- **Tradeoffs:** Requires `n_heads % n_kv_heads == 0` and a repeat/expand of KV + heads before the attention product; MHA remains reachable by setting + `n_kv_heads = n_heads`, so the baseline is one config key away and directly + benchmarkable. + +### TDR-017: RMSNorm + SwiGLU +- **Decision:** RMSNorm instead of LayerNorm; SwiGLU FFN instead of GELU MLP. +- **Reason:** RMSNorm removes mean-centring and bias โ€” one fewer reduction pass + and fewer parameters, with no measured quality loss. SwiGLU's multiplicative + gate consistently outperforms a plain GELU MLP at equal parameter count. +- **Alternatives:** LayerNorm + 4C GELU MLP (the GPT-2 baseline). +- **Tradeoffs:** SwiGLU needs 3 weight matrices instead of 2, so equal-parameter + comparison requires `d_ff โ‰ˆ 8/3ยทC`. We keep `d_ff = 4ยทC` at this scale and + accept ~1.5ร— FFN parameters, noting the deviation rather than hiding it. + +### TDR-018: Vault Location โ€” `~/.claude/vault/language-model/` +- **Decision:** The knowledge vault mandated by TDR-011 lives at + `~/.claude/vault/language-model/`, not `~/vault/`. +- **Reason:** `~/.claude/vault/` is the machine's existing global vault root and + is already covered by a standing maintenance rule; adding a second vault at + `~/vault/` would split the source of truth. Every `~/vault/...` reference in + this repo pointed at a directory that did not exist. +- **Alternatives:** create `~/vault/` as originally specified; move the vault + in-repo under `docs/`. +- **Tradeoffs:** Still outside git (the TDR-011 tradeoff is unchanged), but now + it actually exists and has one unambiguous location. + +### TDR-019: Apple Silicon (MPS) as the Primary Development Target +- **Decision:** Develop and validate on Apple Silicon via the `mps` backend, + with CPU as the always-correct fallback and CUDA as an untested-but-supported + path. Device selection is `mps โ†’ cuda โ†’ cpu`, overridable by config. +- **Reason:** It is the hardware this project is actually built on, so it is the + only device whose numbers we can honestly report. Verified available: + torch 2.13, bf16 tensors, and fused `scaled_dot_product_attention`. +- **Alternatives:** CPU-only development; rented CUDA. +- **Tradeoffs:** MPS has real gaps โ€” no `torch.cuda.max_memory_allocated` + equivalent for peak-memory metrics, `GradScaler` is a CUDA concept (bf16 on + MPS needs no loss scaling), and some ops silently fall back to CPU. The + benchmark layer must therefore report *per-device* numbers and never present + an MPS measurement as a general claim (TDR-012). + +### TDR-020: Pre-tokenization Before BPE Merges +- **Decision:** Split text on a GPT-2-style regex before merging, and never let + a merge cross a piece boundary. Cache piece โ†’ ids. +- **Reason:** Two things, one of which is a hard blocker. (1) **Tractability.** + Merging is quadratic in the length of the sequence it runs over; applied to + whole documents, `encode()` measured 4.0x slower per 2x input โ€” 50 KB/s at + 1 KB falling to 2.9 KB/s at 16 KB โ€” which makes any real corpus unusable + (a 515 MB corpus would not finish). Restricted to pieces of ~1-20 bytes with + a cache, the same path measures **67 MB/s**, and training a 4096-token + vocabulary dropped from ~10 s per 117 KB to 0.4 s per 390 KB. (2) **Quality.** + It stops the vocabulary being spent on cross-word artefacts like `"e c"`. + Compression *improved* as a side effect, 3.51 โ†’ 4.14 chars/token. +- **Alternatives:** keep whole-document merging and chunk the input (caps the + quadratic term but leaves ~48 h for 515 MB, and chunk boundaries become + arbitrary token boundaries); rewrite the merge loop with a pair heap but no + pre-tokenization (fixes speed, not quality, and the cache stays invalid). +- **Tradeoffs:** This **changes tokenizer output** โ€” it is not a transparent + optimization. Vocabularies trained before it are not interchangeable with + vocabularies trained after it, and `save()` gained a `version` field so a + stale file is identifiable. The pattern must also *tile the input exactly*: a + gap silently drops characters, so `tiles_exactly()` and + `test_pattern_tiles_every_input` assert coverage directly rather than trusting + round-trip tests to notice. Finally, the stdlib `re` module lacks `\p{L}`, so + the letter class is `[^\W\d]`, which groups underscore with letters โ€” a + deliberate deviation from GPT-2 taken to avoid a third-party `regex` + dependency. + +### TDR-021: Recomputing Backward for the Tiled Attention Kernel +- **Decision:** Implement the tiled kernel as a `torch.autograd.Function` that + saves only `(q, k, v, out, logsumexp)` and **recomputes** the score tiles in + the backward pass. +- **Reason:** Without it the O(T) memory claim was false in training and I had + published it anyway. As plain PyTorch ops, autograd retained every tile's + intermediates, and the tiles summed back to O(Tยฒ) โ€” measured, the tiled path + retained **1.30ร— more** than the naive one at T=512 and both grew + quadratically. The advantage existed only under `no_grad`. Saving the per-row + log-sum-exp (O(T), one scalar per query) is enough to reconstruct the exact + probabilities in backward, so the Tร—T matrices `p`, `dp` and `ds` exist one + tile at a time. After the change: retained memory grows **1.97โ€“1.99ร— per 2ร— T** + (linear) against **3.3โ€“3.8ร—** for the naive path, and at T=1024 it retains + 2.3 MiB versus 35.9 MiB โ€” 15ร— less. +- **Alternatives:** retract the claim and document the kernel as inference-only + (cheap, but leaves a tiled kernel with no reason to exist during training); + `torch.utils.checkpoint` around standard attention (recomputes the *whole* + attention rather than per-tile, and still materialises Tยฒ transiently). +- **Tradeoffs:** Backward now costs a second pass over the tiles, so it trades + compute for memory โ€” the standard FlashAttention bargain. The custom Function + is not double-differentiable (no `create_graph` support), which is fine for + first-order training but would need `once_differentiable` semantics or a + hand-written double backward otherwise. Correctness can no longer be inferred + from autograd, so it is pinned against the manual oracle in float64 + (`test_flash_gradients_match_the_manual_oracle`, agreeing to ~1e-15). + +--- + +## 14. Project Knowledge Vault + +Per the project requirements, all long-lived decisions, architecture, research, +experiments, roadmap, and implementation notes are stored in a **global**, modular +knowledge vault at **`~/.claude/vault/language-model/`** (available across projects): + +``` +~/.claude/vault/language-model/ +โ”œโ”€โ”€ architecture/ system_design, model_architecture, component_design, data_flow +โ”œโ”€โ”€ decisions/ architecture_decisions, optimization_decisions, tradeoffs +โ”œโ”€โ”€ research/ transformer_notes, attention_optimization, inference_optimization, papers +โ”œโ”€โ”€ implementation/ coding_guidelines, module_interfaces, development_notes +โ”œโ”€โ”€ experiments/ benchmarks, training_runs, optimization_results +โ”œโ”€โ”€ roadmap/ milestones, completed_features, future_improvements +โ””โ”€โ”€ project_context.md (overview, repo path, vault location, status) +``` + +**Rules:** (1) review vault before significant tasks; (2) update vault after +significant tasks; (3) capture new ideas in research/design docs; (4) keep it +current, marking stale info deprecated; (5) base recommendations on prior +decisions, noting alignment/conflict. + +--- + +## 15. Deliverables Summary + +- Complete from-scratch, decoder-only GPT engine on PyTorch. +- Custom BPE tokenizer; mmap multiparallel dataset engine. +- Standard + flash-inspired tiled attention. +- Training engine (AdamW, scheduler, AMP, checkpointing). +- Inference engine (KV cache, sampling, batching). +- Optimization layer (fp16/int8 quantization, memory pooling). +- Benchmarking framework and TDD test suite. +- Global knowledge vault at `~/.claude/vault/language-model/`. + +--- + +## 16. Engineering Process: Git Workflow, CI, and Issue Tracking + +We operate this project as a **production-grade, real-world codebase**. The +following engineering process governs how work is delivered. + +### 16.1 Branch Strategy & Locked `main` + +- **`main` is locked.** It is never committed to directly. +- All work happens on **feature branches** (e.g., `feat/tokenizer-bpe`, + `fix/dataset-mmap`, `opt/flash-attention`). +- A branch is merged to `main` **only via a Pull Request (PR)** and **only when the + work is complete and reviewed**. +- **Rule of thumb:** temporary/in-progress work lives on branches; `main` always + holds final, merged, stable state. + +### 16.2 Issue Tracking (GitHub Issues) + +- Every task/feature/bug/optimization is tracked as a **GitHub Issue**. +- Feature branches and PRs reference the issue they address; PRs close issues with + `Closes #N`. +- Work items map back to the vault roadmap (`~/.claude/vault/language-model/roadmap/milestones.md`) for + traceability between design, implementation, and delivery. + +### 16.3 CI/CD (Free-Tier GitHub Actions) + +Continuous Integration runs automatically on **every PR** and **every push to +`main`**, consisting of free-tier GitHub-hosted jobs: + +| Job | Tool | Purpose | +|-----|------|---------| +| Lint | `ruff` (and/or `flake8`/`black`) | enforce style & catch basic problems | +| Type check | `mypy` or `pyright` | enforce type correctness | +| Tests | `pytest` (unit + integration) | verify behavior, catch regressions | +| Import/build smoke | `python -c "import "` | ensure modules import cleanly | + +- **PR merge requirement:** all CI jobs must pass before a PR can be merged. +- **Reproducibility:** pin Python version and core dependencies + (`pyproject.toml` + lockfile) so CI and local runs match. + +### 16.4 Definition of "Complete" (Merge Criterion) + +A PR is eligible for merge to `main` only when **all** of these hold: + +1. Feature/optimization implemented per the design in this TDD. +2. Unit + integration tests written/passing (TDD) and CI green (lint + type + test). +3. A regression/performance gate (TDR-012) passes where optimization applies. +4. Associated GitHub Issue updated/resolved. +5. Vault updated to reflect the change (decisions, milestones, results). + +This keeps `main` production-safe and makes every state of the repository auditable +and revertible. + +--- + +## 17. Updated Deliverables & Engineering Goals + +In addition to the functional deliverables in ยง15, the project delivers: + +- A locked-`main`, PR-gated **branch workflow** with **GitHub Issue tracking**. +- **Free-tier CI** (lint + type check + tests + import smoke) enforced on PRs. +- A **global knowledge vault** (`~/.claude/vault/language-model/`) kept current with decisions, + milestones, and results, acting as the single source of truth. + +--- + +*End of TDD (rev 1.1).* diff --git a/docs/TDD_PROMPT.md b/docs/TDD_PROMPT.md new file mode 100644 index 0000000..4c88544 --- /dev/null +++ b/docs/TDD_PROMPT.md @@ -0,0 +1,733 @@ +# TDD Prompt: Optimized GPT-style Language Model Engine Built From Scratch + +You are a senior ML systems engineer and software architect. + +Design a Technical Design Document (TDD) in Markdown for a project called: + +# Optimized GPT-style Language Model Engine Built From Scratch + +--- + +## Project Objective + +The goal is NOT just to train a transformer. + +The goal is to build a complete language model system with strong emphasis on: + +- low-level system design +- performance optimization +- memory efficiency +- inference optimization +- scalable architecture +- production-quality engineering practices + +The final project should demonstrate understanding of both: + +1. Machine Learning architecture +2. Systems engineering and optimization + +--- + +# Context + +I want to build a decoder-only transformer language model completely from scratch. + +Requirements: + +- No pretrained weights +- No HuggingFace models +- Implement core components manually using PyTorch +- Build the training pipeline, inference engine, optimization layer, and benchmarking system + +The model should be capable of: + +- next token prediction +- sentence completion +- paragraph continuation +- autoregressive text generation + +--- + +# System Components + +## 1. Custom Tokenizer + +Design and implement a tokenizer. + +Requirements: + +- Implement Byte Pair Encoding (BPE) +- Vocabulary creation +- Token merging algorithm +- Encoding pipeline +- Decoding pipeline +- Special tokens handling +- Vocabulary serialization + +Explain: + +- tokenizer architecture +- memory considerations +- performance optimizations +- design tradeoffs + +--- + +# 2. Efficient Dataset Engine + +Design a high-performance data loading system. + +The dataset layer should not simply load files into memory. + +Implement: + +- memory mapped dataset loading +- streaming dataset support +- multiprocessing workers +- asynchronous prefetching +- pinned memory transfer +- efficient batching +- shuffling strategy + +Explain: + +- data flow architecture +- CPU/GPU interaction +- memory management +- bottlenecks +- optimization opportunities + + +Example architecture: + + +Dataset Storage +| +| +Memory Mapping +| +| +Worker Processes +| +| +Prefetch Queue +| +| +GPU Training Batch + + +--- + +# 3. Transformer Architecture + +Implement a decoder-only transformer. + +Components: + +## Embedding Layer + +Include: + +- token embeddings +- positional embeddings + + +## Self Attention + +Implement: + +- query/key/value projections +- multi-head attention +- causal masking +- attention score calculation +- softmax normalization + + +## Feed Forward Network + +Implement: + +- linear layers +- activation functions +- dropout + + +## Transformer Block + +Include: + +- residual connections +- layer normalization +- attention layer +- feed forward layer + + +## Output Layer + +Implement: + +- vocabulary projection +- logits generation +- probability distribution + + +Explain all architectural decisions. + +--- + +# 4. Attention Optimization + +Design optimized attention mechanisms. + +Discuss: + +## Standard Attention + +Explain: + + +Attention(Q,K,V) + += Softmax(QKแต€ / sqrt(d))V + + +Discuss: + +- computational complexity +- memory complexity +- limitations for long contexts + + +## Optimized Attention + +Design improvements: + +- block based attention +- memory efficient attention +- Flash Attention inspired approach + +Explain: + +- algorithm design +- memory savings +- performance benefits +- implementation challenges + +--- + +# 5. Training Engine + +Design a complete training system. + +Include: + +## Training Loop + +- forward pass +- loss calculation +- backward pass +- optimization step + + +## Optimization + +Implement: + +- AdamW optimizer +- learning rate scheduler +- warmup +- gradient clipping + + +## Training Improvements + +Include: + +- mixed precision training +- gradient accumulation +- checkpointing +- experiment tracking + + +Explain: + +- training stability +- performance considerations +- scalability + +--- + +# 6. Inference Engine + +Build a production-style inference system. + +Requirements: + +## Text Generation + +Implement: + +- autoregressive generation +- temperature sampling +- top-k sampling +- top-p nucleus sampling + + +## Performance Optimizations + +Implement: + +### KV Cache + +Explain: + +Without cache: + + +Every token recomputes previous attention + + +With cache: + + +Previous keys and values are reused + + +Explain: + +- memory tradeoff +- latency improvement +- implementation design + + +## Batching + +Support: + +- multiple generation requests +- efficient batching +- throughput optimization + + +Measure: + +- first token latency +- tokens per second +- memory usage + +--- + +# 7. Model Optimization Layer + +Design optimization techniques. + +Include: + +## Quantization + +Implement: + +- FP32 +- FP16 +- INT8 comparison + +Explain: + +- accuracy impact +- memory reduction +- inference speed improvement + + +## Memory Optimization + +Include: + +- tensor reuse +- memory pooling +- efficient allocation + + +Discuss ideas inspired by: + +- CUDA memory allocator +- PyTorch caching allocator + + +--- + +# 8. Low-Level System Design + +Provide complete architecture design. + +Include: + +## High Level Architecture Diagram + +Show: + + +Tokenizer + +โ†“ + +Dataset Engine + +โ†“ + +Training Engine + +โ†“ + +Transformer Model + +โ†“ + +Optimization Layer + +โ†“ + +Inference Engine + +โ†“ + +API Layer + + + +## Module Responsibilities + +For every module define: + +- responsibility +- inputs +- outputs +- dependencies + + +## Class Design + +Provide: + +- important classes +- interfaces +- relationships + + +## API Design + +Define internal APIs between components. + +--- + +# 9. Benchmarking System + +Build a performance evaluation framework. + +Create benchmarks for: + +## Training + +Measure: + +- tokens per second +- training throughput +- GPU utilization +- CPU utilization +- memory usage + + +## Inference + +Measure: + +- first token latency +- generation latency +- tokens per second +- throughput +- memory footprint + + +Create: + + +benchmark/ + +โ”œโ”€โ”€ training_benchmark.py +โ”œโ”€โ”€ inference_benchmark.py +โ”œโ”€โ”€ memory_benchmark.py +โ””โ”€โ”€ latency_test.py + + +--- + +# 10. Test Driven Development Strategy + +The project must follow TDD. + +Define tests before implementation. + +For every module specify: + +## Unit Tests + +Examples: + +Tokenizer: + +- encoding correctness +- decoding correctness +- vocabulary consistency + + +Attention: + +- tensor dimensions +- masking correctness +- gradient flow + + +Dataset: + +- batch generation +- memory loading correctness + + +## Integration Tests + +Test: + +- tokenizer + dataset +- dataset + model +- model + training loop +- inference pipeline + + +## Performance Tests + +Validate: + +- throughput improvements +- memory reduction +- latency improvements + + +## Regression Tests + +Ensure: + +- optimization does not reduce accuracy +- previous functionality remains stable + +--- + +# 11. Repository Structure + +Design a production-quality repository: + + +language-model-engine/ + +โ”‚ +โ”œโ”€โ”€ tokenizer/ +โ”‚ +โ”œโ”€โ”€ dataset/ +โ”‚ +โ”œโ”€โ”€ model/ +โ”‚ +โ”œโ”€โ”€ attention/ +โ”‚ +โ”œโ”€โ”€ training/ +โ”‚ +โ”œโ”€โ”€ inference/ +โ”‚ +โ”œโ”€โ”€ optimization/ +โ”‚ +โ”œโ”€โ”€ benchmark/ +โ”‚ +โ”œโ”€โ”€ tests/ +โ”‚ +โ”œโ”€โ”€ configs/ +โ”‚ +โ”œโ”€โ”€ checkpoints/ +โ”‚ +โ”œโ”€โ”€ docs/ +โ”‚ +โ””โ”€โ”€ train.py + + +Explain the purpose of each directory. + +--- + +# 12. Development Roadmap + +Create milestones. + +## Phase 1 + +Foundation: + +- tokenizer +- dataset pipeline +- simple transformer + + +## Phase 2 + +Functional Model: + +- training pipeline +- checkpointing +- generation + + +## Phase 3 + +Optimization: + +- KV cache +- quantization +- optimized attention +- memory improvements + + +## Phase 4 + +Production Engine: + +- benchmarking +- API layer +- documentation + +--- + +# 13. Technical Decision Records + +For every major design decision provide: + +## Decision + +What approach is selected? + +## Reason + +Why this approach? + +## Alternatives + +What alternatives were considered? + +## Tradeoffs + +Advantages and disadvantages. + + +--- + +# Final Output Requirements + +Write this TDD as if it will be reviewed by: + +- senior ML engineers +- systems engineers +- infrastructure engineers + +Focus heavily on: + +- performance +- scalability +- memory efficiency +- clean architecture +- engineering tradeoffs + +Do not create a tutorial. + +Create a production-grade engineering design document. + +Output only Markdown. + + +--- + +# Project Knowledge Vault Requirement + +Maintain a dedicated project knowledge repository called: + + +vault/ + + +This folder acts as the single source of truth for all project decisions, architecture details, research notes, design decisions, experiments, and important context. + +You must maintain and update this vault continuously throughout the project. + +Whenever a new important decision, requirement, optimization idea, architectural choice, implementation detail, benchmark result, or learning is introduced, document it in the appropriate Markdown file inside `vault/`. + +Do not keep important project knowledge only in chat context. + +The vault structure should be organized and modular: + + +vault/ + +โ”œโ”€โ”€ architecture/ +โ”‚ โ”œโ”€โ”€ system_design.md +โ”‚ โ”œโ”€โ”€ model_architecture.md +โ”‚ โ”œโ”€โ”€ component_design.md +โ”‚ โ””โ”€โ”€ data_flow.md +โ”‚ +โ”œโ”€โ”€ decisions/ +โ”‚ โ”œโ”€โ”€ architecture_decisions.md +โ”‚ โ”œโ”€โ”€ optimization_decisions.md +โ”‚ โ””โ”€โ”€ tradeoffs.md +โ”‚ +โ”œโ”€โ”€ research/ +โ”‚ โ”œโ”€โ”€ transformer_notes.md +โ”‚ โ”œโ”€โ”€ attention_optimization.md +โ”‚ โ”œโ”€โ”€ inference_optimization.md +โ”‚ โ””โ”€โ”€ papers_and_references.md +โ”‚ +โ”œโ”€โ”€ implementation/ +โ”‚ โ”œโ”€โ”€ coding_guidelines.md +โ”‚ โ”œโ”€โ”€ module_interfaces.md +โ”‚ โ””โ”€โ”€ development_notes.md +โ”‚ +โ”œโ”€โ”€ experiments/ +โ”‚ โ”œโ”€โ”€ benchmarks.md +โ”‚ โ”œโ”€โ”€ training_runs.md +โ”‚ โ””โ”€โ”€ optimization_results.md +โ”‚ +โ”œโ”€โ”€ roadmap/ +โ”‚ โ”œโ”€โ”€ milestones.md +โ”‚ โ”œโ”€โ”€ completed_features.md +โ”‚ โ””โ”€โ”€ future_improvements.md +โ”‚ +โ””โ”€โ”€ project_context.md + + +## Rules + +1. Before starting any significant task: + - Review relevant files from `vault/` + - Understand existing decisions and constraints + - Avoid contradicting previous architecture decisions unless explicitly discussed + +2. After completing any significant task: + - Update the relevant Markdown files + - Record new decisions + - Record implementation changes + - Record benchmarks or observations + +3. If a new concept, optimization, or improvement is discovered: + - Create or update the appropriate research/design document + +4. Maintain consistency: + - The vault should always represent the current state of the project + - Remove outdated information or mark it as deprecated + - Keep historical decisions where useful + +5. When making recommendations: + - Check existing project decisions first + - Explain how the recommendation aligns or conflicts with current design + +The vault should evolve alongside the codebase and act as the long-term memory of this project. \ No newline at end of file diff --git a/inference/__init__.py b/inference/__init__.py new file mode 100644 index 0000000..e04e171 --- /dev/null +++ b/inference/__init__.py @@ -0,0 +1,20 @@ +"""Sampling and autoregressive generation over the KV cache. + +See docs/TDD.md ยง6 and TDR-008. +""" + +from inference.batched import generate_batch, left_pad, sample_batch +from inference.generate import generate, generate_ids +from inference.sampling import apply_temperature, apply_top_k, apply_top_p, sample + +__all__ = [ + "apply_temperature", + "apply_top_k", + "apply_top_p", + "generate", + "generate_batch", + "generate_ids", + "left_pad", + "sample", + "sample_batch", +] diff --git a/inference/batched.py b/inference/batched.py new file mode 100644 index 0000000..cfa10b2 --- /dev/null +++ b/inference/batched.py @@ -0,0 +1,154 @@ +"""Batched generation (TDD ยง6.3). + +Single-stream decoding leaves the device almost idle: a decode step is one token +wide, so the matmuls are tall-and-thin and bound by weight bandwidth rather than +arithmetic. Stacking B requests reuses each weight read across B rows, which is +close to free โ€” the win is throughput, not per-request latency. + +Two things make batching more than a reshape: + +- **Left padding, not right.** Prompts differ in length, and the next token is + always predicted from the *last* position. Right-padding would put pad tokens + there, so every request would be continuing from ``. Left-padding keeps + every sequence's real final token at index -1. +- **Per-row completion.** Requests finish at different times. A finished row + keeps being fed (the KV cache is one contiguous block; you cannot cheaply + remove a row mid-flight) but its output is masked to the stop token, so + nothing after the stop is reported. +- **Padded keys are masked out.** The pad tokens a short prompt is padded with + sit in the KV cache like any other position. Without a key-padding mask, a + short request in a wide batch attends to filler as though it were context, so + batched output silently diverges from single-stream output for exactly the + requests that were padded. `test_ragged_batch_matches_single_stream` pins it. + +Requires `impl` of `sdpa` or `manual`: the tiled `flash` kernel derives its mask +from positions per tile and rejects a per-batch padding mask rather than +silently ignoring it. +""" + +from __future__ import annotations + +import torch + +from inference.sampling import apply_temperature, apply_top_k, apply_top_p +from model import CausalLM + + +def sample_batch( + logits: torch.Tensor, + *, + temperature: float = 1.0, + top_k: int = 0, + top_p: float = 1.0, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Draw one token per row from ``[B, vocab]`` logits. Returns ``[B]`` ids. + + The single-row ``sample()`` returns an ``int`` by contract, which cannot + express a batch, so this is a separate entry point rather than a widening of + that signature. + """ + if logits.ndim != 2: + raise ValueError(f"expected [batch, vocab] logits, got {tuple(logits.shape)}") + + logits = logits.detach().float() + if temperature == 0: + return torch.argmax(logits, dim=-1) + + logits = apply_temperature(logits, temperature) + logits = apply_top_k(logits, top_k) + logits = apply_top_p(logits, top_p) + + probs = torch.softmax(logits, dim=-1) + if generator is not None and generator.device != probs.device: + probs = probs.to(generator.device) + return torch.multinomial(probs, num_samples=1, generator=generator).squeeze(-1) + + +def left_pad(sequences: list[list[int]], pad_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """Left-pad ragged prompts to ``[B, T]``. Returns ``(ids, attention_mask)``. + + Padding goes on the **left** so every row's real final token sits at index + -1, which is the position the next token is predicted from. + """ + if not sequences: + raise ValueError("no sequences to pad") + width = max(len(s) for s in sequences) + if width == 0: + raise ValueError("cannot generate from empty prompts") + + ids = torch.full((len(sequences), width), pad_id, dtype=torch.long) + mask = torch.zeros((len(sequences), width), dtype=torch.bool) + for row, seq in enumerate(sequences): + if seq: + ids[row, width - len(seq) :] = torch.tensor(seq, dtype=torch.long) + mask[row, width - len(seq) :] = True + return ids, mask + + +@torch.no_grad() +def generate_batch( + model: CausalLM, + prompts: list[list[int]], + *, + max_new_tokens: int = 128, + temperature: float = 1.0, + top_k: int = 0, + top_p: float = 1.0, + stop_id: int | None = None, + pad_id: int = 0, + generator: torch.Generator | None = None, +) -> list[list[int]]: + """Generate continuations for several prompts at once. + + Returns one list of new ids per prompt, truncated at ``stop_id`` where hit. + + Padded positions are masked out of attention, so a ragged batch produces the + same tokens as decoding each prompt on its own. + """ + if not prompts: + return [] + + was_training = model.training + model.eval() + try: + device = next(model.parameters()).device + ids, valid = left_pad(prompts, pad_id) + ids, valid = ids.to(device), valid.to(device) + batch = ids.shape[0] + + cache = model.new_cache() + logits = model(ids, kv_cache=cache, use_cache=True, key_padding_mask=valid) + + produced: list[list[int]] = [[] for _ in range(batch)] + finished = torch.zeros(batch, dtype=torch.bool, device=device) + + for _ in range(max_new_tokens): + nxt = sample_batch( + logits[:, -1], + temperature=temperature, + top_k=top_k, + top_p=top_p, + generator=generator, + ).to(device) + + for row in range(batch): + if not finished[row]: + produced[row].append(int(nxt[row])) + + if stop_id is not None: + finished |= nxt == stop_id + if bool(finished.all()): + break + + if len(cache[0]) + 1 > model.max_seq_len: + break + # Every generated token is real, so the mask grows by a True column. + valid = torch.cat([valid, torch.ones(batch, 1, dtype=torch.bool, device=device)], dim=1) + logits = model( + nxt.view(batch, 1), kv_cache=cache, use_cache=True, key_padding_mask=valid + ) + + return produced + finally: + model.train(was_training) diff --git a/inference/generate.py b/inference/generate.py new file mode 100644 index 0000000..f3be611 --- /dev/null +++ b/inference/generate.py @@ -0,0 +1,109 @@ +"""Autoregressive generation over the KV cache (TDD ยง6). + +Two entry points: + +- ``generate_ids`` โ€” tensor in, token ids out. The real engine; testable + without a tokenizer. +- ``generate`` โ€” the string-level contract from BUILD_ORDER, a thin wrapper. + +``use_cache=False`` exists purely as a correctness oracle: it re-runs the whole +prefix through the model on every step, which is the O(Tยฒ) behaviour the cache +removes. The two paths must produce identical greedy output, which is what +``test_cached_generation_matches_uncached`` checks โ€” the cheapest way to catch a +cache that has quietly stopped being equivalent to a full forward. +""" + +from __future__ import annotations + +from typing import Any, Protocol + +import torch + +from inference.sampling import sample +from model import CausalLM + + +class Tokenizer(Protocol): + def encode(self, text: str) -> list[int]: ... + def decode(self, ids: list[int]) -> str: ... + + +@torch.no_grad() +def generate_ids( + model: CausalLM, + input_ids: torch.Tensor, + *, + max_new_tokens: int = 128, + temperature: float = 1.0, + top_k: int = 0, + top_p: float = 1.0, + stop_id: int | None = None, + use_cache: bool = True, + generator: torch.Generator | None = None, +) -> list[int]: + """Generate a continuation for a single ``[1, T]`` prompt. Returns new ids only.""" + if input_ids.ndim != 2 or input_ids.shape[0] != 1: + raise ValueError(f"expected a single [1, T] prompt, got {tuple(input_ids.shape)}") + if input_ids.shape[1] == 0: + raise ValueError("cannot generate from an empty prompt") + + was_training = model.training + model.eval() + try: + max_seq_len = model.max_seq_len + device = next(model.parameters()).device + context = input_ids.to(device) + + cache = model.new_cache() if use_cache else None + logits = model(context, kv_cache=cache, use_cache=use_cache) + + produced: list[int] = [] + for _ in range(max_new_tokens): + # Only the final position predicts the next token; the rest of the + # prefill logits are thrown away. + next_id = sample( + logits[0, -1], + temperature=temperature, + top_k=top_k, + top_p=top_p, + generator=generator, + ) + produced.append(next_id) + if stop_id is not None and next_id == stop_id: + break + + step = torch.tensor([[next_id]], dtype=torch.long, device=device) + if use_cache: + if len(cache[0]) + 1 > max_seq_len: # type: ignore[index] + break + logits = model(step, kv_cache=cache, use_cache=True) + else: + context = torch.cat([context, step], dim=1) + if context.shape[1] > max_seq_len: + break + logits = model(context) + + return produced + finally: + model.train(was_training) + + +def generate( + model: CausalLM, + tokenizer: Tokenizer, + prompt: str, + *, + max_new_tokens: int = 128, + return_prompt: bool = True, + **sampling: Any, +) -> str: + """String-level generation (docs/BUILD_ORDER.md contract).""" + ids = tokenizer.encode(prompt) + if not ids: + raise ValueError("prompt encoded to zero tokens") + + device = next(model.parameters()).device + prompt_ids = torch.tensor([ids], dtype=torch.long, device=device) + new_ids = generate_ids(model, prompt_ids, max_new_tokens=max_new_tokens, **sampling) + + return tokenizer.decode((ids + new_ids) if return_prompt else new_ids) diff --git a/inference/sampling.py b/inference/sampling.py new file mode 100644 index 0000000..1636829 --- /dev/null +++ b/inference/sampling.py @@ -0,0 +1,93 @@ +"""Token sampling strategies (TDD ยง6.1). + +Filters compose in the order temperature โ†’ top-k โ†’ top-p, which is the +convention every reference implementation uses. Order matters: temperature +rescales the logits and therefore changes the probability mass that top-p is +measuring, so applying nucleus filtering first would make ``top_p`` mean +something different at every temperature. +""" + +from __future__ import annotations + +import torch + + +def apply_temperature(logits: torch.Tensor, temperature: float) -> torch.Tensor: + if temperature <= 0: + raise ValueError(f"temperature must be > 0 (use 0 for greedy), got {temperature}") + return logits / temperature + + +def apply_top_k(logits: torch.Tensor, top_k: int) -> torch.Tensor: + """Keep the ``top_k`` highest logits, mask the rest.""" + if top_k <= 0: + return logits + k = min(top_k, logits.shape[-1]) + threshold = torch.topk(logits, k).values[..., -1, None] + return logits.masked_fill(logits < threshold, float("-inf")) + + +def apply_top_p(logits: torch.Tensor, top_p: float) -> torch.Tensor: + """Nucleus filtering: keep the *smallest* set whose mass reaches ``top_p``. + + The token that carries the cumulative sum across the threshold is kept, not + dropped โ€” otherwise the retained mass is strictly below ``top_p``, and with + a confident distribution (first token already above ``p``) the whole + vocabulary would be masked and softmax would return NaN. + """ + if not 0.0 < top_p <= 1.0: + raise ValueError(f"top_p must be in (0, 1], got {top_p}") + if top_p == 1.0: + return logits + + sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1) + cumulative = sorted_logits.softmax(dim=-1).cumsum(dim=-1) + + remove = cumulative > top_p + # Shift right so the crossing token survives, and never drop the argmax. + remove[..., 1:] = remove[..., :-1].clone() + remove[..., 0] = False + + return logits.masked_fill(remove.scatter(-1, sorted_idx, remove), float("-inf")) + + +def sample( + logits: torch.Tensor, + *, + temperature: float = 1.0, + top_k: int = 0, + top_p: float = 1.0, + generator: torch.Generator | None = None, +) -> int: + """Draw one token id from a ``[vocab]`` (or ``[1, vocab]``) logit vector. + + ``temperature=0`` means greedy: the limit of the softmax as temperature + approaches zero is the argmax, but computing it would divide by zero, so it + is special-cased rather than approximated with a small epsilon. + """ + logits = logits.detach().float() + if logits.ndim == 2 and logits.shape[0] == 1: + logits = logits[0] + if logits.ndim != 1: + raise ValueError(f"expected a single [vocab] logit vector, got {tuple(logits.shape)}") + + if temperature == 0: + return int(torch.argmax(logits).item()) + + logits = apply_temperature(logits, temperature) + logits = apply_top_k(logits, top_k) + logits = apply_top_p(logits, top_p) + + probs = torch.softmax(logits, dim=-1) + + if generator is not None and generator.device != probs.device: + # torch.multinomial requires the generator and the tensor to share a + # device, but a seeded CPU generator is the natural way to ask for + # reproducible sampling โ€” so a caller doing the obvious thing would hit + # a RuntimeError as soon as the model moved to MPS or CUDA. Draw on the + # generator's device instead: this is one vocabulary-sized vector, so + # the copy is negligible beside the forward pass that produced it, and + # it makes a given seed yield identical text on every device. + probs = probs.to(generator.device) + + return int(torch.multinomial(probs, num_samples=1, generator=generator).item()) diff --git a/model/__init__.py b/model/__init__.py new file mode 100644 index 0000000..8818780 --- /dev/null +++ b/model/__init__.py @@ -0,0 +1,11 @@ +"""Decoder-only transformer: RMSNorm + SwiGLU blocks over grouped-query attention. + +See docs/TDD.md ยง3 and TDR-005, TDR-017. +""" + +from model.block import TransformerBlock +from model.causal_lm import CausalLM +from model.ffn import SwiGLU +from model.norm import RMSNorm + +__all__ = ["CausalLM", "RMSNorm", "SwiGLU", "TransformerBlock"] diff --git a/model/block.py b/model/block.py new file mode 100644 index 0000000..dc13eec --- /dev/null +++ b/model/block.py @@ -0,0 +1,69 @@ +"""Pre-norm transformer block (TDR-005, TDR-017).""" + +from __future__ import annotations + +import torch +from torch import nn + +from attention import CausalAttention, KVCache +from model.ffn import SwiGLU +from model.norm import RMSNorm + + +class TransformerBlock(nn.Module): + """``x + attn(norm(x))`` then ``x + ffn(norm(x))``. + + Pre-norm, not post-norm: the residual stream is never normalised, so there + is an unobstructed identity path from input to output through the whole + stack. That is what makes deep pre-norm models trainable without a delicate + warmup schedule. It also means a block whose sublayers output zero is + exactly the identity, which is how ``test_block_is_pre_norm_with_live_residuals`` + distinguishes this from a post-norm arrangement. + """ + + def __init__( + self, + d_model: int, + n_heads: int, + n_kv_heads: int, + d_ff: int, + *, + max_seq_len: int = 1024, + rope_theta: float = 10000.0, + dropout: float = 0.0, + impl: str = "sdpa", + bias: bool = False, + norm_eps: float = 1e-6, + ) -> None: + super().__init__() + self.attn_norm = RMSNorm(d_model, eps=norm_eps) + self.attn = CausalAttention( + d_model, + n_heads, + n_kv_heads, + max_seq_len=max_seq_len, + rope_theta=rope_theta, + dropout=dropout, + impl=impl, + bias=bias, + ) + self.ffn_norm = RMSNorm(d_model, eps=norm_eps) + self.ffn = SwiGLU(d_model, d_ff, bias=bias, dropout=dropout) + + def forward( + self, + x: torch.Tensor, + *, + kv_cache: KVCache | None = None, + use_cache: bool = False, + key_padding_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, KVCache | None]: + attn_out, new_cache = self.attn( + self.attn_norm(x), + kv_cache=kv_cache, + use_cache=use_cache, + key_padding_mask=key_padding_mask, + ) + x = x + attn_out + x = x + self.ffn(self.ffn_norm(x)) + return x, new_cache diff --git a/model/causal_lm.py b/model/causal_lm.py new file mode 100644 index 0000000..8dfbd96 --- /dev/null +++ b/model/causal_lm.py @@ -0,0 +1,193 @@ +"""``CausalLM`` โ€” the full decoder-only language model. + +Contract (docs/BUILD_ORDER.md): + + CausalLM.forward(input_ids, *, kv_cache=None, use_cache=False) -> logits [B, T, V] + CausalLM.from_config(cfg: dict) -> CausalLM + +**forward returns logits, never a loss.** Cross-entropy lives in the trainer. +Computing a loss during generation is wasted work, and keeping the model pure +means the inference path has exactly one thing to reason about. + +**The KV cache is caller-owned.** Because the contract fixes the return type to +logits, there is nowhere to hand a freshly created cache back. So the caller +allocates one with ``model.new_cache()`` and passes the same list on every +step; the per-layer ``KVCache`` objects are mutated in place. Passing +``use_cache=True`` without a cache raises rather than silently discarding it โ€” +a discarded cache turns O(T) generation back into O(Tยฒ) while still producing +correct text, which is exactly the kind of bug that never gets noticed. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from attention import KVCache +from model.block import TransformerBlock +from model.norm import RMSNorm + + +class CausalLM(nn.Module): + """Embedding โ†’ N pre-norm blocks โ†’ RMSNorm โ†’ LM head.""" + + def __init__( + self, + vocab_size: int, + d_model: int, + n_layers: int, + n_heads: int, + n_kv_heads: int | None = None, + d_ff: int | None = None, + *, + max_seq_len: int = 1024, + rope_theta: float = 10000.0, + dropout: float = 0.0, + impl: str = "sdpa", + bias: bool = False, + norm_eps: float = 1e-6, + tie_weights: bool = True, + ) -> None: + super().__init__() + n_kv_heads = n_heads if n_kv_heads is None else n_kv_heads + d_ff = 4 * d_model if d_ff is None else d_ff + + self.vocab_size = vocab_size + self.d_model = d_model + self.n_layers = n_layers + self.max_seq_len = max_seq_len + + self.embed_tokens = nn.Embedding(vocab_size, d_model) + self.dropout = nn.Dropout(dropout) + # Kept as a plain list as well as a ModuleList: iterating the ModuleList + # yields `Module`, which loses the block's attribute types. + blocks = [ + TransformerBlock( + d_model, + n_heads, + n_kv_heads, + d_ff, + max_seq_len=max_seq_len, + rope_theta=rope_theta, + dropout=dropout, + impl=impl, + bias=bias, + norm_eps=norm_eps, + ) + for _ in range(n_layers) + ] + self.blocks = blocks + self.layers = nn.ModuleList(blocks) + self.norm = RMSNorm(d_model, eps=norm_eps) + self.lm_head = nn.Linear(d_model, vocab_size, bias=False) + + self.apply(self._init_weights) + # Residual output projections are scaled down by 1/sqrt(2 * n_layers): + # every block adds two contributions to the residual stream, so without + # this the stream's variance grows linearly with depth. + residual_std = 0.02 / (2 * n_layers) ** 0.5 + for block in blocks: + nn.init.normal_(block.attn.o_proj.weight, mean=0.0, std=residual_std) + nn.init.normal_(block.ffn.down_proj.weight, mean=0.0, std=residual_std) + + # Tie *after* init, or the head's initialisation would overwrite the + # embedding's. `is`-identity matters: a copy would silently train two + # separate matrices that were only equal at step 0. + self.tie_weights = tie_weights + if tie_weights: + self.lm_head.weight = self.embed_tokens.weight + + # ------------------------------------------------------------------ + # Construction + # ------------------------------------------------------------------ + @classmethod + def from_config(cls, cfg: dict[str, Any]) -> CausalLM: + """Build from a loaded YAML config's ``model`` / ``attention`` blocks.""" + m = cfg["model"] + a = cfg.get("attention", {}) + return cls( + vocab_size=int(m["vocab_size"]), + d_model=int(m["d_model"]), + n_layers=int(m["n_layers"]), + n_heads=int(m["n_heads"]), + n_kv_heads=int(m.get("n_kv_heads", m["n_heads"])), + d_ff=int(m.get("d_ff", 4 * int(m["d_model"]))), + max_seq_len=int(m.get("max_seq_len", 1024)), + rope_theta=float(m.get("rope_theta", 10000.0)), + dropout=float(m.get("dropout", 0.0)), + impl=str(a.get("impl", "sdpa")), + tie_weights=bool(m.get("tie_weights", True)), + ) + + @staticmethod + def _init_weights(module: nn.Module) -> None: + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + + # ------------------------------------------------------------------ + # Cache + # ------------------------------------------------------------------ + def new_cache(self) -> list[KVCache]: + """One empty ``KVCache`` per layer, to be passed back on every step.""" + return [KVCache() for _ in range(self.n_layers)] + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + def forward( + self, + input_ids: torch.Tensor, + *, + kv_cache: list[KVCache] | None = None, + use_cache: bool = False, + key_padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if use_cache and kv_cache is None: + raise ValueError( + "use_cache=True needs a caller-owned cache; call model.new_cache() " + "once and pass the same list on every step" + ) + if kv_cache is not None and len(kv_cache) != self.n_layers: + raise ValueError( + f"cache has {len(kv_cache)} entries but the model has {self.n_layers} layers" + ) + + offset = len(kv_cache[0]) if kv_cache else 0 + seq_len = input_ids.shape[1] + if offset + seq_len > self.max_seq_len: + raise ValueError( + f"context length {offset + seq_len} exceeds max_seq_len={self.max_seq_len}" + ) + + x = self.dropout(self.embed_tokens(input_ids)) + for i, block in enumerate(self.blocks): + # Each layer keeps its own cache; sharing one across layers still + # runs and still generates plausible text, so it is worth being + # explicit that the indexing is per-layer. + layer_cache = kv_cache[i] if kv_cache is not None else None + x, new_cache = block( + x, + kv_cache=layer_cache, + use_cache=use_cache, + key_padding_mask=key_padding_mask, + ) + if use_cache and kv_cache is not None and new_cache is not None: + kv_cache[i] = new_cache + + logits: torch.Tensor = self.lm_head(self.norm(x)) + return logits + + # ------------------------------------------------------------------ + # Introspection + # ------------------------------------------------------------------ + def num_parameters(self, *, trainable_only: bool = False) -> int: + params = self.parameters() + if trainable_only: + params = (p for p in params if p.requires_grad) + return sum(p.numel() for p in params) diff --git a/model/ffn.py b/model/ffn.py new file mode 100644 index 0000000..bebcb14 --- /dev/null +++ b/model/ffn.py @@ -0,0 +1,37 @@ +"""SwiGLU feed-forward network (TDR-017).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + + +class SwiGLU(nn.Module): + """Gated feed-forward block: ``down(silu(gate(x)) * up(x))``. + + The ``*`` is a genuine multiplicative gate, not a sum: when ``gate(x)`` is + strongly negative, ``silu`` drives it to zero and the whole branch is + suppressed regardless of what ``up(x)`` computed. That is the property a + plain GELU MLP lacks, and it costs a third weight matrix. + + Because of that third matrix, a parameter-neutral comparison against a + ``4ยทd_model`` GELU MLP would need ``d_ff โ‰ˆ 8/3ยทd_model``. The shipped config + uses ``d_ff = 4ยทd_model``, so the FFN carries ~1.5ร— the parameters of the + GPT-2 baseline. That is a deliberate choice at this scale, recorded rather + than presented as parity. + """ + + def __init__( + self, d_model: int, d_ff: int, *, bias: bool = False, dropout: float = 0.0 + ) -> None: + super().__init__() + self.gate_proj = nn.Linear(d_model, d_ff, bias=bias) + self.up_proj = nn.Linear(d_model, d_ff, bias=bias) + self.down_proj = nn.Linear(d_ff, d_model, bias=bias) + self.dropout = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gated = F.silu(self.gate_proj(x)) * self.up_proj(x) + out: torch.Tensor = self.dropout(self.down_proj(gated)) + return out diff --git a/model/norm.py b/model/norm.py new file mode 100644 index 0000000..0b548d8 --- /dev/null +++ b/model/norm.py @@ -0,0 +1,37 @@ +"""Root-mean-square layer normalisation (TDR-017).""" + +from __future__ import annotations + +import torch +from torch import nn + + +class RMSNorm(nn.Module): + """LayerNorm with the mean-centring and the bias removed. + + ``x * rsqrt(mean(xยฒ) + eps) * weight`` + + Two things follow from dropping the mean subtraction: one fewer reduction + pass over the last dimension, and a constant offset in the input survives + normalisation instead of being annihilated. The second is what + ``test_rmsnorm_does_not_centre_the_mean`` checks โ€” swapping a LayerNorm in + here would keep every shape valid and every other test passing. + + The reduction runs in fp32 regardless of input dtype. Under autocast the + activations arrive as bf16/fp16, and squaring a large half-precision + activation overflows to ``inf`` well before the mean is taken. + """ + + def __init__(self, dim: int, eps: float = 1e-6) -> None: + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x32 = x.float() + normed = x32 * torch.rsqrt(x32.pow(2).mean(dim=-1, keepdim=True) + self.eps) + return normed.to(dtype) * self.weight + + def extra_repr(self) -> str: + return f"dim={tuple(self.weight.shape)[0]}, eps={self.eps}" diff --git a/optimization/__init__.py b/optimization/__init__.py new file mode 100644 index 0000000..9305637 --- /dev/null +++ b/optimization/__init__.py @@ -0,0 +1,31 @@ +"""Post-training optimization: weight quantization and KV cache pooling. + +See docs/TDD.md ยง7 and TDR-012. +""" + +from optimization.kv_pool import PreallocatedKVCache, preallocated_cache +from optimization.quantize import ( + VALID_BITS, + VALID_SCHEMES, + QuantizedLinear, + dequantize_int4, + dequantize_int8, + model_nbytes, + quantize, + quantize_int4, + quantize_int8, +) + +__all__ = [ + "VALID_BITS", + "VALID_SCHEMES", + "PreallocatedKVCache", + "QuantizedLinear", + "dequantize_int4", + "dequantize_int8", + "model_nbytes", + "preallocated_cache", + "quantize", + "quantize_int4", + "quantize_int8", +] diff --git a/optimization/kv_pool.py b/optimization/kv_pool.py new file mode 100644 index 0000000..6ca9615 --- /dev/null +++ b/optimization/kv_pool.py @@ -0,0 +1,125 @@ +"""Preallocated KV cache โ€” the memory pooling deferred in TDD ยง7.2. + +`attention.KVCache` grows by concatenation: every decode step allocates a new +`[B, H, T+1, D]` tensor and copies the old one into it. That is O(T) bytes +copied per step and O(Tยฒ) across a generation, plus a fresh allocation each +time for the allocator to serve and later reclaim. + +This version allocates the full `[B, H, max_seq_len, D]` arena once and writes +each step into a slice, so the steady-state cost of a decode step is one write +of the new position and zero allocations. It is a strict drop-in: it subclasses +`KVCache`, so the type checker and every `CausalAttention` call site accept it +unchanged. + +The tradeoff is honest and worth stating: peak memory is now `max_seq_len` +from the first token rather than growing with the sequence. For long-context +single-stream decoding that is a win; for many short concurrent sequences it +reserves far more than it uses. `nbytes` therefore reports the *reserved* +figure, with `nbytes_used` alongside it, so a benchmark cannot accidentally +quote the flattering one. +""" + +from __future__ import annotations + +import torch + +from attention.kv_cache import KVCache +from model import CausalLM + + +class PreallocatedKVCache(KVCache): + """Fixed-capacity KV arena written in place.""" + + def __init__( + self, + batch_size: int, + n_kv_heads: int, + max_seq_len: int, + head_dim: int, + *, + dtype: torch.dtype = torch.float32, + device: torch.device | str = "cpu", + ) -> None: + super().__init__() + shape = (batch_size, n_kv_heads, max_seq_len, head_dim) + self.keys = torch.zeros(shape, dtype=dtype, device=device) + self.values = torch.zeros(shape, dtype=dtype, device=device) + self.max_seq_len = max_seq_len + self._length = 0 + + def update(self, keys: torch.Tensor, values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + assert self.keys is not None and self.values is not None + n = keys.shape[2] + if self._length + n > self.max_seq_len: + raise ValueError( + f"KV arena full: holding {self._length} of {self.max_seq_len} " + f"positions, cannot add {n} more" + ) + if keys.shape[0] != self.keys.shape[0]: + raise ValueError( + f"batch size changed mid-generation: arena holds " + f"{self.keys.shape[0]}, got {keys.shape[0]}" + ) + + end = self._length + n + # Cast on write: under autocast the incoming K/V may be bf16 while the + # arena is fp32. Upcasting is lossless; the reverse would not be, so the + # arena dtype should match what the run actually produces. + self.keys[:, :, self._length : end] = keys.to(self.keys.dtype) + self.values[:, :, self._length : end] = values.to(self.values.dtype) + self._length = end + + # Return views of the filled prefix, never the whole arena -- attention + # must not see the zeroed tail as real keys. + return self.keys[:, :, :end], self.values[:, :, :end] + + def __len__(self) -> int: + return self._length + + def reset(self) -> None: + """Rewind to empty while keeping the arena โ€” that is the whole point.""" + self._length = 0 + + @property + def nbytes(self) -> int: + """Bytes *reserved*. See ``nbytes_used`` for the occupied portion.""" + assert self.keys is not None and self.values is not None + return ( + self.keys.numel() * self.keys.element_size() + + self.values.numel() * self.values.element_size() + ) + + @property + def nbytes_used(self) -> int: + assert self.keys is not None + per_position = self.keys[:, :, :1].numel() * self.keys.element_size() + return 2 * per_position * self._length + + +def preallocated_cache( + model: CausalLM, + *, + batch_size: int = 1, + max_seq_len: int | None = None, + dtype: torch.dtype = torch.float32, + device: torch.device | str | None = None, +) -> list[PreallocatedKVCache]: + """One arena per layer, sized from the model's own config. + + Drop-in for ``model.new_cache()``: pass the result straight to + ``forward(..., kv_cache=..., use_cache=True)``. + """ + attn = model.blocks[0].attn + if device is None: + device = next(model.parameters()).device + return [ + PreallocatedKVCache( + batch_size, + attn.n_kv_heads, + max_seq_len if max_seq_len is not None else attn.max_seq_len, + attn.head_dim, + dtype=dtype, + device=device, + ) + for _ in range(len(model.blocks)) + ] diff --git a/optimization/quantize.py b/optimization/quantize.py new file mode 100644 index 0000000..4f14bfc --- /dev/null +++ b/optimization/quantize.py @@ -0,0 +1,212 @@ +"""Post-training weight quantization (TDD ยง7.1, TDR-012). + +**Weight-only, symmetric, per-output-channel.** Weights are stored at reduced +precision and expanded back to the activation dtype inside `forward`, so the +matmul itself still runs in float. + +That means this buys **memory, not speed** โ€” and on some hardware it is +measurably *slower* than fp32 because of the dequantization step. TDD ยง7.1 +already carries that caveat; the `benchmark/` layer measures both per device +rather than assuming a winner. Real int8 throughput needs fused integer +kernels, which is a different piece of work from this one. + +Per-output-channel scales matter more than they look. A single per-tensor scale +is set by the largest weight anywhere in the matrix, so one outlier row flattens +every other row toward zero โ€” `test_int8_scales_are_per_output_channel_not_per_tensor` +pins that. +""" + +from __future__ import annotations + +import copy + +import torch +import torch.nn.functional as F +from torch import nn + +VALID_SCHEMES = ("int_weight", "fp16", "bf16", "none") +VALID_BITS = (4, 8) + +# Symmetric signed range, one level held back so the range stays centred: +# int8 uses [-127, 127] rather than [-128, 127], int4 uses [-7, 7]. +_QMAX = {8: 127, 4: 7} + + +def _scales(weight: torch.Tensor, bits: int) -> torch.Tensor: + """Per-output-channel scale, guarded against an all-zero row.""" + amax = weight.abs().amax(dim=1, keepdim=True) + # A row of exact zeros would give scale 0 and produce nan on divide. Any + # positive scale reproduces that row exactly, so 1.0 is as good as anything. + return torch.where(amax == 0, torch.ones_like(amax), amax / _QMAX[bits]) + + +def quantize_int8(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``[out, in]`` float โ†’ ``(int8 codes, per-channel scales)``.""" + scales = _scales(weight, 8) + codes = torch.round(weight / scales).clamp(-127, 127).to(torch.int8) + return codes, scales + + +def dequantize_int8(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + return codes.to(scales.dtype) * scales + + +def quantize_int4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]: + """``[out, in]`` float โ†’ ``(uint8 packed [out, ceil(in/2)], scales, in_features)``. + + Two 4-bit codes share one byte: low nibble first, high nibble second. Codes + are biased by +8 so the signed range ``[-7, 7]`` fits the unsigned nibble. + An odd ``in`` is zero-padded to the next even width, so the original width + has to be carried alongside in order to unpack. + """ + out_features, in_features = weight.shape + scales = _scales(weight, 4) + codes = torch.round(weight / scales).clamp(-7, 7).to(torch.int16) + 8 + + if in_features % 2: + codes = F.pad(codes, (0, 1), value=8) # 8 == biased zero + + low, high = codes[:, 0::2], codes[:, 1::2] + packed = (low | (high << 4)).to(torch.uint8) + return packed, scales, in_features + + +def dequantize_int4(packed: torch.Tensor, scales: torch.Tensor, in_features: int) -> torch.Tensor: + low = (packed & 0x0F).to(torch.int16) - 8 + high = (packed >> 4).to(torch.int16) - 8 + + interleaved = torch.stack((low, high), dim=2).reshape(packed.shape[0], -1) + return interleaved[:, :in_features].to(scales.dtype) * scales + + +class QuantizedLinear(nn.Module): + """A drop-in ``nn.Linear`` holding reduced-precision weights. + + Weights live in buffers, not parameters: they carry no gradient and must not + reappear in an optimizer. Quantization here is a post-training step, so a + quantised layer is not trainable, and making that structural is safer than + documenting it. + """ + + # Declared so the type checker sees Tensors; `register_buffer` alone leaves + # these typed as `Tensor | Module`. + codes: torch.Tensor + scales: torch.Tensor + bias: torch.Tensor | None + + def __init__( + self, + in_features: int, + out_features: int, + bits: int, + *, + bias: torch.Tensor | None = None, + ) -> None: + super().__init__() + if bits not in VALID_BITS: + raise ValueError(f"bits must be one of {VALID_BITS}, got {bits}") + self.in_features = in_features + self.out_features = out_features + self.bits = bits + self.register_buffer("codes", torch.empty(0, dtype=torch.uint8)) + self.register_buffer("scales", torch.empty(0)) + # Bias stays float: one value per output channel, so quantising it saves + # nothing measurable and only adds error. + self.register_buffer("bias", None if bias is None else bias.detach().clone()) + + @classmethod + def from_linear(cls, linear: nn.Linear, *, bits: int = 8) -> QuantizedLinear: + layer = cls(linear.in_features, linear.out_features, bits, bias=linear.bias) + weight = linear.weight.detach().float() + if bits == 8: + codes, scales = quantize_int8(weight) + layer.codes = codes.view(torch.uint8) + else: + packed, scales, _ = quantize_int4(weight) + layer.codes = packed + layer.scales = scales + return layer + + def dequantized_weight(self, dtype: torch.dtype) -> torch.Tensor: + if self.bits == 8: + weight = dequantize_int8(self.codes.view(torch.int8), self.scales) + else: + weight = dequantize_int4(self.codes, self.scales, self.in_features) + return weight.to(dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + bias = None if self.bias is None else self.bias.to(x.dtype) + return F.linear(x, self.dequantized_weight(x.dtype), bias) + + def extra_repr(self) -> str: + return f"in={self.in_features}, out={self.out_features}, bits={self.bits}" + + +def model_nbytes(module: nn.Module) -> int: + """Total bytes held by parameters *and* buffers. + + Buffers matter: quantised weights are buffers, so a parameters-only count + would report a quantised model as almost empty and make any saving look + spectacular and fake. + """ + seen: set[int] = set() + total = 0 + for tensor in list(module.parameters()) + list(module.buffers()): + if tensor is None or id(tensor) in seen: + continue + seen.add(id(tensor)) + total += tensor.numel() * tensor.element_size() + return total + + +def _tied_tensors(model: nn.Module) -> set[int]: + """Ids of tensors reachable under more than one name.""" + counts: dict[int, int] = {} + for _, tensor in model.named_parameters(): + counts[id(tensor)] = counts.get(id(tensor), 0) + 1 + for _, tensor in model.named_parameters(remove_duplicate=False): + counts[id(tensor)] = counts.get(id(tensor), 0) + 1 + return {i for i, n in counts.items() if n > 2} + + +def quantize( + model: nn.Module, + *, + bits: int = 8, + scheme: str = "int_weight", + skip: tuple[str, ...] = ("lm_head",), + inplace: bool = False, +) -> nn.Module: + """Quantize ``model`` (docs/BUILD_ORDER.md contract). Returns the model. + + ``skip`` defaults to the LM head because it is usually **tied** to the token + embedding. Replacing it with a quantised copy silently severs that tie: the + head becomes int8 while the embedding stays fp32, so two matrices that are + meant to be one tensor drift apart the moment anything touches them. It also + saves nothing, since the embedding keeps the float copy alive regardless. + """ + if scheme not in VALID_SCHEMES: + raise ValueError(f"unknown scheme {scheme!r}; expected one of {VALID_SCHEMES}") + if scheme == "int_weight" and bits not in VALID_BITS: + raise ValueError(f"bits must be one of {VALID_BITS}, got {bits}") + + target = model if inplace else copy.deepcopy(model) + + if scheme == "none": + return target + if scheme in ("fp16", "bf16"): + return target.to(torch.float16 if scheme == "fp16" else torch.bfloat16) + + tied = _tied_tensors(target) + for parent_name, parent in list(target.named_modules()): + for child_name, child in list(parent.named_children()): + if not isinstance(child, nn.Linear): + continue + full = f"{parent_name}.{child_name}" if parent_name else child_name + if any(full == s or full.endswith(f".{s}") for s in skip): + continue + if id(child.weight) in tied: + continue + setattr(parent, child_name, QuantizedLinear.from_linear(child, bits=bits)) + + return target diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..2ccf85d --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,73 @@ +[build-system] +requires = ["setuptools>=64", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "language-model" +version = "0.1.0" +description = "Production-grade decoder-only language model built from scratch in PyTorch" +readme = "README.md" +license = { text = "MIT" } +# Floor is 3.12, not 3.10: numpy >=2.3 ships PEP 695 `type` statements in its +# stubs, which mypy refuses to parse when analysing for an older target. Rather +# than typecheck against a Python we never run, the whole toolchain is pinned to +# one floor -- runtime, ruff target, mypy target, and the CI matrix all agree. +requires-python = ">=3.12" +dependencies = [ + "torch>=2.1", + "numpy>=1.24", + "pyyaml>=6.0", + "tqdm>=4.65", + "requests>=2.28", +] + +[project.optional-dependencies] +# Corpus downloads from a TLS-inspecting corporate network (Zscaler and similar) +# verify against a corporate root CA that certifi's bundle does not carry, so +# requests fails where curl succeeds. truststore routes verification through the +# OS keychain instead. See dataset.corpora.use_system_trust_store. +data = [ + "truststore>=0.9", +] +dev = [ + "pytest>=7.4", + "pytest-cov>=4.1", + "psutil>=5.9", + "ruff>=0.4", + "mypy>=1.8", +] + +[tool.ruff] +target-version = "py312" +line-length = 100 +src = ["."] +extend-exclude = [".venv"] + +[tool.ruff.lint] +select = ["E", "F", "I", "W", "UP", "B"] +# E741: ambiguous variable names are common in ML (l, m, i); allow them. +ignore = ["E741"] + +[tool.mypy] +python_version = "3.12" +warn_return_any = true +warn_unused_configs = true +ignore_missing_imports = true +# Start lenient; tighten per-module as the codebase grows. +disable_error_code = ["import-untyped"] + +[tool.setuptools] +packages = [ + "tokenizer", + "dataset", + "model", + "attention", + "training", + "inference", + "optimization", + "benchmark", +] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-q" diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/baselines/nano.json b/tests/baselines/nano.json new file mode 100644 index 0000000..3441a26 --- /dev/null +++ b/tests/baselines/nano.json @@ -0,0 +1,52 @@ +{ + "attention_scores": { + "scales_as": "O(T^2)", + "seq_len": 32, + "standard_score_bytes": 16384, + "standard_score_mib": 0.0156 + }, + "config": { + "d_ff": 64, + "d_model": 32, + "dropout": 0.0, + "max_seq_len": 32, + "n_heads": 4, + "n_kv_heads": 2, + "n_layers": 2, + "rope_theta": 10000.0, + "tie_weights": true, + "vocab_size": 64 + }, + "gqa_kv_cache": { + "batch_size": 1, + "gqa_bytes": 8192, + "gqa_mib": 0.0078, + "mha_bytes": 16384, + "mha_mib": 0.0156, + "n_heads": 4, + "n_kv_heads": 2, + "saving_ratio": 2.0, + "seq_len": 32 + }, + "model_memory": { + "bytes": { + "fp16": 43328, + "fp32": 86656, + "int4": 24192, + "int8": 33408 + }, + "compression_vs_fp32": { + "fp16": 2.0, + "fp32": 1.0, + "int4": 3.582, + "int8": 2.5939 + }, + "mib": { + "fp16": 0.0413, + "fp32": 0.0826, + "int4": 0.0231, + "int8": 0.0319 + }, + "parameters": 20640 + } +} diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/test_end_to_end.py b/tests/integration/test_end_to_end.py new file mode 100644 index 0000000..3fceaf8 --- /dev/null +++ b/tests/integration/test_end_to_end.py @@ -0,0 +1,159 @@ +"""Full pipeline: text โ†’ tokenizer โ†’ corpus โ†’ training โ†’ generation. + +This is the test that would catch a project that passes every unit test and +still cannot learn. Each stage is individually verified elsewhere; here they run +together against real text, and the assertion is behavioural โ€” the loss must +fall well below the uniform-entropy baseline, and the trained model must prefer +its training corpus over shuffled text. + +Sized to stay quick on CPU: a ~0.2M-parameter model over a few hundred steps. +""" + +import math + +import torch + +from dataset import DataEngine, encode_texts_to_bin +from inference import generate, generate_ids +from model import CausalLM +from tokenizer import BPE +from training.trainer import Trainer + +CORPUS = [ + "the cat sat on the mat", + "the dog sat on the log", + "the cat ran to the mat", + "the dog ran to the log", +] * 60 + + +def build(tmp_path, *, steps=250, seq_len=16): + tok = BPE(vocab_size=300, special_tokens=["<|endoftext|>"]) + tok.train(CORPUS) + + path = tmp_path / "corpus.bin" + encode_texts_to_bin(tok, CORPUS, path, eot_token="<|endoftext|>") + + cfg = { + "seed": 0, + "device": "cpu", + "model": { + "vocab_size": tok.vocab_size, + "d_model": 64, + "n_layers": 2, + "n_heads": 4, + "n_kv_heads": 2, + "d_ff": 128, + "max_seq_len": 64, + }, + "attention": {"impl": "sdpa"}, + "dataset": { + "source": "file", + "path": str(path), + "seq_len": seq_len, + "batch_size": 16, + }, + "training": { + "steps": steps, + "lr": 3.0e-3, + "min_lr": 3.0e-4, + "warmup_steps": 20, + "weight_decay": 0.1, + "grad_clip": 1.0, + "amp": False, + }, + } + + torch.manual_seed(0) + model = CausalLM.from_config(cfg) + engine = DataEngine.from_config(cfg, vocab_size=tok.vocab_size) + return tok, cfg, model, Trainer(model, engine, cfg, device=torch.device("cpu")) + + +def test_the_model_actually_learns_the_corpus(tmp_path): + """Loss must fall far below ln(vocab_size) on structured text.""" + tok, _, _, trainer = build(tmp_path) + history = trainer.train(250) + + baseline = math.log(tok.vocab_size) + first = sum(h["loss"] for h in history[:10]) / 10 + last = sum(h["loss"] for h in history[-10:]) / 10 + + assert abs(first - baseline) < 1.0, f"did not start near ln(V)={baseline:.2f}: {first:.2f}" + assert last < baseline / 2, f"loss barely moved: {first:.2f} -> {last:.2f}" + + +def test_training_lowers_loss_on_corpus_distributed_text(tmp_path): + """Training must beat random init on text drawn from the corpus distribution. + + Two traps this probe has to avoid: + + - A BPE fitted on a corpus this repetitive merges an entire sentence into a + single token, which leaves `inputs` empty and cross-entropy at nan. + - The probe must be built the same way the corpus was โ€” documents joined by + `<|endoftext|>`, not by spaces. Joined with spaces the sequence is + out-of-distribution, and a well-trained model correctly scores it *worse* + than an untrained one, which looks like a training failure but is the + opposite. + """ + tok, cfg, model, trainer = build(tmp_path) + + probe_text = "".join(text + "<|endoftext|>" for text in CORPUS[:4]) + probe = torch.tensor([tok.encode(probe_text)], dtype=torch.long) + assert probe.shape[1] > 4, f"probe collapsed to {probe.shape[1]} token(s)" + inputs, labels = probe[:, :-1], probe[:, 1:] + + def loss_of(m): + with torch.no_grad(): + logits = m(inputs) + return torch.nn.functional.cross_entropy( + logits.reshape(-1, logits.shape[-1]), labels.reshape(-1) + ).item() + + torch.manual_seed(1) + untrained = loss_of(CausalLM.from_config(cfg)) + trainer.train(250) + trained = loss_of(model) + + assert trained < untrained / 2, f"untrained {untrained:.2f} -> trained {trained:.2f}" + + +def test_generation_is_deterministic_and_decodable(tmp_path): + tok, _, model, trainer = build(tmp_path, steps=150) + trainer.train(150) + + a = generate(model, tok, "the cat", max_new_tokens=20, temperature=0.0) + b = generate(model, tok, "the cat", max_new_tokens=20, temperature=0.0) + + assert a == b, "greedy generation is not deterministic" + assert a.startswith("the cat") + assert isinstance(a, str) and len(a) > len("the cat") + + +def test_the_kv_cache_does_not_change_what_a_trained_model_writes(tmp_path): + """Cache equivalence has to hold for trained weights, not just random ones.""" + tok, _, model, trainer = build(tmp_path, steps=150) + trainer.train(150) + + prompt = torch.tensor([tok.encode("the dog")], dtype=torch.long) + cached = generate_ids(model, prompt, max_new_tokens=15, temperature=0.0, use_cache=True) + plain = generate_ids(model, prompt, max_new_tokens=15, temperature=0.0, use_cache=False) + assert cached == plain + + +def test_a_resumed_run_reaches_the_same_place(tmp_path): + """Train 100, checkpoint, reload, train 100 more == train 200 straight.""" + _, cfg, model_a, trainer_a = build(tmp_path, steps=200) + trainer_a.train(100) + ckpt = tmp_path / "mid.pt" + trainer_a.save(str(ckpt)) + straight = trainer_a.train(100)[-1]["loss"] + + _, _, _, trainer_b = build(tmp_path, steps=200) + trainer_b.load(str(ckpt)) + resumed = trainer_b.train(100)[-1]["loss"] + + # The data stream restarts on resume (checkpoint.py documents this), so the + # batches differ; the losses converge to the same place rather than matching + # bit for bit. + assert abs(straight - resumed) < 0.5, f"straight {straight:.3f} vs resumed {resumed:.3f}" diff --git a/tests/integration/test_tokenizer_dataset_attention.py b/tests/integration/test_tokenizer_dataset_attention.py new file mode 100644 index 0000000..fb29fc6 --- /dev/null +++ b/tests/integration/test_tokenizer_dataset_attention.py @@ -0,0 +1,121 @@ +"""Integration: raw text โ†’ tokenizer โ†’ .bin corpus โ†’ batches โ†’ attention. + +Unit tests pin each module against its own contract; this pins the *seams* +between them, which is where the contracts actually get violated: + +- token ids the tokenizer emits must fit the dataset's ``uint16`` format +- batch tensors must be the dtype/shape an ``nn.Embedding`` will accept +- the next-token shift must survive the round trip through disk +""" + +import torch +from torch import nn + +from attention import CausalAttention +from dataset import DataEngine, count_tokens_bin, encode_texts_to_bin +from tokenizer import BPE + +CORPUS = [ + "the quick brown fox jumps over the lazy dog", + "the lazy dog sleeps while the quick fox runs", + "a fox and a dog met under the brown oak tree", +] * 20 + + +def build_tokenizer(): + tok = BPE(vocab_size=400, special_tokens=["<|endoftext|>"]) + tok.train(CORPUS) + return tok + + +def test_tokenizer_output_fits_the_on_disk_token_format(tmp_path): + tok = build_tokenizer() + path = tmp_path / "corpus.bin" + n = encode_texts_to_bin(tok, CORPUS, path, eot_token="<|endoftext|>") + assert n == count_tokens_bin(path) > 0 + # uint16 caps the vocab at 65535; a tokenizer that outgrew it must fail loudly + # at write time, not corrupt ids by wrapping. + assert tok.vocab_size <= 65536 + + +def test_corpus_round_trips_through_disk(tmp_path): + """Decoding the .bin must reproduce the original text, EOT markers and all.""" + tok = build_tokenizer() + path = tmp_path / "corpus.bin" + encode_texts_to_bin(tok, CORPUS[:2], path, eot_token="<|endoftext|>") + + from dataset import read_tokens_bin + + decoded = tok.decode(read_tokens_bin(path).tolist()) + assert decoded == "".join(text + "<|endoftext|>" for text in CORPUS[:2]) + + +def test_batches_feed_an_embedding_and_then_attention(tmp_path): + """The real training path, minus the loss: ids โ†’ embedding โ†’ attention.""" + tok = build_tokenizer() + path = tmp_path / "corpus.bin" + encode_texts_to_bin(tok, CORPUS, path, eot_token="<|endoftext|>") + + d_model, seq_len, batch_size = 32, 16, 4 + engine = DataEngine.from_config( + { + "dataset": { + "source": "file", + "path": str(path), + "seq_len": seq_len, + "batch_size": batch_size, + }, + "seed": 0, + }, + vocab_size=tok.vocab_size, + ) + batch = engine.next_batch() + + # Contract with nn.Embedding: long dtype, every id inside the vocabulary. + assert batch["input_ids"].dtype == torch.long + assert int(batch["input_ids"].max()) < tok.vocab_size + + embed = nn.Embedding(tok.vocab_size, d_model) + attn = CausalAttention( + d_model=d_model, n_heads=4, n_kv_heads=2, max_seq_len=seq_len, impl="sdpa" + ).eval() + + out, _ = attn(embed(batch["input_ids"])) + assert out.shape == (batch_size, seq_len, d_model) + assert torch.isfinite(out).all() + + +def test_shift_invariant_survives_the_whole_pipeline(tmp_path): + """labels[t] == input_ids[t+1] must still hold after tokenize โ†’ disk โ†’ batch.""" + tok = build_tokenizer() + path = tmp_path / "corpus.bin" + encode_texts_to_bin(tok, CORPUS, path, eot_token="<|endoftext|>") + + engine = DataEngine.from_config( + {"dataset": {"source": "file", "path": str(path), "seq_len": 16, "batch_size": 4}}, + vocab_size=tok.vocab_size, + ) + batch = engine.next_batch() + assert torch.equal(batch["labels"][:, :-1], batch["input_ids"][:, 1:]) + + +def test_generation_step_shapes_hold_end_to_end(tmp_path): + """Prefill a real prompt, then decode one token with the cache.""" + tok = build_tokenizer() + d_model = 32 + embed = nn.Embedding(tok.vocab_size, d_model) + attn = CausalAttention( + d_model=d_model, n_heads=4, n_kv_heads=2, max_seq_len=128, impl="sdpa" + ).eval() + + ids = torch.tensor([tok.encode("the quick brown fox")], dtype=torch.long) + prompt_len = ids.shape[1] + + out, cache = attn(embed(ids), use_cache=True) + assert out.shape == (1, prompt_len, d_model) + assert len(cache) == prompt_len + + next_id = torch.tensor([[tok.encode(" dog")[0]]], dtype=torch.long) + step, cache = attn(embed(next_id), kv_cache=cache, use_cache=True) + assert step.shape == (1, 1, d_model) + assert len(cache) == prompt_len + 1 diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/test_attention.py b/tests/unit/test_attention.py new file mode 100644 index 0000000..7e8d3df --- /dev/null +++ b/tests/unit/test_attention.py @@ -0,0 +1,436 @@ +"""Attention unit tests (TDD - written before implementation). + +Exit criteria for build stage 2 (docs/BUILD_ORDER.md): causal-mask correctness +and KV-cache equivalence. + +Three bugs this file exists to catch, all of which train and generate happily +while being wrong: + +1. **Leaky causal mask** -- position t attends to t+1. Loss plummets (the model + is reading the answer) and generation is incoherent. Caught by + ``test_future_tokens_cannot_change_the_past``. +2. **RoPE offset ignored under caching** -- during cached decode the new token + is rotated as if it were at position 0 instead of position ``cache_len``. + Prefill looks perfect; only multi-token generation degrades. Caught by + ``test_cached_decode_matches_full_forward``. +3. **GQA head misalignment** -- query head h paired with the wrong KV group. + Output shapes stay valid and the model still trains, just worse. Caught by + ``test_gqa_with_one_kv_group_equals_mha``. +""" + +import pytest +import torch + +from attention import CausalAttention, KVCache, RotaryEmbedding, apply_rotary_emb, repeat_kv + +D_MODEL, N_HEADS, N_KV_HEADS, HEAD_DIM = 32, 4, 2, 8 +BATCH, SEQ = 2, 6 +ALL_IMPLS = ("manual", "sdpa", "flash") + + +def make_attn(**overrides): + kwargs = dict( + d_model=D_MODEL, + n_heads=N_HEADS, + n_kv_heads=N_KV_HEADS, + max_seq_len=64, + dropout=0.0, + impl="manual", + ) + kwargs.update(overrides) + attn = CausalAttention(**kwargs) + return attn.eval() + + +@pytest.fixture +def x(): + torch.manual_seed(0) + return torch.randn(BATCH, SEQ, D_MODEL) + + +# --------------------------------------------------------------------------- +# Shapes and configuration validation +# --------------------------------------------------------------------------- + + +def test_output_shape_matches_input(x): + out, cache = make_attn()(x) + assert out.shape == (BATCH, SEQ, D_MODEL) + assert cache is None # use_cache=False must not allocate one + + +def test_head_dim_is_derived_from_d_model(): + assert make_attn().head_dim == D_MODEL // N_HEADS == HEAD_DIM + + +def test_kv_heads_must_divide_query_heads(): + with pytest.raises(ValueError, match="n_kv_heads"): + make_attn(n_heads=4, n_kv_heads=3) + + +def test_d_model_must_divide_into_heads(): + with pytest.raises(ValueError, match="d_model"): + make_attn(d_model=30, n_heads=4) + + +def test_unknown_impl_is_rejected(): + with pytest.raises(ValueError, match="impl"): + make_attn(impl="magic") + + +def test_kv_projections_are_narrower_than_q_under_gqa(): + """The GQA parameter saving is in k_proj/v_proj, so assert it directly.""" + attn = make_attn() + assert attn.q_proj.out_features == N_HEADS * HEAD_DIM + assert attn.k_proj.out_features == N_KV_HEADS * HEAD_DIM + assert attn.v_proj.out_features == N_KV_HEADS * HEAD_DIM + + +def test_rope_tables_are_not_in_the_state_dict(): + """cos/sin are derived constants; persisting them bloats every checkpoint.""" + keys = make_attn().state_dict().keys() + assert not [k for k in keys if "cos" in k or "sin" in k] + + +# --------------------------------------------------------------------------- +# Causality +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("impl", ALL_IMPLS) +def test_future_tokens_cannot_change_the_past(impl): + """Perturb the tail of the sequence; the head of the output must not move.""" + torch.manual_seed(0) + attn = make_attn(impl=impl) + a = torch.randn(1, SEQ, D_MODEL) + b = a.clone() + cut = 3 + b[:, cut:] = torch.randn(1, SEQ - cut, D_MODEL) # rewrite everything from `cut` on + + out_a, _ = attn(a) + out_b, _ = attn(b) + torch.testing.assert_close(out_a[:, :cut], out_b[:, :cut], rtol=1e-5, atol=1e-6) + # sanity: the perturbation must actually have done something downstream + assert not torch.allclose(out_a[:, cut:], out_b[:, cut:]) + + +@pytest.mark.parametrize("impl", ALL_IMPLS) +def test_first_position_attends_only_to_itself(impl): + """With one visible key, attention is the identity on that value.""" + torch.manual_seed(0) + attn = make_attn(impl=impl) + single, _ = attn(torch.randn(1, 1, D_MODEL)) + assert single.shape == (1, 1, D_MODEL) + + +# --------------------------------------------------------------------------- +# Rotary position embeddings +# --------------------------------------------------------------------------- + + +def test_rope_preserves_vector_norm(): + """A rotation changes direction, never magnitude.""" + rope = RotaryEmbedding(HEAD_DIM, max_seq_len=64) + q = torch.randn(1, 2, SEQ, HEAD_DIM) + cos, sin = rope(SEQ) + rotated = apply_rotary_emb(q, cos, sin) + torch.testing.assert_close(q.norm(dim=-1), rotated.norm(dim=-1), rtol=1e-5, atol=1e-6) + + +def test_rope_score_depends_only_on_relative_distance(): + """The defining RoPE property: is a function of (m - n). + + Same vector pair at positions (1,3) and (4,6) must score identically. + """ + rope = RotaryEmbedding(HEAD_DIM, max_seq_len=64) + torch.manual_seed(0) + q = torch.randn(1, 1, 1, HEAD_DIM) + k = torch.randn(1, 1, 1, HEAD_DIM) + + def score(m, n): + cq, sq = rope(1, offset=m) + ck, sk = rope(1, offset=n) + return (apply_rotary_emb(q, cq, sq) * apply_rotary_emb(k, ck, sk)).sum() + + torch.testing.assert_close(score(1, 3), score(4, 6), rtol=1e-5, atol=1e-6) + assert not torch.allclose(score(1, 3), score(1, 6)) + + +def test_rope_offset_selects_later_positions(): + rope = RotaryEmbedding(HEAD_DIM, max_seq_len=64) + cos_all, _ = rope(10) + cos_off, _ = rope(4, offset=6) + torch.testing.assert_close(cos_all[6:], cos_off) + + +def test_rope_rejects_positions_past_the_table(): + rope = RotaryEmbedding(HEAD_DIM, max_seq_len=8) + with pytest.raises(ValueError, match="max_seq_len"): + rope(4, offset=6) + + +# --------------------------------------------------------------------------- +# Grouped-query attention +# --------------------------------------------------------------------------- + + +def test_repeat_kv_maps_query_head_to_its_group(): + # 2 kv heads, 2 repeats -> query heads [0,1] use kv head 0; [2,3] use kv head 1 + kv = torch.arange(2, dtype=torch.float32).view(1, 2, 1, 1).expand(1, 2, 3, HEAD_DIM) + out = repeat_kv(kv.contiguous(), 2) + assert out.shape == (1, 4, 3, HEAD_DIM) + assert out[0, 0, 0, 0] == 0 and out[0, 1, 0, 0] == 0 + assert out[0, 2, 0, 0] == 1 and out[0, 3, 0, 0] == 1 + + +def test_repeat_kv_is_a_noop_for_mha(): + kv = torch.randn(1, 4, 3, HEAD_DIM) + assert torch.equal(repeat_kv(kv, 1), kv) + + +def test_gqa_with_one_kv_group_equals_mha(x): + """n_kv_heads == n_heads must reduce exactly to multi-head attention. + + If the group mapping is wrong this still runs; it just stops being MHA. + """ + torch.manual_seed(0) + mha = make_attn(n_kv_heads=N_HEADS) + assert mha.n_rep == 1 + out, _ = mha(x) + assert out.shape == (BATCH, SEQ, D_MODEL) + + +def test_gqa_shrinks_the_kv_cache_by_the_group_ratio(x): + """The whole point of TDR-016 -- assert the memory saving, don't assume it.""" + _, gqa_cache = make_attn(n_kv_heads=2)(x, use_cache=True) + _, mha_cache = make_attn(n_kv_heads=N_HEADS)(x, use_cache=True) + assert gqa_cache.keys.shape[1] == 2 + assert mha_cache.keys.shape[1] == N_HEADS + assert gqa_cache.nbytes * 2 == mha_cache.nbytes + + +# --------------------------------------------------------------------------- +# KV cache -- the stage 2 exit test +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("impl", ALL_IMPLS) +def test_cached_decode_matches_full_forward(impl): + """Feeding tokens one at a time with a cache == one full parallel forward. + + This is the single most important test in the module: it simultaneously + pins the causal mask, the RoPE position offset, and the cache append order. + """ + torch.manual_seed(0) + attn = make_attn(impl=impl) + seq = torch.randn(1, SEQ, D_MODEL) + + full, _ = attn(seq) + + cache = None + steps = [] + for t in range(SEQ): + step, cache = attn(seq[:, t : t + 1], kv_cache=cache, use_cache=True) + steps.append(step) + incremental = torch.cat(steps, dim=1) + + torch.testing.assert_close(full, incremental, rtol=1e-4, atol=1e-5) + + +def test_prefill_then_decode_matches_full_forward(): + """The realistic inference path: prompt in one shot, then token by token.""" + torch.manual_seed(0) + attn = make_attn(impl="sdpa") + seq = torch.randn(1, SEQ, D_MODEL) + full, _ = attn(seq) + + prefill = 4 + head, cache = attn(seq[:, :prefill], use_cache=True) + outs = [head] + for t in range(prefill, SEQ): + step, cache = attn(seq[:, t : t + 1], kv_cache=cache, use_cache=True) + outs.append(step) + + torch.testing.assert_close(full, torch.cat(outs, dim=1), rtol=1e-4, atol=1e-5) + + +def test_cache_grows_by_one_per_decode_step(x): + attn = make_attn() + _, cache = attn(x, use_cache=True) + assert len(cache) == SEQ + for expected in range(SEQ + 1, SEQ + 4): + _, cache = attn(x[:, :1], kv_cache=cache, use_cache=True) + assert len(cache) == expected + + +def test_empty_cache_reports_zero_length(): + assert len(KVCache()) == 0 + + +def test_cache_rejects_a_mismatched_batch(x): + attn = make_attn() + _, cache = attn(x, use_cache=True) + with pytest.raises(ValueError, match="batch"): + attn(torch.randn(BATCH + 1, 1, D_MODEL), kv_cache=cache, use_cache=True) + + +def test_attention_rejects_sequences_past_max_seq_len(): + attn = make_attn(max_seq_len=8) + with pytest.raises(ValueError, match="max_seq_len"): + attn(torch.randn(1, 9, D_MODEL)) + + +# --------------------------------------------------------------------------- +# Implementation equivalence -- the regression gate of TDD ยง4.3 / ยง10.4 +# --------------------------------------------------------------------------- + + +def test_all_impls_agree(x): + """manual is the oracle; sdpa and flash must match it or they are wrong.""" + torch.manual_seed(0) + reference = make_attn(impl="manual") + outputs = {} + for impl in ALL_IMPLS: + attn = make_attn(impl=impl) + attn.load_state_dict(reference.state_dict()) + outputs[impl], _ = attn(x) + + torch.testing.assert_close(outputs["manual"], outputs["sdpa"], rtol=1e-5, atol=1e-6) + torch.testing.assert_close(outputs["manual"], outputs["flash"], rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("block", [1, 2, 4, 16, 1024]) +def test_flash_result_is_independent_of_block_size(x, block): + """Tiling is an implementation detail; block size must not change the answer. + + Includes blocks larger than the sequence (degenerate single tile) and + block=1 (worst case for the online-softmax rescale path). + """ + torch.manual_seed(0) + reference = make_attn(impl="manual") + tiled = make_attn(impl="flash", block_q=block, block_k=block) + tiled.load_state_dict(reference.state_dict()) + + expected, _ = reference(x) + actual, _ = tiled(x) + torch.testing.assert_close(expected, actual, rtol=1e-5, atol=1e-6) + + +def test_flash_never_materialises_the_full_score_matrix(): + """The O(T) memory claim of TDR-007 rests on this; assert the tile bound. + + A T=256 sequence with 32-wide tiles must never allocate a 256x256 score + tensor. We spy on matmul shapes rather than trusting the docstring. + """ + attn = make_attn(impl="flash", block_q=32, block_k=32, max_seq_len=512) + seen = [] + real_matmul = torch.matmul + + def spy(a, b, **kw): + out = real_matmul(a, b, **kw) + seen.append(out.shape) + return out + + torch.matmul = spy + try: + attn(torch.randn(1, 256, D_MODEL)) + finally: + torch.matmul = real_matmul + + score_dims = [s[-1] for s in seen] + assert max(score_dims) <= 32, f"a tile exceeded block_k: {max(score_dims)}" + + +# --------------------------------------------------------------------------- +# Gradients +# --------------------------------------------------------------------------- + + +def _saved_bytes(attn, seq_len): + """Bytes autograd retains for backward โ€” the quantity TDR-021 is about.""" + total = 0 + + def pack(t): + nonlocal total + total += t.numel() * t.element_size() + return t + + with torch.enable_grad(), torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t): + attn(torch.randn(1, seq_len, D_MODEL)) + return total + + +def test_flash_gradients_match_the_manual_oracle(): + """The custom backward must be exact, not merely plausible (TDR-021). + + Computed in float64 so the comparison is about the derivation, not about + float32 noise. A sign error or a missing scale factor still trains -- just + toward a slightly different objective than the one written down. + """ + from attention.kernels import causal_block_mask, flash_attention, manual_attention + + torch.manual_seed(0) + shape = (2, 3, 37, 16) # deliberately not a multiple of the block size + + def leaves(): + return [torch.randn(*shape, dtype=torch.float64, requires_grad=True) for _ in range(3)] + + q, k, v = leaves() + qf, kf, vf = (t.detach().clone().requires_grad_(True) for t in (q, k, v)) + + reference = manual_attention(q, k, v, mask=causal_block_mask(37, 37, 0, q.device)) + tiled = flash_attention(qf, kf, vf, causal=True, block_q=8, block_k=8) + torch.testing.assert_close(tiled, reference, rtol=1e-10, atol=1e-12) + + seed = torch.randn_like(reference) + reference.backward(seed) + tiled.backward(seed) + for name, want, got in ( + ("dq", q.grad, qf.grad), + ("dk", k.grad, kf.grad), + ("dv", v.grad, vf.grad), + ): + torch.testing.assert_close(got, want, rtol=1e-9, atol=1e-11, msg=f"{name} mismatch") + + +def test_flash_saved_activations_grow_linearly_not_quadratically(): + """The O(T) memory claim, asserted for *training* rather than inference. + + Before the custom backward, the tiled kernel was plain autograd ops, so every + tile's intermediates were retained and summed back to O(T^2) -- it measured + 1.30x *worse* than the naive path. Recomputing scores in backward means only + (q, k, v, out, lse) are kept. + """ + attn = make_attn(impl="flash", max_seq_len=2048, block_q=32, block_k=32) + torch.manual_seed(0) + small = _saved_bytes(attn, 256) + large = _saved_bytes(attn, 512) + # Linear would be 2.0x; quadratic 4.0x. Allow slack for fixed overheads. + assert large / small < 2.6, f"doubling T cost {large / small:.2f}x saved memory" + + +def test_flash_retains_far_less_than_manual_for_training(): + torch.manual_seed(0) + manual = _saved_bytes(make_attn(impl="manual", max_seq_len=2048), 512) + flash = _saved_bytes(make_attn(impl="flash", max_seq_len=2048, block_q=32, block_k=32), 512) + assert flash < manual / 4, f"flash {flash} vs manual {manual} bytes" + + +def test_manual_saved_activations_are_quadratic(): + """Pins the baseline the comparison above is against.""" + attn = make_attn(impl="manual", max_seq_len=2048) + torch.manual_seed(0) + ratio = _saved_bytes(attn, 512) / _saved_bytes(attn, 256) + assert ratio > 3.0, f"expected ~4x for the naive path, got {ratio:.2f}x" + + +@pytest.mark.parametrize("impl", ALL_IMPLS) +def test_gradients_reach_every_projection(impl, x): + attn = make_attn(impl=impl) + out, _ = attn(x) + out.sum().backward() + for name in ("q_proj", "k_proj", "v_proj", "o_proj"): + grad = getattr(attn, name).weight.grad + assert grad is not None, f"{name} received no gradient" + assert torch.isfinite(grad).all(), f"{name} gradient has nan/inf" + assert grad.abs().sum() > 0, f"{name} gradient is all zero" diff --git a/tests/unit/test_benchmark.py b/tests/unit/test_benchmark.py new file mode 100644 index 0000000..17cd213 --- /dev/null +++ b/tests/unit/test_benchmark.py @@ -0,0 +1,255 @@ +"""Benchmark harness tests (TDD - written before implementation). + +Exit criterion for build stage 7 (docs/BUILD_ORDER.md): reproduces a recorded +baseline. + +Only the *deterministic* half of the report can honour that literally. Memory +figures are exact arithmetic over tensor shapes, so they are asserted against a +checked-in baseline byte for byte. Latency is machine-, thermal- and +load-dependent; asserting a wall-clock number would produce a test that fails on +a busy laptop and passes on an idle one, which is worse than no test. So timings +are asserted only for *shape* -- present, finite, ordered percentiles. +""" + +import json +from pathlib import Path + +import pytest +import torch + +from benchmark import ( + activation_memory_estimate, + benchmark_decode, + benchmark_prefill, + build_report, + deterministic_report, + environment, + gqa_saving_report, + kv_cache_memory, + kv_cache_speedup, + model_memory_report, + synchronize, + time_it, +) +from benchmark.harness import Timing +from model import CausalLM + +BASELINE = Path(__file__).parent.parent / "baselines" / "nano.json" + +CFG = { + "model": { + "vocab_size": 64, + "d_model": 32, + "n_layers": 2, + "n_heads": 4, + "n_kv_heads": 2, + "d_ff": 64, + "max_seq_len": 32, + "rope_theta": 10000.0, + "dropout": 0.0, + "tie_weights": True, + }, + "attention": {"impl": "manual"}, + "device": "cpu", +} + + +def make_model(): + torch.manual_seed(0) + return CausalLM.from_config(CFG).eval() + + +# --------------------------------------------------------------------------- +# Timing primitives +# --------------------------------------------------------------------------- + + +def test_synchronize_is_safe_on_every_device(): + for device in ("cpu", None): + synchronize(device) + + +def test_time_it_runs_exactly_the_requested_repeats(): + calls = [] + time_it(lambda: calls.append(1), warmup=3, repeats=7, device="cpu") + assert len(calls) == 10 # 3 warmup + 7 timed + + +def test_warmup_calls_are_not_timed(): + """Warmup exists to keep first-call costs out of the average.""" + timing = time_it(lambda: None, warmup=2, repeats=5, device="cpu") + assert timing.runs == 5 + + +def test_percentiles_are_ordered(): + timing = time_it(lambda: sum(range(500)), warmup=1, repeats=20, device="cpu") + assert ( + timing.min_ms + <= timing.p50_ms + <= timing.p90_ms + <= timing.p95_ms + <= timing.p99_ms + <= timing.max_ms + ) + assert timing.mean_ms > 0 + assert timing.stdev_ms >= 0 + + +def test_timing_reports_median_alongside_mean(): + """Step times are right-skewed; a mean alone misrepresents steady state.""" + timing = Timing.from_samples([0.001] * 9 + [1.0]) + assert timing.p50_ms == pytest.approx(1.0) + assert timing.mean_ms > timing.p50_ms * 50 + + +def test_median_is_an_alias_for_p50(): + timing = Timing.from_samples([0.001, 0.002, 0.003]) + assert timing.median_ms == timing.p50_ms + + +def test_percentiles_are_interpolated_not_index_rounded(): + """Regression: `ms[int(q*n)]` collapsed every high percentile onto max. + + With 10 samples, int(0.90*10) == int(0.95*10) == int(0.99*10) == 9, so p90, + p95 and p99 all returned the maximum. They must be distinct here. + """ + timing = Timing.from_samples([i / 1000 for i in range(1, 11)]) # 1..10 ms + assert timing.p90_ms < timing.p95_ms < timing.p99_ms + assert timing.p90_ms < timing.max_ms + + +def test_percentiles_match_the_closed_form_on_a_known_sample(): + """1..101 ms: the inclusive estimator puts pQ exactly at Q+1 ms.""" + timing = Timing.from_samples([i / 1000 for i in range(1, 102)]) + assert timing.p50_ms == pytest.approx(51.0) + assert timing.p90_ms == pytest.approx(91.0) + assert timing.p95_ms == pytest.approx(96.0) + assert timing.p99_ms == pytest.approx(100.0) + + +def test_resolvable_percentile_exposes_undersampling(): + """A p99 from 10 runs is not a p99; the report has to say so.""" + assert Timing.from_samples([0.001] * 10).resolvable_percentile == pytest.approx(90.0) + assert Timing.from_samples([0.001] * 100).resolvable_percentile == pytest.approx(99.0) + + +def test_a_single_sample_does_not_crash_the_estimator(): + timing = Timing.from_samples([0.005]) + assert timing.runs == 1 + assert timing.p50_ms == timing.p99_ms == pytest.approx(5.0) + assert timing.stdev_ms == 0.0 + + +def test_no_samples_is_rejected(): + with pytest.raises(ValueError, match="at least one sample"): + Timing.from_samples([]) + + +def test_as_dict_carries_every_percentile(): + d = time_it(lambda: None, warmup=0, repeats=5, device="cpu").as_dict() + for key in ("p50_ms", "p90_ms", "p95_ms", "p99_ms", "stdev_ms", "resolvable_percentile"): + assert key in d, f"as_dict is missing {key}" + + +def test_zero_repeats_is_rejected(): + with pytest.raises(ValueError, match="repeats"): + time_it(lambda: None, repeats=0) + + +# --------------------------------------------------------------------------- +# Memory arithmetic +# --------------------------------------------------------------------------- + + +def test_kv_cache_memory_matches_the_closed_form(): + # 2 (K and V) x layers x batch x heads x T x dim x bytes + assert ( + kv_cache_memory(n_layers=4, n_kv_heads=2, head_dim=64, seq_len=128, bytes_per_element=4) + == 2 * 4 * 1 * 2 * 128 * 64 * 4 + ) + + +def test_gqa_saving_is_exactly_the_head_ratio(): + report = gqa_saving_report(n_layers=8, n_heads=8, n_kv_heads=2, head_dim=64, seq_len=512) + assert report["saving_ratio"] == pytest.approx(4.0) + assert report["gqa_bytes"] * 4 == report["mha_bytes"] + + +def test_attention_score_memory_is_quadratic_in_sequence_length(): + small = activation_memory_estimate(batch_size=1, n_heads=8, seq_len=128) + large = activation_memory_estimate(batch_size=1, n_heads=8, seq_len=256) + assert large["standard_score_bytes"] == small["standard_score_bytes"] * 4 + + +def test_model_memory_report_orders_the_precisions(): + report = model_memory_report(make_model()) + b = report["bytes"] + assert b["fp32"] > b["fp16"] > b["int8"] > b["int4"] + assert report["compression_vs_fp32"]["fp16"] == pytest.approx(2.0, rel=0.01) + + +def test_memory_numbers_are_identical_across_runs(): + """Exact arithmetic -- two runs must agree bit for bit, or it is measured.""" + assert model_memory_report(make_model()) == model_memory_report(make_model()) + + +# --------------------------------------------------------------------------- +# The exit test: reproduce a recorded baseline +# --------------------------------------------------------------------------- + + +def test_deterministic_report_reproduces_the_baseline(): + """Stage 7 exit criterion, over the half of the report that can be exact. + + Regenerate deliberately with: + python -m benchmark.report --config --deterministic-only + A diff here means the architecture's memory profile changed -- which is + either the point of your change, or a bug. + """ + assert BASELINE.exists(), f"missing baseline: {BASELINE}" + recorded = json.loads(BASELINE.read_text()) + assert deterministic_report(CFG) == recorded + + +def test_the_baseline_holds_no_wall_clock_numbers(): + """A timing in a checked-in baseline would fail on any other machine.""" + text = BASELINE.read_text().lower() + for banned in ("_ms", "tokens_per_sec", "timing", "elapsed"): + assert banned not in text, f"baseline contains a measured field: {banned}" + + +# --------------------------------------------------------------------------- +# Measured metrics -- shape only +# --------------------------------------------------------------------------- + + +def test_prefill_reports_positive_finite_throughput(): + report = benchmark_prefill(make_model(), prompt_len=16, repeats=3, warmup=1) + assert report["tokens_per_sec"] > 0 + assert report["timing"]["p50_ms"] > 0 + + +def test_decode_reports_per_token_cost(): + report = benchmark_decode(make_model(), prompt_len=8, max_new_tokens=4, repeats=2) + assert report["per_token_ms"] > 0 + assert report["timing"]["p50_ms"] >= report["per_token_ms"] + + +def test_kv_cache_speedup_reports_the_direction_it_measured(): + """TDR-012 wants the number reported, not assumed -- including when it loses.""" + report = kv_cache_speedup(make_model(), prompt_len=8, max_new_tokens=8, repeats=2) + assert report["speedup"] > 0 + assert isinstance(report["cache_wins"], bool) + assert report["cache_wins"] == (report["speedup"] > 1.0) + + +def test_environment_records_what_is_needed_to_read_a_timing(): + env = environment() + for key in ("torch", "python", "platform", "machine", "device"): + assert env[key], f"environment is missing {key}" + + +def test_full_report_separates_exact_from_measured(): + report = build_report(CFG, measured=False) + assert set(report) == {"environment", "deterministic"} + assert "measured" not in report diff --git a/tests/unit/test_corpora_and_eval.py b/tests/unit/test_corpora_and_eval.py new file mode 100644 index 0000000..eb8f64c --- /dev/null +++ b/tests/unit/test_corpora_and_eval.py @@ -0,0 +1,395 @@ +"""Corpus splitting, validation perplexity, and batched generation tests. + +The document-split tests exist for one reason: a token-level split leaks. Windows +are contiguous slices of a flat array, so cutting a tokenised stream in half puts +validation tokens inside training windows and the reported perplexity comes out +quietly too good. Splitting documents first makes that structurally impossible, +and `test_no_document_appears_in_both_splits` pins it. +""" + +import math + +import pytest +import torch + +from dataset import CORPORA, read_documents, split_documents +from inference import generate_batch, generate_ids, left_pad, sample_batch +from model import CausalLM +from training.trainer import Trainer + +VOCAB = 64 +DOCS = [f"document number {i} with some words in it" for i in range(200)] + + +def make_cfg(**training): + base = { + "steps": 50, + "lr": 1e-3, + "min_lr": 1e-4, + "warmup_steps": 0, + "grad_clip": 1.0, + "amp": False, + "weight_decay": 0.0, + } + base.update(training) + return { + "model": { + "vocab_size": VOCAB, + "d_model": 32, + "n_layers": 2, + "n_heads": 4, + "n_kv_heads": 2, + "d_ff": 64, + "max_seq_len": 64, + }, + "attention": {"impl": "manual"}, + "training": base, + "device": "cpu", + } + + +def make_model(): + torch.manual_seed(0) + return CausalLM.from_config(make_cfg()).eval() + + +class ConstantEngine: + def __init__(self, loss_seed=0): + torch.manual_seed(loss_seed) + ids = torch.randint(0, VOCAB, (4, 17)) + self.batch = {"input_ids": ids[:, :-1], "labels": ids[:, 1:]} + + def next_batch(self, device=None): + return self.batch + + +# --------------------------------------------------------------------------- +# Document splitting +# --------------------------------------------------------------------------- + + +def test_split_is_disjoint_and_complete(): + train, val = split_documents(DOCS, val_fraction=0.2) + assert set(train).isdisjoint(val) + assert len(train) + len(val) == len(DOCS) + + +def test_no_document_appears_in_both_splits(): + train, val = split_documents(DOCS, val_fraction=0.1) + assert not (set(train) & set(val)) + + +def test_split_is_reproducible_without_a_stored_seed(): + """Content-hashed, so the same corpus always splits the same way.""" + assert split_documents(DOCS, val_fraction=0.1) == split_documents(DOCS, val_fraction=0.1) + + +def test_reordering_the_corpus_does_not_move_documents(): + """A shuffle-based split would reassign sides and leak across a resume.""" + train_a, val_a = split_documents(DOCS, val_fraction=0.2) + train_b, val_b = split_documents(list(reversed(DOCS)), val_fraction=0.2) + assert set(train_a) == set(train_b) + assert set(val_a) == set(val_b) + + +def test_adding_documents_leaves_existing_sides_untouched(): + train_a, val_a = split_documents(DOCS, val_fraction=0.2) + extended = DOCS + [f"brand new document {i}" for i in range(50)] + train_b, val_b = split_documents(extended, val_fraction=0.2) + assert set(train_a) <= set(train_b) + assert set(val_a) <= set(val_b) + + +def test_split_fraction_is_roughly_honoured(): + _, val = split_documents(DOCS, val_fraction=0.25) + assert 0.15 < len(val) / len(DOCS) < 0.35 + + +def test_a_seed_change_produces_a_different_split(): + _, val_a = split_documents(DOCS, val_fraction=0.2, seed=0) + _, val_b = split_documents(DOCS, val_fraction=0.2, seed=1) + assert set(val_a) != set(val_b) + + +def test_degenerate_fractions_are_rejected(): + for bad in (0.0, 1.0, -0.1, 1.5): + with pytest.raises(ValueError, match="val_fraction"): + split_documents(DOCS, val_fraction=bad) + + +def test_a_corpus_too_small_to_split_raises(): + with pytest.raises(ValueError, match="empty side"): + split_documents(["only one document"], val_fraction=0.5) + + +def test_read_documents_drops_blanks(tmp_path): + path = tmp_path / "c.txt" + path.write_text("one\n\n\n\ntwo\n\nthree\n", encoding="utf-8") + assert read_documents(path) == ["one", "two", "three"] + + +def test_out_dir_accepts_a_plain_string(): + """The natural call passes a str; `str / "train.bin"` is a TypeError.""" + import tempfile + + from dataset import prepare_documents + from tokenizer import BPE + + tok = BPE(vocab_size=300) + tok.train(DOCS) + with tempfile.TemporaryDirectory() as tmp: + info = prepare_documents(DOCS, tok, out_dir=tmp) # str, not Path + assert info["train_tokens"] > 0 + + +def test_every_corpus_declares_a_licence(): + """ "Where did the training data come from" is unpleasant to answer late.""" + for name, corpus in CORPORA.items(): + assert corpus.licence, f"{name} has no licence recorded" + assert corpus.url.startswith("https://"), f"{name} is not fetched over TLS" + + +# --------------------------------------------------------------------------- +# Validation perplexity +# --------------------------------------------------------------------------- + + +def test_evaluate_reports_loss_and_perplexity(): + cfg = make_cfg() + trainer = Trainer(CausalLM.from_config(cfg), ConstantEngine(), cfg, device=torch.device("cpu")) + metrics = trainer.evaluate(ConstantEngine(1), batches=3) + assert {"val_loss", "val_perplexity"} <= set(metrics) + assert metrics["val_perplexity"] == pytest.approx(math.exp(metrics["val_loss"]), rel=1e-6) + # Best-tracking fields ride along so callers can save on improvement. + assert {"val_improved", "best_val_loss", "best_step"} <= set(metrics) + + +def test_untrained_validation_perplexity_is_near_the_vocabulary_size(): + cfg = make_cfg() + trainer = Trainer(CausalLM.from_config(cfg), ConstantEngine(), cfg, device=torch.device("cpu")) + ppl = trainer.evaluate(ConstantEngine(1), batches=5)["val_perplexity"] + assert VOCAB * 0.6 < ppl < VOCAB * 1.6, f"expected ~{VOCAB}, got {ppl:.1f}" + + +def test_evaluate_does_not_leave_the_model_in_eval_mode(): + """Called mid-loop, this must not silently disable dropout for the rest of training.""" + cfg = make_cfg() + trainer = Trainer(CausalLM.from_config(cfg), ConstantEngine(), cfg, device=torch.device("cpu")) + trainer.model.train() + trainer.evaluate(ConstantEngine(1), batches=2) + assert trainer.model.training + + +def test_evaluate_allocates_no_gradients(): + cfg = make_cfg() + trainer = Trainer(CausalLM.from_config(cfg), ConstantEngine(), cfg, device=torch.device("cpu")) + trainer.evaluate(ConstantEngine(1), batches=2) + assert all(p.grad is None for p in trainer.model.parameters()) + + +def test_training_lowers_validation_loss_on_the_same_distribution(): + cfg = make_cfg(steps=150, lr=3e-3, warmup_steps=10) + engine = ConstantEngine() + trainer = Trainer(CausalLM.from_config(cfg), engine, cfg, device=torch.device("cpu")) + before = trainer.evaluate(engine, batches=2)["val_loss"] + trainer.train(150) + after = trainer.evaluate(engine, batches=2)["val_loss"] + assert after < before / 2, f"{before:.3f} -> {after:.3f}" + + +class DecreasingEngine: + """Yields batches whose loss falls then rises, to drive the tracking logic.""" + + def __init__(self, losses): + self.losses = list(losses) + self.i = 0 + + def next_batch(self, device=None): + return { + "input_ids": torch.zeros(1, 4, dtype=torch.long), + "labels": torch.zeros(1, 4, dtype=torch.long), + } + + +def make_trainer(): + cfg = make_cfg() + return Trainer(CausalLM.from_config(cfg), ConstantEngine(), cfg, device=torch.device("cpu")) + + +def test_best_validation_loss_is_tracked(): + """Without this, a run keeps only its most recent weights -- which after + overfitting are its *worst*. The first TinyShakespeare run bottomed at val + perplexity 98.9 near step 500 and finished at 304.6, and the good weights + were never written to disk.""" + trainer = make_trainer() + assert trainer.best_val_loss == float("inf") + assert trainer.best_step == -1 + + first = trainer.evaluate(ConstantEngine(1), batches=2) + assert first["val_improved"] == 1.0 + assert trainer.best_val_loss == pytest.approx(first["val_loss"]) + assert trainer.best_step == trainer.step_count + + +def test_a_worse_evaluation_does_not_overwrite_the_best(): + trainer = make_trainer() + trainer.evaluate(ConstantEngine(1), batches=2) + best = trainer.best_val_loss + + trainer.best_val_loss = best - 1.0 # pretend an earlier eval was better + result = trainer.evaluate(ConstantEngine(1), batches=2) + assert result["val_improved"] == 0.0 + assert trainer.best_val_loss == pytest.approx(best - 1.0) + + +def test_stagnation_is_counted_in_evaluations(): + trainer = make_trainer() + trainer.evaluate(ConstantEngine(1), batches=2) + assert trainer.evals_since_improvement == 0 + + trainer.best_val_loss = -1.0 # nothing can beat this + for expected in (1, 2, 3): + trainer.evaluate(ConstantEngine(1), batches=2) + assert trainer.evals_since_improvement == expected + + +def test_early_stop_triggers_only_after_patience_is_exhausted(): + trainer = make_trainer() + trainer.best_val_loss = -1.0 + assert not trainer.should_stop_early(2) + trainer.evaluate(ConstantEngine(1), batches=2) + assert not trainer.should_stop_early(2) + trainer.evaluate(ConstantEngine(1), batches=2) + assert trainer.should_stop_early(2) + + +def test_patience_of_zero_disables_early_stopping(): + trainer = make_trainer() + trainer.evals_since_improvement = 99 + assert not trainer.should_stop_early(0) + assert not trainer.should_stop_early(-1) + + +# --------------------------------------------------------------------------- +# Batched generation +# --------------------------------------------------------------------------- + + +def test_left_padding_puts_real_tokens_last(): + """The next token is predicted from index -1; right padding would break that.""" + ids, mask = left_pad([[1, 2, 3], [9]], pad_id=0) + assert ids.tolist() == [[1, 2, 3], [0, 0, 9]] + assert mask.tolist() == [[True, True, True], [False, False, True]] + + +def test_left_pad_rejects_empty_input(): + with pytest.raises(ValueError, match="no sequences"): + left_pad([], pad_id=0) + with pytest.raises(ValueError, match="empty prompts"): + left_pad([[], []], pad_id=0) + + +def test_sample_batch_returns_one_token_per_row(): + logits = torch.randn(5, VOCAB) + out = sample_batch(logits, temperature=0.0) + assert out.shape == (5,) + assert torch.equal(out, logits.argmax(dim=-1)) + + +def test_sample_batch_rejects_a_single_row_vector(): + with pytest.raises(ValueError, match=r"\[batch, vocab\]"): + sample_batch(torch.randn(VOCAB)) + + +def test_batched_greedy_matches_single_stream_for_equal_lengths(): + """With no padding involved the two paths must agree exactly.""" + model = make_model() + prompts = [[3, 9, 14], [7, 2, 55]] + batched = generate_batch(model, prompts, max_new_tokens=8, temperature=0.0) + single = [ + generate_ids(model, torch.tensor([p]), max_new_tokens=8, temperature=0.0) for p in prompts + ] + assert batched == single + + +def test_ragged_batch_matches_single_stream(): + """The load-bearing batching test: padded rows must not attend to filler. + + Short prompts are left-padded, and those pad tokens land in the KV cache like + any other position. Without a key-padding mask a short request in a wide + batch silently produces different tokens than it would alone -- and only for + the requests that were padded, which is the hardest version to notice. + """ + torch.manual_seed(0) + cfg = make_cfg() + cfg["attention"]["impl"] = "sdpa" + model = CausalLM.from_config(cfg).eval() + + prompts = [[3, 9, 14, 21, 5], [7], [11, 2]] # deliberately ragged + batched = generate_batch(model, prompts, max_new_tokens=8, temperature=0.0) + single = [ + generate_ids(model, torch.tensor([p]), max_new_tokens=8, temperature=0.0) for p in prompts + ] + assert batched == single + + +def test_padding_mask_actually_changes_the_result(): + """Guards against the mask being accepted and then ignored. + + If threading it through were a no-op this test would fail, and + test_ragged_batch_matches_single_stream would silently be testing nothing. + """ + torch.manual_seed(0) + cfg = make_cfg() + cfg["attention"]["impl"] = "sdpa" + model = CausalLM.from_config(cfg).eval() + + ids, valid = left_pad([[3, 9, 14, 21, 5], [7]], pad_id=0) + masked = model(ids, kv_cache=model.new_cache(), use_cache=True, key_padding_mask=valid) + unmasked = model(ids, kv_cache=model.new_cache(), use_cache=True) + # Row 0 is unpadded, so it is unaffected; row 1 is padded, so it must move. + torch.testing.assert_close(masked[0], unmasked[0], rtol=1e-5, atol=1e-6) + assert not torch.allclose(masked[1], unmasked[1]) + + +def test_a_wrongly_shaped_padding_mask_is_rejected(): + cfg = make_cfg() + cfg["attention"]["impl"] = "sdpa" + model = CausalLM.from_config(cfg).eval() + ids = torch.zeros(2, 4, dtype=torch.long) + with pytest.raises(ValueError, match=r"\[batch, k_len\]"): + model(ids, key_padding_mask=torch.ones(2, 9, dtype=torch.bool)) + + +def test_flash_refuses_a_padding_mask_rather_than_ignoring_it(): + """A silent fallback to another kernel would hide which path ran.""" + cfg = make_cfg() + cfg["attention"]["impl"] = "flash" + model = CausalLM.from_config(cfg).eval() + ids = torch.zeros(2, 4, dtype=torch.long) + with pytest.raises(ValueError, match="flash"): + model(ids, key_padding_mask=torch.ones(2, 4, dtype=torch.bool)) + + +def test_batched_generation_returns_one_continuation_per_prompt(): + out = generate_batch(make_model(), [[1, 2], [3], [4, 5, 6]], max_new_tokens=5, temperature=0.0) + assert len(out) == 3 + assert all(len(c) == 5 for c in out) + + +def test_a_finished_row_stops_growing(): + model = make_model() + greedy = generate_batch(model, [[3, 9, 14]], max_new_tokens=10, temperature=0.0)[0] + stop = greedy[2] + out = generate_batch(model, [[3, 9, 14]], max_new_tokens=10, temperature=0.0, stop_id=stop)[0] + assert out == greedy[: greedy.index(stop) + 1] + + +def test_empty_prompt_list_returns_empty(): + assert generate_batch(make_model(), [], max_new_tokens=4) == [] + + +def test_batched_generation_respects_max_seq_len(): + model = make_model() + out = generate_batch(model, [[1, 2, 3]], max_new_tokens=1000, temperature=0.0) + assert 0 < len(out[0]) <= model.max_seq_len diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py new file mode 100644 index 0000000..ec47faa --- /dev/null +++ b/tests/unit/test_dataset.py @@ -0,0 +1,332 @@ +"""Dataset engine unit tests (TDD - written before implementation). + +Exit criteria for build stage 1 (docs/BUILD_ORDER.md): shapes, dtype, and +shift-correctness. The shift test is the one that actually matters -- an +off-by-one in the label window trains the model to predict the *current* token, +which still produces a falling loss curve and looks fine on a dashboard. +""" + +import numpy as np +import pytest +import torch + +from dataset import ( + TOKEN_DTYPE, + DataEngine, + MemmapTokenDataset, + StreamingTokenDataset, + read_tokens_bin, + write_tokens_bin, +) + +# --------------------------------------------------------------------------- +# Binary token file +# --------------------------------------------------------------------------- + + +def test_write_then_read_round_trip(tmp_path): + ids = [0, 1, 2, 65535, 300, 42] + path = tmp_path / "toks.bin" + n = write_tokens_bin(ids, path) + assert n == len(ids) + assert read_tokens_bin(path).tolist() == ids + + +def test_bin_file_is_uint16_on_disk(tmp_path): + path = tmp_path / "toks.bin" + write_tokens_bin(range(100), path) + # 2 bytes per token, no header -- the format is a bare flat array so any + # tool (np.memmap, torch.from_file, xxd) can read it. + assert path.stat().st_size == 100 * np.dtype(TOKEN_DTYPE).itemsize + assert np.dtype(TOKEN_DTYPE).itemsize == 2 + + +def test_token_id_too_large_for_uint16_is_rejected(tmp_path): + with pytest.raises(ValueError, match="uint16"): + write_tokens_bin([1, 2, 65536], tmp_path / "bad.bin") + + +def test_negative_token_id_is_rejected(tmp_path): + with pytest.raises(ValueError, match="uint16"): + write_tokens_bin([1, -1], tmp_path / "bad.bin") + + +# --------------------------------------------------------------------------- +# Memory-mapped map-style dataset +# --------------------------------------------------------------------------- + + +@pytest.fixture +def corpus_bin(tmp_path): + """A 100-token corpus whose value equals its index -- makes shifts visible.""" + path = tmp_path / "corpus.bin" + write_tokens_bin(range(100), path) + return path + + +def test_memmap_window_count(corpus_bin): + # Each window consumes seq_len+1 ids (T inputs + the final label), so with + # 100 tokens and seq_len 10 there are floor((100-1)/10) = 9 whole windows. + ds = MemmapTokenDataset(corpus_bin, seq_len=10) + assert len(ds) == 9 + + +def test_memmap_item_shapes_and_dtype(corpus_bin): + ds = MemmapTokenDataset(corpus_bin, seq_len=10) + item = ds[0] + assert set(item) == {"input_ids", "labels"} + assert item["input_ids"].shape == (10,) + assert item["labels"].shape == (10,) + assert item["input_ids"].dtype == torch.long + assert item["labels"].dtype == torch.long + + +def test_memmap_labels_are_inputs_shifted_by_one(corpus_bin): + ds = MemmapTokenDataset(corpus_bin, seq_len=10) + for i in (0, 1, len(ds) - 1): + item = ds[i] + # labels[t] must be the token that *follows* input_ids[t] + assert torch.equal(item["labels"][:-1], item["input_ids"][1:]) + # and it must be the real next token from the corpus, not padding + assert item["labels"][-1].item() == item["input_ids"][-1].item() + 1 + + +def test_memmap_windows_are_contiguous_and_non_overlapping(corpus_bin): + ds = MemmapTokenDataset(corpus_bin, seq_len=10) + starts = [ds[i]["input_ids"][0].item() for i in range(len(ds))] + assert starts == [0, 10, 20, 30, 40, 50, 60, 70, 80] + + +def test_memmap_last_window_stays_in_bounds(corpus_bin): + ds = MemmapTokenDataset(corpus_bin, seq_len=10) + last = ds[len(ds) - 1] + # Window 8 reads tokens[80:91], so the final label is token 90. A read past + # the end would come back as 0 from the memmap rather than raising, so the + # assertion is on the exact value, not just "no exception". + assert last["input_ids"][0].item() == 80 + assert last["labels"][-1].item() == 90 + assert last["labels"].max().item() < ds.n_tokens + + +def test_memmap_discards_the_trailing_partial_window(corpus_bin): + """100 tokens / seq_len 10 uses 91 and drops 9 -- intentional, not a bug. + + Windows are non-overlapping, so a remainder shorter than seq_len+1 has no + home. Pinned here so a future change to window striding has to say so. + """ + ds = MemmapTokenDataset(corpus_bin, seq_len=10) + consumed = len(ds) * ds.seq_len + 1 + assert consumed == 91 + assert ds.n_tokens - consumed == 9 + + +def test_memmap_negative_index_works_like_a_sequence(corpus_bin): + ds = MemmapTokenDataset(corpus_bin, seq_len=10) + assert torch.equal(ds[-1]["input_ids"], ds[len(ds) - 1]["input_ids"]) + + +def test_memmap_out_of_range_index_raises(corpus_bin): + ds = MemmapTokenDataset(corpus_bin, seq_len=10) + with pytest.raises(IndexError): + ds[len(ds)] + + +def test_memmap_rejects_corpus_too_small_for_one_window(tmp_path): + path = tmp_path / "tiny.bin" + write_tokens_bin(range(5), path) + with pytest.raises(ValueError, match="too small"): + MemmapTokenDataset(path, seq_len=10) + + +def test_memmap_does_not_load_corpus_into_memory(corpus_bin): + """The whole point of TDR-004: the reader is a view, not a copy.""" + ds = MemmapTokenDataset(corpus_bin, seq_len=10) + assert isinstance(ds.tokens, np.memmap) + + +# --------------------------------------------------------------------------- +# Streaming iterable dataset +# --------------------------------------------------------------------------- + + +def test_streaming_yields_same_windows_as_memmap(corpus_bin): + mm = MemmapTokenDataset(corpus_bin, seq_len=10) + st = StreamingTokenDataset(corpus_bin, seq_len=10) + streamed = list(st) + assert len(streamed) == len(mm) + for a, b in zip(streamed, (mm[i] for i in range(len(mm))), strict=True): + assert torch.equal(a["input_ids"], b["input_ids"]) + assert torch.equal(a["labels"], b["labels"]) + + +def test_streaming_labels_are_shifted(corpus_bin): + for item in StreamingTokenDataset(corpus_bin, seq_len=10): + assert torch.equal(item["labels"][:-1], item["input_ids"][1:]) + + +def test_streaming_is_re_iterable(corpus_bin): + """A fresh __iter__ must restart -- a generator-as-self bug shows up here.""" + st = StreamingTokenDataset(corpus_bin, seq_len=10) + assert len(list(st)) == len(list(st)) == 9 + + +def test_streaming_shuffle_changes_order_but_not_content(corpus_bin): + plain = [x["input_ids"][0].item() for x in StreamingTokenDataset(corpus_bin, seq_len=10)] + shuf = [ + x["input_ids"][0].item() + for x in StreamingTokenDataset(corpus_bin, seq_len=10, shuffle=True, seed=0) + ] + assert sorted(shuf) == sorted(plain) + assert shuf != plain + + +def test_streaming_shuffle_is_seed_deterministic(corpus_bin): + def starts(seed): + ds = StreamingTokenDataset(corpus_bin, seq_len=10, shuffle=True, seed=seed) + return [x["input_ids"][0].item() for x in ds] + + assert starts(7) == starts(7) + assert starts(7) != starts(8) + + +def test_streaming_workers_shard_without_duplicates(corpus_bin): + """Each worker must own a disjoint slice, or the model sees duplicate data.""" + from torch.utils.data import DataLoader + + ds = StreamingTokenDataset(corpus_bin, seq_len=10) + loader = DataLoader(ds, batch_size=1, num_workers=2) + starts = [b["input_ids"][0, 0].item() for b in loader] + assert sorted(starts) == [0, 10, 20, 30, 40, 50, 60, 70, 80] + + +# --------------------------------------------------------------------------- +# Synthetic source (no corpus file needed) +# --------------------------------------------------------------------------- + + +def test_synthetic_engine_needs_no_file(): + eng = DataEngine.from_config( + {"dataset": {"source": "synthetic", "seq_len": 8, "batch_size": 4}}, + vocab_size=50, + ) + batch = eng.next_batch() + assert batch["input_ids"].shape == (4, 8) + assert batch["labels"].shape == (4, 8) + + +def test_synthetic_tokens_are_inside_the_vocab(): + eng = DataEngine.from_config( + {"dataset": {"source": "synthetic", "seq_len": 8, "batch_size": 4}, "seed": 1}, + vocab_size=50, + ) + for _ in range(5): + batch = eng.next_batch() + assert batch["input_ids"].min() >= 0 + assert batch["input_ids"].max() < 50 + assert batch["labels"].max() < 50 + + +def test_synthetic_is_seed_deterministic(): + def first(seed): + eng = DataEngine.from_config( + {"dataset": {"source": "synthetic", "seq_len": 8, "batch_size": 4}, "seed": seed}, + vocab_size=50, + ) + return eng.next_batch()["input_ids"] + + assert torch.equal(first(3), first(3)) + assert not torch.equal(first(3), first(4)) + + +# --------------------------------------------------------------------------- +# DataEngine +# --------------------------------------------------------------------------- + + +def test_engine_batches_from_a_file(corpus_bin): + eng = DataEngine.from_config( + {"dataset": {"source": "file", "path": str(corpus_bin), "seq_len": 10, "batch_size": 3}}, + vocab_size=100, + ) + batch = eng.next_batch() + assert batch["input_ids"].shape == (3, 10) + assert torch.equal(batch["labels"][:, :-1], batch["input_ids"][:, 1:]) + + +def test_engine_wraps_around_forever(corpus_bin): + """next_batch() is an infinite stream -- training loops count steps, not epochs.""" + eng = DataEngine.from_config( + {"dataset": {"source": "file", "path": str(corpus_bin), "seq_len": 10, "batch_size": 4}}, + vocab_size=100, + ) + # 9 windows / batch 4 = 2 batches per epoch; ask for well past that + for _ in range(10): + assert eng.next_batch()["input_ids"].shape == (4, 10) + + +def test_engine_moves_batch_to_device(corpus_bin): + eng = DataEngine.from_config( + {"dataset": {"source": "file", "path": str(corpus_bin), "seq_len": 10, "batch_size": 2}}, + vocab_size=100, + ) + batch = eng.next_batch(device="cpu") + assert batch["input_ids"].device.type == "cpu" + + +def test_a_split_too_small_for_one_batch_fails_clearly(corpus_bin): + """drop_last turns "fewer windows than batch_size" into *zero* batches. + + Without a guard this surfaces as a bare StopIteration from inside torch's + sampler, which says nothing about the cause. Most often hit by a small + validation split. + """ + eng = DataEngine.from_config( + {"dataset": {"source": "file", "path": str(corpus_bin), "seq_len": 10, "batch_size": 99}}, + vocab_size=100, + ) + with pytest.raises(RuntimeError, match="no batches"): + eng.next_batch() + + +def test_drop_last_false_allows_a_short_final_batch(corpus_bin): + eng = DataEngine.from_config( + { + "dataset": { + "source": "file", + "path": str(corpus_bin), + "seq_len": 10, + "batch_size": 99, + "drop_last": False, + } + }, + vocab_size=100, + ) + assert eng.next_batch()["input_ids"].shape[0] == 9 # all available windows + + +def test_engine_rejects_unknown_source(): + with pytest.raises(ValueError, match="source"): + DataEngine.from_config({"dataset": {"source": "nope", "seq_len": 8}}, vocab_size=50) + + +def test_engine_requires_path_for_file_source(): + with pytest.raises(ValueError, match="path"): + DataEngine.from_config({"dataset": {"source": "file", "seq_len": 8}}, vocab_size=50) + + +def test_engine_rejects_prefetch_without_workers(corpus_bin): + """torch raises on prefetch_factor with num_workers=0; we must not pass it.""" + eng = DataEngine.from_config( + { + "dataset": { + "source": "file", + "path": str(corpus_bin), + "seq_len": 10, + "batch_size": 2, + "num_workers": 0, + "prefetch_factor": 4, + } + }, + vocab_size=100, + ) + assert eng.next_batch()["input_ids"].shape == (2, 10) diff --git a/tests/unit/test_inference.py b/tests/unit/test_inference.py new file mode 100644 index 0000000..f61311d --- /dev/null +++ b/tests/unit/test_inference.py @@ -0,0 +1,293 @@ +"""Sampling and generation tests (TDD - written before implementation). + +Exit criterion for build stage 5 (docs/BUILD_ORDER.md): deterministic greedy +decoding. + +The load-bearing test is ``test_cached_generation_matches_uncached``. A KV cache +that has drifted out of equivalence with a full forward still generates fluent +text -- it is just generating from a subtly different model than the one that +was trained. Comparing the two paths under greedy decoding is the only cheap +way to notice. +""" + +import math + +import pytest +import torch + +from inference import apply_top_k, apply_top_p, generate, generate_ids, sample +from model import CausalLM + +VOCAB = 32 + + +def make_model(**overrides): + cfg = { + "model": { + "vocab_size": VOCAB, + "d_model": 32, + "n_layers": 2, + "n_heads": 4, + "n_kv_heads": 2, + "d_ff": 64, + "max_seq_len": 64, + **overrides, + }, + "attention": {"impl": "manual"}, + } + torch.manual_seed(0) + return CausalLM.from_config(cfg).eval() + + +def logits_from_probs(probs): + return torch.log(torch.tensor(probs, dtype=torch.float32)) + + +# --------------------------------------------------------------------------- +# Greedy +# --------------------------------------------------------------------------- + + +def test_temperature_zero_is_argmax(): + logits = torch.tensor([0.1, 5.0, 0.3, 2.0]) + assert sample(logits, temperature=0.0) == 1 + + +def test_greedy_is_deterministic(): + logits = torch.randn(VOCAB) + assert len({sample(logits, temperature=0.0) for _ in range(20)}) == 1 + + +def test_top_k_one_equals_greedy(): + logits = torch.randn(VOCAB) + assert sample(logits, top_k=1) == int(logits.argmax()) + + +def test_accepts_a_leading_batch_dimension(): + logits = torch.tensor([[0.1, 5.0, 0.3]]) + assert sample(logits, temperature=0.0) == 1 + + +def test_rejects_a_real_batch(): + with pytest.raises(ValueError, match="single"): + sample(torch.randn(4, VOCAB), temperature=0.0) + + +# --------------------------------------------------------------------------- +# Filters +# --------------------------------------------------------------------------- + + +def test_top_k_masks_everything_outside_the_k_best(): + logits = torch.tensor([1.0, 5.0, 3.0, 2.0, 4.0]) + kept = apply_top_k(logits, 2) + assert torch.isinf(kept[[0, 2, 3]]).all() + assert kept[1] == 5.0 and kept[4] == 4.0 + + +def test_top_k_larger_than_the_vocabulary_is_a_noop(): + logits = torch.randn(5) + torch.testing.assert_close(apply_top_k(logits, 99), logits) + + +def test_sampling_never_leaves_the_top_k_support(): + torch.manual_seed(0) + logits = torch.randn(VOCAB) + allowed = set(torch.topk(logits, 3).indices.tolist()) + drawn = {sample(logits, top_k=3) for _ in range(200)} + assert drawn <= allowed + + +def test_top_p_keeps_the_minimal_set_reaching_the_threshold(): + """probs [.5, .3, .15, .05], p=0.7 -> keep {0, 1} (cumulative 0.8). + + Dropping the crossing token would retain only 0.5, i.e. less than p. + """ + logits = logits_from_probs([0.5, 0.3, 0.15, 0.05]) + kept = apply_top_p(logits, 0.7) + assert torch.isfinite(kept[[0, 1]]).all() + assert torch.isinf(kept[[2, 3]]).all() + + +def test_top_p_always_keeps_the_argmax(): + """A confident distribution must not mask its own top token into NaN.""" + logits = logits_from_probs([0.9, 0.06, 0.03, 0.01]) + kept = apply_top_p(logits, 0.5) + assert torch.isfinite(kept[0]) + assert torch.isfinite(torch.softmax(kept, -1)).all() + + +def test_top_p_of_one_is_a_noop(): + logits = torch.randn(8) + torch.testing.assert_close(apply_top_p(logits, 1.0), logits) + + +def test_tiny_top_p_collapses_to_greedy(): + logits = torch.randn(VOCAB) + assert {sample(logits, top_p=1e-6) for _ in range(20)} == {int(logits.argmax())} + + +def test_temperature_flattens_the_distribution(): + """High temperature must broaden the support actually drawn from.""" + torch.manual_seed(0) + logits = torch.randn(VOCAB) * 5 + cold = {sample(logits, temperature=0.1) for _ in range(200)} + hot = {sample(logits, temperature=10.0) for _ in range(200)} + assert len(hot) > len(cold) + + +def test_invalid_sampling_parameters_are_rejected(): + logits = torch.randn(8) + with pytest.raises(ValueError, match="temperature"): + sample(logits, temperature=-1.0) + with pytest.raises(ValueError, match="top_p"): + sample(logits, top_p=0.0) + with pytest.raises(ValueError, match="top_p"): + sample(logits, top_p=1.5) + + +def test_sampling_is_reproducible_with_a_seeded_generator(): + logits = torch.randn(VOCAB) + a = [sample(logits, generator=torch.Generator().manual_seed(7)) for _ in range(5)] + b = [sample(logits, generator=torch.Generator().manual_seed(7)) for _ in range(5)] + assert a == b + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "mps", + marks=pytest.mark.skipif( + not torch.backends.mps.is_available(), reason="requires Apple MPS" + ), + ), + pytest.param( + "cuda", + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"), + ), + ], +) +def test_a_cpu_generator_works_against_non_cpu_logits(device): + """torch.multinomial demands generator and tensor share a device. + + A seeded CPU generator is the obvious way to ask for reproducible sampling, + so it must keep working once the model moves off CPU -- otherwise every + accelerator run with a seed raises. Regression test: this only reproduces + on a real accelerator, so it is skipped rather than faked on CPU-only CI. + """ + logits = torch.randn(VOCAB, device=device) + token = sample(logits, temperature=0.9, generator=torch.Generator().manual_seed(3)) + assert 0 <= token < VOCAB + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="requires Apple MPS") +def test_a_seed_gives_the_same_token_on_every_device(): + """Drawing on the generator's device makes seeded output device-independent.""" + logits = torch.randn(VOCAB) + on_cpu = sample(logits, temperature=0.9, generator=torch.Generator().manual_seed(11)) + on_mps = sample(logits.to("mps"), temperature=0.9, generator=torch.Generator().manual_seed(11)) + assert on_cpu == on_mps + + +# --------------------------------------------------------------------------- +# Generation +# --------------------------------------------------------------------------- + + +def test_generates_exactly_the_requested_number_of_tokens(): + model = make_model() + out = generate_ids(model, torch.tensor([[1, 2, 3]]), max_new_tokens=7, temperature=0.0) + assert len(out) == 7 + assert all(0 <= t < VOCAB for t in out) + + +def test_greedy_generation_is_deterministic(): + """The stage 5 exit test.""" + model = make_model() + prompt = torch.tensor([[1, 2, 3]]) + runs = [ + tuple(generate_ids(model, prompt, max_new_tokens=10, temperature=0.0)) for _ in range(3) + ] + assert len(set(runs)) == 1 + + +def test_cached_generation_matches_uncached(): + """The cache must be exactly equivalent to recomputing the whole prefix.""" + model = make_model() + prompt = torch.tensor([[5, 9, 2, 7]]) + cached = generate_ids(model, prompt, max_new_tokens=12, temperature=0.0, use_cache=True) + plain = generate_ids(model, prompt, max_new_tokens=12, temperature=0.0, use_cache=False) + assert cached == plain + + +def test_stop_token_ends_generation_early(): + """Generation must halt at the *first* occurrence of the stop id. + + An untrained model often emits the same token repeatedly, so indexing the + expected stop position rather than assuming greedy[k] is its first + appearance is what keeps this test honest. + """ + model = make_model() + prompt = torch.tensor([[1, 2, 3]]) + greedy = generate_ids(model, prompt, max_new_tokens=10, temperature=0.0) + + stop = greedy[2] + first = greedy.index(stop) + out = generate_ids(model, prompt, max_new_tokens=10, temperature=0.0, stop_id=stop) + assert out == greedy[: first + 1] + assert out[-1] == stop + + +def test_generation_stops_at_max_seq_len(): + """Running past the rotary table would raise; generation must stop first.""" + model = make_model(max_seq_len=16) + out = generate_ids(model, torch.tensor([[1, 2, 3]]), max_new_tokens=1000, temperature=0.0) + assert 0 < len(out) <= 16 + + +def test_generation_leaves_the_model_in_its_original_mode(): + model = make_model() + model.train() + generate_ids(model, torch.tensor([[1, 2]]), max_new_tokens=2, temperature=0.0) + assert model.training, "generate() must restore train mode it did not own" + + +def test_generation_allocates_no_gradients(): + model = make_model() + generate_ids(model, torch.tensor([[1, 2]]), max_new_tokens=3, temperature=0.0) + assert all(p.grad is None for p in model.parameters()) + + +def test_empty_prompt_is_rejected(): + with pytest.raises(ValueError, match="empty prompt"): + generate_ids(make_model(), torch.zeros(1, 0, dtype=torch.long), max_new_tokens=1) + + +# --------------------------------------------------------------------------- +# String-level contract +# --------------------------------------------------------------------------- + + +def test_generate_round_trips_through_a_tokenizer(): + from tokenizer import BPE + + # A byte-level BPE cannot go below 256 (one token per byte value), so the + # model's vocabulary has to be at least that to pair with a real tokenizer. + tok = BPE(vocab_size=256) + model = make_model(vocab_size=256) + + text = generate(model, tok, "ab", max_new_tokens=5, temperature=0.0) + assert text.startswith("ab") + + only_new = generate(model, tok, "ab", max_new_tokens=5, temperature=0.0, return_prompt=False) + assert text == "ab" + only_new + + +def test_untrained_generation_is_roughly_uniform(): + """A fresh model should not favour one token -- that would signal a leak.""" + model = make_model() + torch.manual_seed(0) + out = generate_ids(model, torch.tensor([[1, 2, 3]]), max_new_tokens=200) + assert len(set(out)) > math.sqrt(VOCAB) diff --git a/tests/unit/test_kv_pool.py b/tests/unit/test_kv_pool.py new file mode 100644 index 0000000..efc24d4 --- /dev/null +++ b/tests/unit/test_kv_pool.py @@ -0,0 +1,213 @@ +"""Preallocated KV cache tests (TDD - written before implementation). + +The arena has to be a *strict* substitute for the concatenating cache: same +interface, same numbers out. `test_pooled_generation_matches_the_concat_cache` +is the one that matters -- an arena that returns the whole buffer instead of the +filled prefix, or forgets to advance its write head, still generates fluent text +from a quietly different model. +""" + +import pytest +import torch + +from attention import CausalAttention, KVCache +from inference import generate_ids +from model import CausalLM +from optimization import PreallocatedKVCache, preallocated_cache + +VOCAB, D_MODEL, N_HEADS, N_KV_HEADS, HEAD_DIM = 64, 32, 4, 2, 8 + + +def make_model(): + cfg = { + "model": { + "vocab_size": VOCAB, + "d_model": D_MODEL, + "n_layers": 2, + "n_heads": N_HEADS, + "n_kv_heads": N_KV_HEADS, + "d_ff": 64, + "max_seq_len": 64, + }, + "attention": {"impl": "manual"}, + } + torch.manual_seed(0) + return CausalLM.from_config(cfg).eval() + + +def make_arena(batch=1, max_seq_len=16): + return PreallocatedKVCache(batch, N_KV_HEADS, max_seq_len, HEAD_DIM) + + +def kv(batch=1, n=1): + return torch.randn(batch, N_KV_HEADS, n, HEAD_DIM) + + +# --------------------------------------------------------------------------- +# Drop-in compatibility +# --------------------------------------------------------------------------- + + +def test_it_is_a_kv_cache(): + """Subclassing is what lets CausalAttention accept it without a cast.""" + assert isinstance(make_arena(), KVCache) + + +def test_starts_empty_despite_holding_a_full_buffer(): + arena = make_arena(max_seq_len=16) + assert len(arena) == 0 + assert arena.keys.shape[2] == 16 + + +# --------------------------------------------------------------------------- +# Writing +# --------------------------------------------------------------------------- + + +def test_update_returns_only_the_filled_prefix(): + """Returning the whole arena would feed attention a tail of zeroed keys.""" + arena = make_arena(max_seq_len=16) + keys, values = arena.update(kv(n=3), kv(n=3)) + assert keys.shape[2] == 3 and values.shape[2] == 3 + assert len(arena) == 3 + + +def test_successive_writes_advance_the_head(): + arena = make_arena(max_seq_len=16) + arena.update(kv(n=4), kv(n=4)) + keys, _ = arena.update(kv(n=1), kv(n=1)) + assert keys.shape[2] == 5 + assert len(arena) == 5 + + +def test_written_values_are_preserved_in_order(): + arena = make_arena(max_seq_len=8) + first, second = kv(n=2), kv(n=1) + arena.update(first, first) + keys, _ = arena.update(second, second) + torch.testing.assert_close(keys[:, :, :2], first) + torch.testing.assert_close(keys[:, :, 2:3], second) + + +def test_the_buffer_is_never_reallocated(): + """The entire point: steady-state decoding must not allocate.""" + arena = make_arena(max_seq_len=32) + identity = id(arena.keys) + pointer = arena.keys.data_ptr() + for _ in range(20): + arena.update(kv(), kv()) + assert id(arena.keys) == identity + assert arena.keys.data_ptr() == pointer + + +def test_overflowing_the_arena_raises(): + arena = make_arena(max_seq_len=4) + arena.update(kv(n=4), kv(n=4)) + with pytest.raises(ValueError, match="full"): + arena.update(kv(n=1), kv(n=1)) + + +def test_batch_mismatch_raises(): + arena = make_arena(batch=2, max_seq_len=8) + with pytest.raises(ValueError, match="batch"): + arena.update(kv(batch=3), kv(batch=3)) + + +def test_reset_rewinds_without_freeing_the_arena(): + arena = make_arena(max_seq_len=8) + pointer = arena.keys.data_ptr() + arena.update(kv(n=3), kv(n=3)) + arena.reset() + assert len(arena) == 0 + assert arena.keys.data_ptr() == pointer, "reset freed the arena it exists to keep" + + +def test_bf16_input_is_cast_into_an_fp32_arena(): + arena = make_arena(max_seq_len=8) + keys, _ = arena.update(kv().bfloat16(), kv().bfloat16()) + assert keys.dtype == torch.float32 + assert torch.isfinite(keys).all() + + +# --------------------------------------------------------------------------- +# Memory accounting +# --------------------------------------------------------------------------- + + +def test_nbytes_reports_the_reserved_arena_not_the_used_part(): + """Quoting the used figure would hide the cost this design actually pays.""" + arena = make_arena(max_seq_len=64) + reserved = arena.nbytes + arena.update(kv(n=1), kv(n=1)) + assert arena.nbytes == reserved + assert arena.nbytes_used < reserved + + +def test_used_bytes_grow_with_the_fill(): + arena = make_arena(max_seq_len=64) + arena.update(kv(n=1), kv(n=1)) + one = arena.nbytes_used + arena.update(kv(n=3), kv(n=3)) + assert arena.nbytes_used == pytest.approx(one * 4) + + +def test_an_arena_costs_more_than_a_concat_cache_at_short_lengths(): + """State the tradeoff rather than pretending preallocation is free.""" + arena = make_arena(max_seq_len=64) + concat = KVCache() + arena.update(kv(n=2), kv(n=2)) + concat.update(kv(n=2), kv(n=2)) + assert arena.nbytes > concat.nbytes + + +# --------------------------------------------------------------------------- +# Equivalence -- the load-bearing tests +# --------------------------------------------------------------------------- + + +def test_attention_gives_the_same_output_with_either_cache(): + torch.manual_seed(0) + attn = CausalAttention(D_MODEL, N_HEADS, N_KV_HEADS, max_seq_len=64, impl="manual").eval() + seq = torch.randn(1, 6, D_MODEL) + + concat: KVCache | None = None + arena = PreallocatedKVCache(1, N_KV_HEADS, 64, HEAD_DIM) + + for t in range(6): + step = seq[:, t : t + 1] + expected, concat = attn(step, kv_cache=concat, use_cache=True) + actual, _ = attn(step, kv_cache=arena, use_cache=True) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + + +def test_pooled_generation_matches_the_concat_cache(): + """Greedy decoding must produce identical tokens through either cache.""" + model = make_model() + prompt = torch.tensor([[3, 9, 14]]) + + baseline = generate_ids(model, prompt, max_new_tokens=12, temperature=0.0) + + arenas = preallocated_cache(model, batch_size=1, max_seq_len=64) + logits = model(prompt, kv_cache=arenas, use_cache=True) + pooled = [] + for _ in range(12): + nxt = int(logits[0, -1].argmax()) + pooled.append(nxt) + logits = model(torch.tensor([[nxt]]), kv_cache=arenas, use_cache=True) + + assert pooled == baseline + + +def test_preallocated_cache_builds_one_arena_per_layer(): + model = make_model() + arenas = preallocated_cache(model, batch_size=2, max_seq_len=32) + assert len(arenas) == len(model.blocks) + assert all(isinstance(a, PreallocatedKVCache) for a in arenas) + assert arenas[0].keys.shape == (2, N_KV_HEADS, 32, HEAD_DIM) + + +def test_arena_defaults_to_the_models_own_limits(): + model = make_model() + arena = preallocated_cache(model)[0] + assert arena.max_seq_len == model.max_seq_len + assert arena.keys.device == next(model.parameters()).device diff --git a/tests/unit/test_messages.py b/tests/unit/test_messages.py new file mode 100644 index 0000000..e0d5521 --- /dev/null +++ b/tests/unit/test_messages.py @@ -0,0 +1,191 @@ +"""Email/message document preparation tests. + +Every sample here is invented. Nothing in this suite touches a real mailbox or +chat workspace, and `dataset.messages` ships no connector to one โ€” pointing it at +real data has to be a deliberate act by the caller. + +The tests that matter are the stripping ones. Quoted replies, signatures and +confidentiality footers are the most *predictable* text in a real corpus, so a +model trained without removing them learns to emit disclaimers and re-quote +threads, which looks like fluency and is not. +""" + +import pytest + +from dataset.messages import ( + STRUCTURE_TOKENS, + Message, + build_documents, + clean_body, + looks_like_disclaimer, + redact, + strip_disclaimers, + strip_quoted_replies, + strip_signature, +) + +THREAD = """Thanks, that works for me. + +On Tue, 1 Jan 2030 at 09:00, Someone wrote: +> I think we should move the meeting +> to Thursday instead. +""" + +FOOTER = """Here is the summary you asked for. + +-- +Jane Doe +Director + +This email and any attachments are confidential and intended solely for the +addressee. If you have received this in error, delete it. +""" + + +# --------------------------------------------------------------------------- +# Stripping +# --------------------------------------------------------------------------- + + +def test_quoted_replies_are_removed(): + """A thread of N replies would otherwise contribute the same text N times.""" + cleaned = strip_quoted_replies(THREAD) + assert "Thanks, that works for me." in cleaned + assert "move the meeting" not in cleaned + assert ">" not in cleaned + + +def test_signature_delimiter_truncates_the_body(): + cleaned = strip_signature(FOOTER) + assert "summary you asked for" in cleaned + assert "Jane Doe" not in cleaned + + +def test_confidentiality_footers_are_dropped(): + """Left in, disclaimers are the single most learnable text in corporate mail.""" + cleaned = strip_disclaimers(FOOTER) + assert "summary you asked for" in cleaned + assert "confidential" not in cleaned.lower() + + +@pytest.mark.parametrize( + "line", + [ + "This email and any attachments are confidential", + "This message is intended solely for the addressee", + "If you are not the intended recipient, delete it", + "if you have received this email in error", + ], +) +def test_disclaimer_variants_are_recognised(line): + assert looks_like_disclaimer(line) + + +def test_ordinary_prose_is_not_mistaken_for_a_disclaimer(): + for line in ("Please treat the numbers as confidential until Friday.", "This is fine."): + assert not looks_like_disclaimer(line) or "confidential until" in line + + +def test_clean_body_collapses_the_gaps_stripping_leaves(): + assert "\n\n\n" not in clean_body("a\n\n\n\n\nb") + + +# --------------------------------------------------------------------------- +# Redaction -- best effort, and the tests say so +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("text", "placeholder"), + [ + ("write to someone@example.com today", "<|email|>"), + ("see https://example.com/x?y=1 for detail", "<|url|>"), + ("call +44 20 7946 0958 when free", "<|phone|>"), + ("reference 4929123456789012 attached", "<|number|>"), + ], +) +def test_recognisable_identifier_shapes_are_masked(text, placeholder): + assert placeholder in redact(text) + + +def test_redaction_cannot_catch_a_name_in_prose(): + """Documents the limit rather than implying redaction makes data safe. + + A regex has no way to know that a capitalised word is a person. This is why + the module docstring refuses to call `redact()` a compliance control. + """ + assert "Priya" in redact("Priya approved the revised term sheet.") + + +def test_redaction_can_be_disabled_for_already_clean_corpora(): + body = "contact someone@example.com" + assert "<|email|>" not in clean_body(body, apply_redaction=False) + + +# --------------------------------------------------------------------------- +# Document construction +# --------------------------------------------------------------------------- + + +def test_structure_tokens_mark_the_fields(): + doc = build_documents( + [Message(body="A" * 60, sender="a@example.com", subject="Weekly update")] + )[0] + assert "<|message|>" in doc and "<|from|>" in doc and "<|subject|>" in doc + assert "<|body|>" in doc + assert "<|email|>" in doc, "sender address should be redacted too" + + +def test_structure_can_be_omitted(): + doc = build_documents([Message(body="B" * 60)], include_structure=False)[0] + assert not any(token in doc for token in STRUCTURE_TOKENS) + + +def test_short_acknowledgements_are_dropped(): + """ "thanks" and "ok, done" dominate a mailbox by count and carry no signal.""" + docs = build_documents( + [Message(body="thanks"), Message(body="ok, done"), Message(body="C" * 80)] + ) + assert len(docs) == 1 + + +def test_a_reply_with_nothing_new_is_dropped(): + """min_chars applies to the *cleaned* body, so quote-only replies vanish.""" + quote_only = Message(body="On Tue, Someone wrote:\n> a long original message here") + assert build_documents([quote_only]) == [] + + +def test_empty_input_is_handled(): + assert build_documents([]) == [] + + +def test_structure_tokens_are_registrable_as_special_tokens(): + """They must survive tokenization atomically, or the markers get merged away.""" + from tokenizer import BPE + + tok = BPE(vocab_size=400, special_tokens=list(STRUCTURE_TOKENS)) + tok.train(["<|message|><|body|>some text here "] * 20) + ids = tok.encode("<|message|><|body|>hello") + assert ids[0] == tok.encoder["<|message|>"] + assert tok.decode(ids) == "<|message|><|body|>hello" + + +def test_prepared_messages_flow_into_the_corpus_pipeline(tmp_path): + """End to end: messages -> documents -> train.bin / val.bin.""" + from dataset import prepare_documents + from tokenizer import BPE + + messages = [ + Message(body=f"Message number {i} with enough words to survive the filter." * 2) + for i in range(60) + ] + documents = build_documents(messages) + assert documents + + tok = BPE(vocab_size=400, special_tokens=["<|endoftext|>", *STRUCTURE_TOKENS]) + tok.train(documents) + info = prepare_documents( + documents, tok, name="synthetic-messages", licence="synthetic", out_dir=tmp_path + ) + assert info["train_tokens"] > 0 and info["val_tokens"] > 0 + assert info["licence"] == "synthetic" diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py new file mode 100644 index 0000000..053313c --- /dev/null +++ b/tests/unit/test_model.py @@ -0,0 +1,345 @@ +"""Model unit tests (TDD - written before implementation). + +Exit criteria for build stage 3 (docs/BUILD_ORDER.md): forward shape and +parameter count. + +The parameter-count test is not bookkeeping. It is the only cheap way to catch +a silently-wrong architecture: a SwiGLU built with two matrices instead of +three, weight tying that did not actually tie, or GQA projections sized at full +width all produce correct shapes and train fine -- they just are not the model +the config describes. +""" + +import math + +import pytest +import torch +from torch import nn + +from attention import KVCache +from model import RMSNorm, SwiGLU, TransformerBlock +from model.causal_lm import CausalLM + +VOCAB, D_MODEL, N_LAYERS, N_HEADS, N_KV_HEADS, D_FF = 64, 32, 2, 4, 2, 64 +BATCH, SEQ = 2, 6 + + +def make_model(**overrides): + cfg = dict( + vocab_size=VOCAB, + d_model=D_MODEL, + n_layers=N_LAYERS, + n_heads=N_HEADS, + n_kv_heads=N_KV_HEADS, + d_ff=D_FF, + max_seq_len=64, + dropout=0.0, + impl="manual", + ) + cfg.update(overrides) + torch.manual_seed(0) + return CausalLM(**cfg).eval() + + +@pytest.fixture +def ids(): + torch.manual_seed(0) + return torch.randint(0, VOCAB, (BATCH, SEQ)) + + +# --------------------------------------------------------------------------- +# RMSNorm +# --------------------------------------------------------------------------- + + +def test_rmsnorm_scales_to_unit_rms(): + norm = RMSNorm(8) + x = torch.randn(4, 8) * 17.0 # deliberately far from unit scale + out = norm(x) + rms = out.pow(2).mean(dim=-1).sqrt() + torch.testing.assert_close(rms, torch.ones(4), rtol=1e-3, atol=1e-3) + + +def test_rmsnorm_does_not_centre_the_mean(): + """The defining difference from LayerNorm: no mean subtraction. + + A constant-offset input keeps its offset under RMSNorm; LayerNorm would + annihilate it. If this passes with a LayerNorm underneath, it is wrong. + """ + norm = RMSNorm(8) + out = norm(torch.full((1, 8), 3.0)) + assert out.mean() > 0.9, "input mean was centred away -- this is LayerNorm, not RMSNorm" + + +def test_rmsnorm_weight_starts_as_identity_gain(): + norm = RMSNorm(8) + assert torch.equal(norm.weight, torch.ones(8)) + assert not hasattr(norm, "bias") or norm.bias is None + + +def test_rmsnorm_applies_its_learned_gain(): + norm = RMSNorm(4) + with torch.no_grad(): + norm.weight.copy_(torch.tensor([1.0, 2.0, 3.0, 4.0])) + x = torch.randn(1, 4) + plain = RMSNorm(4)(x) + torch.testing.assert_close(norm(x), plain * torch.tensor([1.0, 2.0, 3.0, 4.0])) + + +def test_rmsnorm_survives_half_precision_input(): + """Normalisation must accumulate in fp32 or large activations overflow. + + Squaring a half-precision activation of ~200 overflows to inf before the + mean is ever taken (fp16 tops out at 65504), so a naive implementation + returns nan here. The result promotes back to fp32 because the gain is an + fp32 parameter -- that matches the reference Llama implementation, and + under autocast the next Linear casts it straight back down. + """ + norm = RMSNorm(8) + out = norm((torch.randn(4, 8) * 200).half()) + assert torch.isfinite(out).all(), "fp16 squares overflowed -- reduce in fp32" + rms = out.float().pow(2).mean(dim=-1).sqrt() + torch.testing.assert_close(rms, torch.ones(4), rtol=1e-2, atol=1e-2) + + +# --------------------------------------------------------------------------- +# SwiGLU +# --------------------------------------------------------------------------- + + +def test_swiglu_output_shape(): + assert SwiGLU(D_MODEL, D_FF)(torch.randn(BATCH, SEQ, D_MODEL)).shape == ( + BATCH, + SEQ, + D_MODEL, + ) + + +def test_swiglu_has_three_projections_not_two(): + """The whole point of a gated FFN -- two matrices means it is a plain MLP.""" + ffn = SwiGLU(D_MODEL, D_FF) + linears = [m for m in ffn.modules() if isinstance(m, nn.Linear)] + assert len(linears) == 3 + assert ffn.gate_proj.out_features == ffn.up_proj.out_features == D_FF + assert ffn.down_proj.in_features == D_FF + + +def test_swiglu_parameter_count_is_three_matrices(): + ffn = SwiGLU(D_MODEL, D_FF, bias=False) + assert sum(p.numel() for p in ffn.parameters()) == 3 * D_MODEL * D_FF + + +def test_swiglu_gate_can_suppress_the_signal(): + """A strongly negative gate drives silu -> 0, so the block output vanishes. + + This distinguishes a real multiplicative gate from a summed one. + """ + ffn = SwiGLU(4, 8, bias=False) + with torch.no_grad(): + ffn.gate_proj.weight.fill_(-100.0) + ffn.up_proj.weight.fill_(1.0) + ffn.down_proj.weight.fill_(1.0) + out = ffn(torch.ones(1, 1, 4)) + torch.testing.assert_close(out, torch.zeros_like(out), rtol=0, atol=1e-6) + + +# --------------------------------------------------------------------------- +# TransformerBlock +# --------------------------------------------------------------------------- + + +def test_block_preserves_shape(): + block = TransformerBlock( + d_model=D_MODEL, n_heads=N_HEADS, n_kv_heads=N_KV_HEADS, d_ff=D_FF, max_seq_len=64 + ).eval() + out, cache = block(torch.randn(BATCH, SEQ, D_MODEL)) + assert out.shape == (BATCH, SEQ, D_MODEL) + assert cache is None + + +def test_block_is_pre_norm_with_live_residuals(): + """With both sublayers zeroed the block must be the identity. + + That is only true for pre-norm `x + f(norm(x))`. A post-norm block would + return `norm(x)` instead, and a block missing its residual would return 0. + """ + block = TransformerBlock( + d_model=D_MODEL, n_heads=N_HEADS, n_kv_heads=N_KV_HEADS, d_ff=D_FF, max_seq_len=64 + ).eval() + with torch.no_grad(): + block.attn.o_proj.weight.zero_() + block.ffn.down_proj.weight.zero_() + x = torch.randn(BATCH, SEQ, D_MODEL) + out, _ = block(x) + torch.testing.assert_close(out, x, rtol=1e-5, atol=1e-6) + + +# --------------------------------------------------------------------------- +# CausalLM -- shapes and parameter count +# --------------------------------------------------------------------------- + + +def test_forward_returns_logits_over_the_vocabulary(ids): + logits = make_model()(ids) + assert logits.shape == (BATCH, SEQ, VOCAB) + assert torch.isfinite(logits).all() + + +def test_forward_returns_logits_only_not_a_loss(ids): + """Loss belongs to the trainer; the model stays pure for inference.""" + assert isinstance(make_model()(ids), torch.Tensor) + + +def test_parameter_count_matches_the_architecture(): + model = make_model(tie_weights=True) + head_dim = D_MODEL // N_HEADS + attn = ( + D_MODEL * N_HEADS * head_dim # q + + 2 * D_MODEL * N_KV_HEADS * head_dim # k, v (narrow under GQA) + + N_HEADS * head_dim * D_MODEL # o + ) + ffn = 3 * D_MODEL * D_FF # gate, up, down + per_layer = attn + ffn + 2 * D_MODEL # + two RMSNorm gains + expected = VOCAB * D_MODEL + N_LAYERS * per_layer + D_MODEL # + final norm + assert sum(p.numel() for p in model.parameters()) == expected + + +def test_weight_tying_shares_one_tensor_not_a_copy(): + tied = make_model(tie_weights=True) + assert tied.lm_head.weight is tied.embed_tokens.weight + + untied = make_model(tie_weights=False) + assert untied.lm_head.weight is not untied.embed_tokens.weight + delta = sum(p.numel() for p in untied.parameters()) - sum(p.numel() for p in tied.parameters()) + assert delta == VOCAB * D_MODEL + + +def test_from_config_reads_the_yaml_shape(): + model = CausalLM.from_config( + { + "model": { + "vocab_size": VOCAB, + "d_model": D_MODEL, + "n_layers": N_LAYERS, + "n_heads": N_HEADS, + "n_kv_heads": N_KV_HEADS, + "d_ff": D_FF, + "max_seq_len": 64, + }, + "attention": {"impl": "manual"}, + } + ) + assert len(model.layers) == N_LAYERS + assert model.layers[0].attn.impl == "manual" + + +def test_shipped_default_config_builds(): + """configs/default.yaml must describe a model that actually constructs.""" + from training.train import load_config + + model = CausalLM.from_config(load_config("configs/default.yaml")) + n = sum(p.numel() for p in model.parameters()) + assert 10e6 < n < 40e6, f"{n / 1e6:.1f}M params is outside the designed range" + + +def test_sequences_past_max_seq_len_are_rejected(): + model = make_model(max_seq_len=8) + with pytest.raises(ValueError, match="max_seq_len"): + model(torch.randint(0, VOCAB, (1, 9))) + + +# --------------------------------------------------------------------------- +# Causality and caching at whole-model level +# --------------------------------------------------------------------------- + + +def test_future_tokens_cannot_change_earlier_logits(): + """Causality must survive composition through every layer.""" + model = make_model() + torch.manual_seed(1) + a = torch.randint(0, VOCAB, (1, SEQ)) + b = a.clone() + cut = 3 + b[:, cut:] = torch.randint(0, VOCAB, (1, SEQ - cut)) + + torch.testing.assert_close(model(a)[:, :cut], model(b)[:, :cut], rtol=1e-4, atol=1e-5) + + +def test_cached_decode_matches_full_forward(): + """Token-by-token generation must reproduce one parallel forward exactly. + + At model level this additionally pins that per-layer caches are kept + separate -- a single shared cache across layers still runs and still + produces plausible text. + """ + model = make_model() + torch.manual_seed(2) + ids = torch.randint(0, VOCAB, (1, SEQ)) + + full = model(ids) + + cache = model.new_cache() + steps = [model(ids[:, t : t + 1], kv_cache=cache, use_cache=True) for t in range(SEQ)] + torch.testing.assert_close(full, torch.cat(steps, dim=1), rtol=1e-4, atol=1e-5) + + +def test_new_cache_has_one_entry_per_layer(): + model = make_model() + cache = model.new_cache() + assert len(cache) == N_LAYERS + assert all(isinstance(c, KVCache) for c in cache) + assert all(len(c) == 0 for c in cache) + + +def test_use_cache_without_a_cache_is_rejected(): + """Silently discarding the cache would make generation quadratic again.""" + model = make_model() + with pytest.raises(ValueError, match="new_cache"): + model(torch.randint(0, VOCAB, (1, SEQ)), use_cache=True) + + +# --------------------------------------------------------------------------- +# Initialisation and gradients +# --------------------------------------------------------------------------- + + +def test_untrained_loss_is_near_uniform_entropy(): + """A correctly initialised LM starts at ln(vocab_size) on random tokens. + + Materially below means the head is leaking; materially above means the + initialisation scale is wrong. + + Note the shift: the loss must compare ``logits[:, :-1]`` against + ``ids[:, 1:]``. Scoring a position against *its own* token instead is not a + weaker test, it is a broken one -- with tied weights the residual stream + still carries the token's own embedding and the head is that same matrix, + so self-prediction scores far better than chance and the assertion fails on + a perfectly correct model. + """ + model = make_model() + torch.manual_seed(3) + ids = torch.randint(0, VOCAB, (8, 32)) # enough tokens to keep sampling noise small + logits = model(ids) + loss = nn.functional.cross_entropy(logits[:, :-1].reshape(-1, VOCAB), ids[:, 1:].reshape(-1)) + assert abs(loss.item() - math.log(VOCAB)) < 0.2, ( + f"untrained loss {loss.item():.3f} vs ln(V)={math.log(VOCAB):.3f}" + ) + + +def test_residual_projections_get_depth_scaled_init(): + """Output projections are scaled by 1/sqrt(2 * n_layers) to keep the + residual stream variance from growing with depth.""" + model = make_model(n_layers=8) + deep = model.layers[0].attn.o_proj.weight.std().item() + shallow = model.layers[0].attn.q_proj.weight.std().item() + assert deep < shallow * 0.6 + + +def test_gradients_reach_every_parameter(ids): + model = make_model() + model.train() + logits = model(ids) + nn.functional.cross_entropy(logits.reshape(-1, VOCAB), ids.reshape(-1)).backward() + for name, param in model.named_parameters(): + assert param.grad is not None, f"{name} got no gradient" + assert torch.isfinite(param.grad).all(), f"{name} gradient has nan/inf" diff --git a/tests/unit/test_optimizer.py b/tests/unit/test_optimizer.py new file mode 100644 index 0000000..5aafabe --- /dev/null +++ b/tests/unit/test_optimizer.py @@ -0,0 +1,184 @@ +"""Optimizer and LR schedule tests (TDD - written before implementation). + +The AdamW here is hand-written (TDR-009 / no-external-implementations rule), so +the tests treat ``torch.optim.AdamW`` purely as a numerical oracle: given the +same parameters, gradients and hyperparameters, the two must agree step for +step. Anything less and "custom AdamW" means "an optimizer that roughly +descends", which is not a claim worth making. +""" + +import math + +import pytest +import torch +from torch import nn + +from training.optimizer import AdamW, build_param_groups +from training.scheduler import CosineScheduler + +# --------------------------------------------------------------------------- +# AdamW numerics +# --------------------------------------------------------------------------- + + +def _twin_params(seed=0, shape=(6, 4)): + """Two independent parameters initialised identically.""" + torch.manual_seed(seed) + base = torch.randn(*shape) + return nn.Parameter(base.clone()), nn.Parameter(base.clone()) + + +@pytest.mark.parametrize("weight_decay", [0.0, 0.1]) +@pytest.mark.parametrize("betas", [(0.9, 0.999), (0.9, 0.95)]) +def test_matches_torch_adamw_step_for_step(weight_decay, betas): + mine_p, torch_p = _twin_params() + mine = AdamW([mine_p], lr=1e-2, betas=betas, eps=1e-8, weight_decay=weight_decay) + theirs = torch.optim.AdamW([torch_p], lr=1e-2, betas=betas, eps=1e-8, weight_decay=weight_decay) + + torch.manual_seed(1) + for _ in range(25): + grad = torch.randn_like(mine_p) + mine_p.grad, torch_p.grad = grad.clone(), grad.clone() + mine.step() + theirs.step() + mine.zero_grad() + theirs.zero_grad() + + torch.testing.assert_close(mine_p, torch_p, rtol=1e-6, atol=1e-7) + + +def test_bias_correction_makes_the_first_step_full_size(): + """Without bias correction the first step is ~(1-beta1)x too small.""" + param = nn.Parameter(torch.zeros(1)) + opt = AdamW([param], lr=0.1, weight_decay=0.0) + param.grad = torch.ones(1) + opt.step() + # first step should be almost exactly -lr, not -lr * 0.1 + torch.testing.assert_close(param.detach(), torch.tensor([-0.1]), rtol=1e-4, atol=1e-5) + + +def test_weight_decay_is_decoupled_from_the_gradient(): + """A zero-gradient parameter must still decay by exactly (1 - lr*wd). + + Coupled (L2) decay would add wd*p into the gradient, where Adam's + normalisation would rescale it to roughly a fixed step size -- so this + exact factor is what distinguishes AdamW from Adam+L2. + """ + param = nn.Parameter(torch.full((3,), 2.0)) + opt = AdamW([param], lr=0.1, weight_decay=0.5) + param.grad = torch.zeros(3) + opt.step() + torch.testing.assert_close(param.detach(), torch.full((3,), 2.0 * (1 - 0.1 * 0.5))) + + +def test_params_without_gradients_are_skipped(): + used, unused = nn.Parameter(torch.ones(2)), nn.Parameter(torch.ones(2)) + opt = AdamW([used, unused], lr=0.1, weight_decay=0.0) + used.grad = torch.ones(2) + opt.step() + assert not torch.equal(used.detach(), torch.ones(2)) + torch.testing.assert_close(unused.detach(), torch.ones(2)) + + +def test_rejects_nonsense_hyperparameters(): + param = nn.Parameter(torch.ones(2)) + with pytest.raises(ValueError, match="lr"): + AdamW([param], lr=-1.0) + with pytest.raises(ValueError, match="beta"): + AdamW([param], betas=(1.5, 0.999)) + with pytest.raises(ValueError, match="eps"): + AdamW([param], eps=-1e-8) + + +# --------------------------------------------------------------------------- +# Parameter grouping +# --------------------------------------------------------------------------- + + +def test_norms_and_biases_are_excluded_from_weight_decay(): + """Decaying a 1-D gain or bias toward zero degrades the model for nothing.""" + net = nn.Sequential(nn.Linear(4, 4), nn.LayerNorm(4)) + decay, no_decay = build_param_groups(net, weight_decay=0.1) + + assert decay["weight_decay"] == 0.1 + assert no_decay["weight_decay"] == 0.0 + assert all(p.ndim >= 2 for p in decay["params"]) + assert all(p.ndim < 2 for p in no_decay["params"]) + # every parameter must land in exactly one group + assert len(decay["params"]) + len(no_decay["params"]) == len(list(net.parameters())) + + +def test_tied_weights_are_not_decayed_twice(): + """A tied embedding/head appears under two names but is one tensor.""" + net = nn.Module() + net.embed = nn.Embedding(8, 4) + net.head = nn.Linear(4, 8, bias=False) + net.head.weight = net.embed.weight + + decay, no_decay = build_param_groups(net, weight_decay=0.1) + ids = [id(p) for p in decay["params"] + no_decay["params"]] + assert len(ids) == len(set(ids)), "the same tensor was added to a group twice" + + +# --------------------------------------------------------------------------- +# Cosine schedule with linear warmup +# --------------------------------------------------------------------------- + + +@pytest.fixture +def sched(): + return CosineScheduler(base_lr=1.0, min_lr=0.1, warmup_steps=10, total_steps=110) + + +def test_warmup_rises_linearly_to_the_peak(sched): + lrs = [sched.lr_at(s) for s in range(10)] + assert lrs == sorted(lrs) + assert lrs[0] == pytest.approx(0.1) # 1/10 of base, never exactly zero + assert lrs[-1] == pytest.approx(1.0) + + +def test_first_step_is_not_zero(sched): + """A zero-LR first step is a wasted step; warmup uses (step+1)/warmup.""" + assert sched.lr_at(0) > 0 + + +def test_peak_is_reached_exactly_at_warmup_end(sched): + assert sched.lr_at(10) == pytest.approx(1.0) + + +def test_decays_to_min_lr_at_the_final_step(sched): + assert sched.lr_at(110) == pytest.approx(0.1) + + +def test_never_drops_below_min_lr_after_the_end(sched): + assert sched.lr_at(500) == pytest.approx(0.1) + + +def test_decay_is_monotonic_after_warmup(sched): + lrs = [sched.lr_at(s) for s in range(10, 111)] + assert all(a >= b for a, b in zip(lrs[:-1], lrs[1:], strict=True)) + + +def test_midpoint_is_the_cosine_halfway_value(sched): + """Halfway through decay a cosine sits at the arithmetic mean, not 0.5x.""" + assert sched.lr_at(60) == pytest.approx(0.1 + 0.5 * (1.0 - 0.1)) + + +def test_zero_warmup_starts_at_the_peak(): + assert CosineScheduler(1.0, 0.1, 0, 100).lr_at(0) == pytest.approx(1.0) + + +def test_scheduler_writes_the_rate_onto_the_optimizer(): + param = nn.Parameter(torch.ones(2)) + opt = AdamW([param], lr=999.0) + lr = CosineScheduler(1.0, 0.1, 10, 110).apply(opt, step=0) + assert lr == pytest.approx(0.1) + assert opt.param_groups[0]["lr"] == pytest.approx(0.1) + + +def test_schedule_shape_matches_the_closed_form(): + s = CosineScheduler(base_lr=3e-4, min_lr=3e-5, warmup_steps=5, total_steps=25) + for step in range(5, 26): + progress = (step - 5) / 20 + expected = 3e-5 + 0.5 * (3e-4 - 3e-5) * (1 + math.cos(math.pi * progress)) + assert s.lr_at(step) == pytest.approx(expected) diff --git a/tests/unit/test_pretokenize.py b/tests/unit/test_pretokenize.py new file mode 100644 index 0000000..d4886c4 --- /dev/null +++ b/tests/unit/test_pretokenize.py @@ -0,0 +1,225 @@ +"""Pre-tokenization tests (TDR-020). + +The load-bearing test is ``test_pattern_tiles_every_input``. If the pattern has a +gap, ``encode()`` silently drops characters โ€” and a round-trip test only notices +if it happens to use a dropped character. So coverage is asserted directly +against a spread of inputs rather than inferred. +""" + +import time + +import pytest + +from tokenizer import BPE, GPT2_PATTERN, pretokenize, tiles_exactly + +SAMPLES = [ + "hello world", + "The quick brown fox.", + "it's don't we've I'm you'll he'd", + "snake_case and CamelCase and kebab-case", + "numbers 123 4567 0.5 -3", + "punctuation!!! ??? ...---", + " leading and trailing ", + "tabs\tand\nnewlines\r\n", + "unicode: hรฉllo ฮฉ ไฝ ๅฅฝ ๆ—ฅๆœฌ่ชž", + "emoji ๐ŸŽ‰ mixed โ˜• with text", + "", + " ", + "\n\n\n", + "a", + "_", + "_leading_underscore", + "url https://example.com/a?b=c#d", + "math โˆ‘ โ‰ˆ โˆž ร— รท", + "quotes \"double\" 'single' `back`", + "mixed123abc456", +] + + +# --------------------------------------------------------------------------- +# Coverage -- the pattern must lose nothing +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("text", SAMPLES) +def test_pattern_tiles_every_input(text): + """Concatenating the pieces must reproduce the input exactly.""" + assert tiles_exactly(text), f"pattern dropped characters from {text!r}" + + +def test_pattern_tiles_a_large_real_document(): + from pathlib import Path + + text = Path("docs/TDD.md").read_text(encoding="utf-8") + assert tiles_exactly(text) + + +@pytest.mark.parametrize("codepoint", [0x5F, 0x7F, 0xA0, 0x2603, 0x1F600]) +def test_awkward_codepoints_are_not_dropped(codepoint): + """Underscore, DEL, NBSP, snowman, emoji -- each has bitten a naive pattern.""" + text = f"a{chr(codepoint)}b" + assert tiles_exactly(text) + + +# --------------------------------------------------------------------------- +# Splitting behaviour +# --------------------------------------------------------------------------- + + +def test_a_leading_space_attaches_to_its_word(): + """GPT-2 convention: " world" is one piece, so spacing is learned not guessed.""" + assert pretokenize("hello world") == ["hello", " world"] + + +def test_letters_digits_and_symbols_split_apart(): + assert pretokenize("abc123!!") == ["abc", "123", "!!"] + + +def test_contractions_stay_whole(): + assert "'s" in pretokenize("it's") + + +def test_pattern_is_reachable_from_the_tokenizer(): + assert BPE(vocab_size=256).pattern is GPT2_PATTERN + + +# --------------------------------------------------------------------------- +# Merges must not cross piece boundaries +# --------------------------------------------------------------------------- + + +def test_no_merge_spans_a_piece_boundary(): + """ "the cat" repeated cannot produce a token containing the space-joined pair. + + Before TDR-020 the merge loop ran over whole documents, so "e c" was a + learnable token. Restricting merges to pieces is what makes the encode cache + valid -- a cached piece would otherwise depend on its neighbours. + """ + tok = BPE(vocab_size=400) + tok.train(["the cat sat on the mat "] * 50) + + for idx in tok.decoder: + raw = tok.token_bytes(idx) + # A token may *start* with a space (GPT-2 style " the"), but must never + # contain an interior space -- that would be a cross-piece merge. + assert b" " not in raw[1:], f"token {raw!r} spans a piece boundary" + + +def test_pieces_encode_independently_of_context(): + """The cache is only sound if a piece's ids never depend on what surrounds it.""" + tok = BPE(vocab_size=400) + tok.train(["the cat sat on the mat "] * 50) + + alone = tok.encode(" cat") + in_context = tok.encode("the cat sat") + assert alone == in_context[len(tok.encode("the")) :][: len(alone)] + + +# --------------------------------------------------------------------------- +# Caching and complexity +# --------------------------------------------------------------------------- + + +def test_each_unique_piece_is_merged_only_once(): + """Repetition must cost a dict lookup, not a re-merge. + + Deterministic stand-in for the throughput claim: 200 copies of the same + sentence must invoke the merge routine only once per distinct piece. + """ + tok = BPE(vocab_size=400) + tok.train(["the cat sat on the mat "] * 20) + tok._piece_cache.clear() + + calls = [] + original = tok._encode_piece + tok._encode_piece = lambda piece: calls.append(piece) or original(piece) # type: ignore[method-assign] + + text = "the cat sat on the mat " * 200 + tok.encode(text) + + assert len(calls) == len(set(calls)), "a piece was merged more than once" + assert len(calls) < 10, f"expected a handful of distinct pieces, got {len(calls)}" + + +def test_encode_scales_far_better_than_quadratically(): + """Regression for the O(n^2) whole-document merge loop. + + The old implementation measured 4.0x time per 2x input. A 16x input increase + would therefore have cost ~256x. The bound here is deliberately loose -- it + only needs to fail for a quadratic implementation, not to pin a throughput. + """ + from pathlib import Path + + text = Path("docs/TDD.md").read_text(encoding="utf-8") * 4 + tok = BPE(vocab_size=1024) + tok.train([text[:100_000]]) + + def encode_time(n): + tok._piece_cache.clear() + chunk = text[:n] + start = time.perf_counter() + tok.encode(chunk) + return time.perf_counter() - start + + small = encode_time(8_192) + large = encode_time(131_072) # 16x the input + # Quadratic would be ~256x. The bound is loose on purpose: this runs on + # shared CI runners where a 2x timing margin flakes, and the test only needs + # to separate linear-ish from quadratic. + assert large < small * 64, f"16x input cost {large / small:.0f}x time -- looks quadratic" + + +def test_the_cache_does_not_change_the_answer(): + tok = BPE(vocab_size=400) + tok.train(["the cat sat on the mat "] * 20) + text = "the cat sat on the mat and the dog ran" + + warm = tok.encode(text) + tok._piece_cache.clear() + cold = tok.encode(text) + assert warm == cold + + +def test_training_clears_a_stale_cache(): + """Cached ids from a previous vocabulary would be silently wrong.""" + tok = BPE(vocab_size=400) + tok.encode("hello") + assert tok._piece_cache + tok.train(["hello world "] * 20) + assert not tok._piece_cache + + +# --------------------------------------------------------------------------- +# Special tokens under the new splitter +# --------------------------------------------------------------------------- + + +def test_special_tokens_survive_pretokenization(): + tok = BPE(vocab_size=400, special_tokens=["<|endoftext|>", "<|pad|>"]) + tok.train(["hello world "] * 20) + eot = tok.encoder["<|endoftext|>"] + + ids = tok.encode("hello<|endoftext|>world") + assert eot in ids + assert tok.decode(ids) == "hello<|endoftext|>world" + + +def test_adjacent_special_tokens_are_separate_ids(): + tok = BPE(vocab_size=400, special_tokens=["<|endoftext|>"]) + tok.train(["hello "] * 20) + eot = tok.encoder["<|endoftext|>"] + assert tok.encode("<|endoftext|><|endoftext|>") == [eot, eot] + + +def test_a_longer_special_token_wins_over_its_prefix(): + """Alternation is ordered longest-first, or "<|end|>" would shadow "<|endoftext|>".""" + tok = BPE(vocab_size=400, special_tokens=["<|end|>", "<|endoftext|>"]) + tok.train(["hello "] * 20) + assert tok.encode("<|endoftext|>") == [tok.encoder["<|endoftext|>"]] + + +def test_special_token_at_each_boundary(): + tok = BPE(vocab_size=400, special_tokens=["<|endoftext|>"]) + tok.train(["hello world "] * 20) + for text in ("<|endoftext|>hello", "hello<|endoftext|>", "<|endoftext|>"): + assert tok.decode(tok.encode(text)) == text diff --git a/tests/unit/test_quantization.py b/tests/unit/test_quantization.py new file mode 100644 index 0000000..025df8b --- /dev/null +++ b/tests/unit/test_quantization.py @@ -0,0 +1,291 @@ +"""Quantization tests (TDD - written before implementation). + +Exit criterion for build stage 6 (docs/BUILD_ORDER.md): perplexity within +tolerance of the fp32 reference. + +Quantization is unusually easy to fake. A `quantize()` that returns the model +untouched passes every shape and perplexity check trivially, so the tests here +assert the *memory actually shrank* by the expected factor -- measured in bytes +off the real tensors, not inferred from the dtype name. +""" + +import math + +import pytest +import torch +from torch import nn + +from model import CausalLM +from optimization import ( + QuantizedLinear, + dequantize_int4, + dequantize_int8, + model_nbytes, + quantize, + quantize_int4, + quantize_int8, +) + +VOCAB = 64 + + +def make_model(**overrides): + cfg = { + "model": { + "vocab_size": VOCAB, + "d_model": 32, + "n_layers": 2, + "n_heads": 4, + "n_kv_heads": 2, + "d_ff": 64, + "max_seq_len": 32, + **overrides, + }, + "attention": {"impl": "manual"}, + } + torch.manual_seed(0) + return CausalLM.from_config(cfg).eval() + + +# --------------------------------------------------------------------------- +# int8 tensor round trip +# --------------------------------------------------------------------------- + + +def test_int8_round_trip_stays_close_to_the_original(): + torch.manual_seed(0) + w = torch.randn(16, 32) + restored = dequantize_int8(*quantize_int8(w)) + # Symmetric per-channel int8 has 127 levels per side, so the worst-case + # error is half a step: max|w_row| / 254. + tolerance = w.abs().amax(dim=1, keepdim=True) / 254 + assert (restored - w).abs().le(tolerance * 1.01).all() + + +def test_int8_scales_are_per_output_channel_not_per_tensor(): + """One row 1000x larger than the rest must not destroy the small rows. + + A single per-tensor scale would quantise the small rows to nearly all + zeros; per-channel scales keep each row's own dynamic range. + """ + w = torch.randn(4, 32) + w[0] *= 1000.0 + q, scales = quantize_int8(w) + assert scales.shape == (4, 1) + + restored = dequantize_int8(q, scales) + small_rows_error = (restored[1:] - w[1:]).abs().max() + assert small_rows_error < w[1:].abs().max() * 0.02 + + +def test_int8_values_stay_in_range(): + q, _ = quantize_int8(torch.randn(8, 16) * 50) + assert q.dtype == torch.int8 + assert q.min() >= -127 and q.max() <= 127 + + +def test_int8_handles_an_all_zero_row_without_dividing_by_zero(): + w = torch.randn(4, 8) + w[2] = 0.0 + restored = dequantize_int8(*quantize_int8(w)) + assert torch.isfinite(restored).all() + torch.testing.assert_close(restored[2], torch.zeros(8)) + + +# --------------------------------------------------------------------------- +# int4 packing +# --------------------------------------------------------------------------- + + +def test_int4_packs_two_values_per_byte(): + """The memory claim is the point -- assert the byte count, not the dtype.""" + w = torch.randn(8, 32) + packed, scales, in_features = quantize_int4(w) + assert in_features == 32 + assert packed.dtype == torch.uint8 + assert packed.shape == (8, 16) + assert packed.numel() * packed.element_size() == w.numel() // 2 + + +def test_int4_round_trip_is_coarser_than_int8_but_bounded(): + torch.manual_seed(0) + w = torch.randn(16, 32) + err8 = (dequantize_int8(*quantize_int8(w)) - w).abs().mean() + err4 = (dequantize_int4(*quantize_int4(w)) - w).abs().mean() + assert err8 < err4, "int4 should be coarser than int8" + tolerance = w.abs().amax(dim=1, keepdim=True) / 14 + assert (dequantize_int4(*quantize_int4(w)) - w).abs().le(tolerance * 1.01).all() + + +def test_int4_pads_an_odd_number_of_columns(): + """Two values share a byte, so an odd width needs a pad column.""" + w = torch.randn(4, 7) + packed, scales, in_features = quantize_int4(w) + assert in_features == 7 + assert packed.shape == (4, 4) + assert dequantize_int4(packed, scales, in_features).shape == (4, 7) + + +# --------------------------------------------------------------------------- +# QuantizedLinear +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bits", [8, 4]) +def test_quantized_linear_tracks_the_float_layer(bits): + torch.manual_seed(0) + linear = nn.Linear(32, 16, bias=False) + q = QuantizedLinear.from_linear(linear, bits=bits) + + x = torch.randn(4, 32) + expected, actual = linear(x), q(x) + assert actual.shape == expected.shape + relative = (actual - expected).norm() / expected.norm() + assert relative < (0.02 if bits == 8 else 0.2), f"relative error {relative:.4f}" + + +def test_quantized_linear_preserves_the_bias_in_full_precision(): + """Bias is one value per output -- quantising it saves nothing and costs accuracy.""" + linear = nn.Linear(8, 4, bias=True) + q = QuantizedLinear.from_linear(linear, bits=8) + torch.testing.assert_close(q.bias, linear.bias) + + +def test_quantized_linear_stores_no_float_weight(): + q = QuantizedLinear.from_linear(nn.Linear(32, 16, bias=False), bits=8) + assert not any(p.dtype.is_floating_point and p.numel() > 16 for p in q.buffers()) + assert not list(q.parameters()), "quantised weights must not stay trainable" + + +# --------------------------------------------------------------------------- +# Whole-model quantization +# --------------------------------------------------------------------------- + + +def test_quantize_shrinks_the_model(): + model = make_model() + before = model_nbytes(model) + after = model_nbytes(quantize(make_model(), bits=8)) + assert after < before, f"{before} -> {after} bytes: no saving at all" + + +def test_int4_is_smaller_than_int8(): + eight = model_nbytes(quantize(make_model(), bits=8)) + four = model_nbytes(quantize(make_model(), bits=4)) + assert four < eight + + +def test_quantize_replaces_the_linear_layers(): + model = quantize(make_model(), bits=8) + quantised = [m for m in model.modules() if isinstance(m, QuantizedLinear)] + assert quantised, "no layer was actually replaced" + # every attention/FFN projection in a 2-layer model: 4 attn + 3 ffn + assert len(quantised) == 2 * 7 + + +def test_the_tied_lm_head_is_left_alone_by_default(): + """Quantising a tied head silently breaks the tie with the embedding. + + The head would become an int8 copy while the embedding stayed fp32, so the + two matrices -- which are supposed to be one tensor -- would diverge. + """ + model = quantize(make_model(tie_weights=True), bits=8) + assert isinstance(model.lm_head, nn.Linear) + assert model.lm_head.weight is model.embed_tokens.weight + + +def test_quantized_model_still_produces_valid_logits(): + model = quantize(make_model(), bits=8) + logits = model(torch.randint(0, VOCAB, (2, 8))) + assert logits.shape == (2, 8, VOCAB) + assert torch.isfinite(logits).all() + + +def test_quantized_model_still_generates_with_a_cache(): + from inference import generate_ids + + model = quantize(make_model(), bits=8) + out = generate_ids(model, torch.tensor([[1, 2, 3]]), max_new_tokens=5, temperature=0.0) + assert len(out) == 5 + + +def test_fp16_halves_the_model(): + model = make_model() + before = model_nbytes(model) + after = model_nbytes(quantize(make_model(), scheme="fp16")) + assert after == pytest.approx(before / 2, rel=0.01) + + +def test_unknown_scheme_and_bit_width_are_rejected(): + with pytest.raises(ValueError, match="scheme"): + quantize(make_model(), scheme="int3_magic") + with pytest.raises(ValueError, match="bits"): + quantize(make_model(), bits=7) + + +# --------------------------------------------------------------------------- +# The exit test: perplexity within tolerance +# --------------------------------------------------------------------------- + + +def _perplexity(model, ids): + with torch.no_grad(): + logits = model(ids[:, :-1]) + loss = nn.functional.cross_entropy(logits.reshape(-1, logits.shape[-1]), ids[:, 1:].reshape(-1)) + return math.exp(loss.item()) + + +@pytest.fixture(scope="module") +def trained(): + """A briefly-trained model, so quantization error has real structure to damage.""" + from training.trainer import Trainer + + cfg = { + "model": { + "vocab_size": VOCAB, + "d_model": 64, + "n_layers": 2, + "n_heads": 4, + "n_kv_heads": 2, + "d_ff": 128, + "max_seq_len": 64, + }, + "attention": {"impl": "manual"}, + "training": { + "steps": 150, + "lr": 3e-3, + "min_lr": 3e-4, + "warmup_steps": 10, + "grad_clip": 1.0, + "amp": False, + "weight_decay": 0.0, + }, + } + torch.manual_seed(0) + # A learnable pattern: token t+1 is always (t * 7 + 3) mod VOCAB. + seq = torch.tensor([[(i * 7 + 3) % VOCAB for i in range(33)]]).repeat(8, 1) + batch = {"input_ids": seq[:, :-1], "labels": seq[:, 1:]} + + class OneBatch: + def next_batch(self, device=None): + return batch + + model = CausalLM.from_config(cfg) + Trainer(model, OneBatch(), cfg, device=torch.device("cpu")).train(150) + return model.eval(), seq + + +def test_int8_perplexity_stays_within_tolerance(trained): + """The stage 6 exit criterion.""" + model, seq = trained + reference = _perplexity(model, seq) + quantised = _perplexity(quantize(model, bits=8), seq) + assert quantised < reference * 1.10, f"ppl {reference:.3f} -> {quantised:.3f} (>10% worse)" + + +def test_int4_degrades_more_than_int8(trained): + """A sanity check on the tradeoff -- fewer bits must cost something.""" + model, seq = trained + eight = _perplexity(quantize(model, bits=8), seq) + four = _perplexity(quantize(model, bits=4), seq) + assert eight <= four diff --git a/tests/unit/test_tokenizer.py b/tests/unit/test_tokenizer.py new file mode 100644 index 0000000..d0a1804 --- /dev/null +++ b/tests/unit/test_tokenizer.py @@ -0,0 +1,171 @@ +"""Tokenizer unit tests (TDD - written before implementation).""" + +import os + +import pytest + +from tokenizer.bpe import BPE + +# --------------------------------------------------------------------------- +# Encoding / decoding correctness +# --------------------------------------------------------------------------- + + +def test_encode_basic_words(): + tok = BPE(vocab_size=256) + # Base byte-level tokenizer: each byte is a separate token (ASCII chars) + encoded = tok.encode("hello") + assert encoded == [ord(c) for c in "hello"] + + +def test_decode_basic_words(): + tok = BPE(vocab_size=256) + assert tok.decode([104, 101, 108, 108, 111]) == "hello" + + +def test_round_trip_encode_decode(): + tok = BPE(vocab_size=300) + tok.train(["hello world", "this is a test sentence"]) + text = "hello world" + assert tok.decode(tok.encode(text)) == text + + +def test_round_trip_unicode(): + tok = BPE(vocab_size=300) + tok.train(["hello", "hรฉllo", "ไฝ ๅฅฝ", "emoji ๐ŸŽ‰"]) + for text in ["hello", "hรฉllo", "ไฝ ๅฅฝ", "emoji ๐ŸŽ‰"]: + assert tok.decode(tok.encode(text)) == text, f"failed for {text!r}" + + +def test_encode_no_unknown_tokens(): + """Byte-level BPE should never produce unknown tokens for any UTF-8 text.""" + tok = BPE(vocab_size=256) + for text in ["", "a", "ฮฉ", "ๆ—ฅๆœฌ่ชž", "emoji ๐Ÿ˜€", "mixed cafรฉ โ˜•"]: + tokens = tok.encode(text) + assert all(0 <= t < tok.vocab_size for t in tokens) + + +# --------------------------------------------------------------------------- +# Vocabulary / BPE merging +# --------------------------------------------------------------------------- + + +def test_training_grows_vocabulary(): + text = "the quick brown fox jumps over the lazy dog. " * 50 + before = BPE(vocab_size=256).vocab_size + tok = BPE(vocab_size=512) + tok.train([text]) + # merges should have created new tokens beyond the base 256 bytes + assert tok.vocab_size > before + assert tok.vocab_size <= 512 + + +def test_most_frequent_pair_merged_first(): + """'ab' appears more than other pairs, so it should become a token. + + Canonical representation: decoder values are GPT-2-style rendered-unicode + strings, NOT raw bytes. We therefore verify the merge via token_bytes(), + which recovers the original byte sequence backing each token id. + """ + text = "ab" * 100 + " xy " * 10 + tok = BPE(vocab_size=258) # +2 merges + tok.train([text]) + # 'ab' should be one of the merged tokens (compare original bytes). + merged_byte_seqs = {tok.token_bytes(i) for i in tok.decoder} + assert b"ab" in merged_byte_seqs + # And it must actually be used when encoding. + assert ( + tok.encode("ab") == [tok.encoder[tok._render_bytes(list(b"ab"))]] + or b"ab" in merged_byte_seqs + ) + + +def test_vocab_consistency(): + """encoder and decoder must be consistent inverses (for token bytes).""" + tok = BPE(vocab_size=300) + tok.train(["hello world", "the cat sat", "abcdefg"]) + assert len(tok.encoder) == len(tok.decoder) == tok.vocab_size + for token_bytes, idx in tok.encoder.items(): + assert tok.decoder[idx] == token_bytes + + +# --------------------------------------------------------------------------- +# Special tokens +# --------------------------------------------------------------------------- + + +def test_special_tokens_reserved(): + tok = BPE(vocab_size=300, special_tokens=["<|endoftext|>", "<|pad|>"]) + tok.train(["hello world"]) + assert "<|endoftext|>" in tok.encoder + assert "<|pad|>" in tok.encoder + # special tokens should occupy low ids + assert tok.encoder["<|endoftext|>"] == 0 + assert tok.encoder["<|pad|>"] == 1 + + +def test_encode_decode_special_token(): + tok = BPE(vocab_size=300, special_tokens=["<|endoftext|>"]) + tok.train(["hello"]) + eot = tok.encoder["<|endoftext|>"] + text = "hello" + result = tok.decode(tok.encode(text) + [eot]) + assert result == "hello<|endoftext|>" or result == "hello" + + +# --------------------------------------------------------------------------- +# Serialization +# --------------------------------------------------------------------------- + + +@pytest.fixture +def trained_tokenizer(tmp_path): + tok = BPE(vocab_size=300, special_tokens=["<|endoftext|>", "<|pad|>"]) + tok.train(["the quick brown fox jumps over the lazy dog"] * 20) + return tok + + +def test_save_and_load_round_trip(trained_tokenizer, tmp_path): + path = os.path.join(tmp_path, "vocab.json") + trained_tokenizer.save(path) + loaded = BPE(vocab_size=300, special_tokens=["<|endoftext|>", "<|pad|>"]) + loaded.load(path) + + assert loaded.vocab_size == trained_tokenizer.vocab_size + assert loaded.encoder == trained_tokenizer.encoder + assert loaded.decoder == trained_tokenizer.decoder + + # Encode/decode must be identical after reload + sample = "the quick brown fox jumps" + assert loaded.encode(sample) == trained_tokenizer.encode(sample) + assert loaded.decode(loaded.encode(sample)) == sample + + +def test_save_creates_file(trained_tokenizer, tmp_path): + path = tmp_path / "vocab.json" + trained_tokenizer.save(str(path)) + assert path.exists() + + +def test_save_is_silent(trained_tokenizer, tmp_path, capsys): + """A tokenizer saved every checkpoint must not narrate into the training log.""" + trained_tokenizer.save(str(tmp_path / "vocab.json")) + assert capsys.readouterr().out == "" + + +# --------------------------------------------------------------------------- +# Edge cases +# --------------------------------------------------------------------------- + + +def test_empty_text_encoding(): + tok = BPE(vocab_size=256) + assert tok.encode("") == [] + + +def test_encoding_deterministic(): + tok = BPE(vocab_size=300) + tok.train(["repeat repeat repeat"]) + a = tok.encode("repeat repeat repeat") + b = tok.encode("repeat repeat repeat") + assert a == b diff --git a/tests/unit/test_train_config.py b/tests/unit/test_train_config.py new file mode 100644 index 0000000..2fb1087 --- /dev/null +++ b/tests/unit/test_train_config.py @@ -0,0 +1,110 @@ +"""Config loading and CLI override tests. + +Overrides are the easiest place in the project to introduce a silent bug: a +mistyped value does not crash, it just trains with the wrong hyperparameter. +``test_scientific_notation_becomes_a_float`` exists because PyYAML implements +YAML 1.1, which does not accept "3e-4" as a float -- so the natural +implementation quietly hands the trainer a string learning rate. +""" + +import pytest +import yaml + +from training.train import apply_overrides, coerce, load_config, set_deep + +# --------------------------------------------------------------------------- +# Value coercion +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("12", 12), + ("-3", -3), + ("0.5", 0.5), + ("3.0e-4", 3.0e-4), + ("true", True), + ("false", False), + ("checkpoints", "checkpoints"), + ("sdpa", "sdpa"), + ("bfloat16", "bfloat16"), + ], +) +def test_coerce_matches_yaml_types(raw, expected): + value = coerce(raw) + assert value == expected + assert type(value) is type(expected) + + +@pytest.mark.parametrize("raw", ["3e-4", "1e-4", "1E5"]) +def test_scientific_notation_becomes_a_float(raw): + value = coerce(raw) + assert isinstance(value, float), f"{raw!r} must be a float, got {type(value).__name__}" + assert value == float(raw) + + +def test_override_type_matches_the_config_it_replaces(): + """An override must not change a key's type -- that is the whole risk.""" + cfg = yaml.safe_load("training:\n lr: 3.0e-4\n amp: true\n steps: 2000\n") + apply_overrides(cfg, ["training.lr=1e-4", "training.amp=false", "training.steps=10"]) + assert isinstance(cfg["training"]["lr"], float) + assert isinstance(cfg["training"]["amp"], bool) + assert isinstance(cfg["training"]["steps"], int) + assert cfg["training"] == {"lr": 1e-4, "amp": False, "steps": 10} + + +# --------------------------------------------------------------------------- +# Dotted-key assignment +# --------------------------------------------------------------------------- + + +def test_set_deep_creates_missing_levels(): + cfg: dict = {} + set_deep(cfg, "a.b.c", 1) + assert cfg == {"a": {"b": {"c": 1}}} + + +def test_set_deep_overwrites_without_clobbering_siblings(): + cfg = {"model": {"d_model": 512, "n_layers": 8}} + set_deep(cfg, "model.n_layers", 12) + assert cfg == {"model": {"d_model": 512, "n_layers": 12}} + + +def test_override_without_equals_is_rejected(): + with pytest.raises(ValueError, match="dotted.key=value"): + apply_overrides({}, ["training.lr"]) + + +def test_empty_override_list_is_a_noop(): + cfg = {"seed": 42} + assert apply_overrides(cfg, []) == {"seed": 42} + + +# --------------------------------------------------------------------------- +# Config file +# --------------------------------------------------------------------------- + + +def test_load_config_reads_yaml(tmp_path): + path = tmp_path / "c.yaml" + path.write_text("seed: 7\nmodel:\n d_model: 64\n", encoding="utf-8") + assert load_config(str(path)) == {"seed": 7, "model": {"d_model": 64}} + + +def test_load_config_tolerates_an_empty_file(tmp_path): + path = tmp_path / "empty.yaml" + path.write_text("", encoding="utf-8") + assert load_config(str(path)) == {} + + +def test_shipped_default_config_has_every_key_the_code_reads(): + """Guards against the config and the code drifting apart.""" + cfg = load_config("configs/default.yaml") + assert cfg["model"]["n_heads"] % cfg["model"]["n_kv_heads"] == 0 + assert cfg["model"]["d_model"] % cfg["model"]["n_heads"] == 0 + assert cfg["dataset"]["seq_len"] <= cfg["model"]["max_seq_len"] + for key in ("d_model", "n_layers", "n_heads", "n_kv_heads", "d_ff", "vocab_size"): + assert key in cfg["model"], f"configs/default.yaml is missing model.{key}" + for key in ("steps", "lr", "warmup_steps", "out_dir", "amp_dtype"): + assert key in cfg["training"], f"configs/default.yaml is missing training.{key}" diff --git a/tests/unit/test_trainer.py b/tests/unit/test_trainer.py new file mode 100644 index 0000000..7d1a863 --- /dev/null +++ b/tests/unit/test_trainer.py @@ -0,0 +1,312 @@ +"""Trainer and checkpoint tests (TDD - written before implementation). + +Exit criterion for build stage 4 (docs/BUILD_ORDER.md): loss falls when +overfitting a single batch. That test is the one that proves the whole stack +composes -- a broken causal mask, a mis-shifted label, a dead gradient path or +a scheduler stuck at zero all show up as a flat loss curve here. +""" + +import math + +import pytest +import torch + +from model import CausalLM +from training.checkpoint import load_checkpoint, save_checkpoint +from training.trainer import Trainer + +VOCAB, SEQ, BATCH = 32, 8, 4 + + +class FixedEngine: + """Replays a fixed batch list so tests are deterministic. + + Satisfies the same duck-typed `next_batch` interface as DataEngine, which + is all Trainer depends on. + """ + + def __init__(self, batches): + self.batches = batches + self.i = 0 + self.served = 0 + + def next_batch(self, device=None): + batch = self.batches[self.i % len(self.batches)] + self.i += 1 + self.served += 1 + if device is None: + return batch + return {k: v.to(device) for k, v in batch.items()} + + +def make_batch(seed=0, batch=BATCH, seq=SEQ): + torch.manual_seed(seed) + ids = torch.randint(0, VOCAB, (batch, seq + 1)) + return {"input_ids": ids[:, :-1], "labels": ids[:, 1:]} + + +def make_cfg(**overrides): + training = { + "steps": 100, + "lr": 1e-3, + "min_lr": 1e-4, + "warmup_steps": 0, + "weight_decay": 0.0, + "grad_clip": 1.0, + "grad_accum_steps": 1, + "amp": False, + "log_every": 10, + } + training.update(overrides) + return { + "model": { + "vocab_size": VOCAB, + "d_model": 32, + "n_layers": 2, + "n_heads": 4, + "n_kv_heads": 2, + "d_ff": 64, + "max_seq_len": 32, + }, + "attention": {"impl": "manual"}, + "training": training, + "device": "cpu", + } + + +def make_trainer(batches=None, **cfg_overrides): + cfg = make_cfg(**cfg_overrides) + torch.manual_seed(0) + model = CausalLM.from_config(cfg) + engine = FixedEngine(batches or [make_batch()]) + return Trainer(model, engine, cfg, device=torch.device("cpu")) + + +# --------------------------------------------------------------------------- +# A single step +# --------------------------------------------------------------------------- + + +def test_step_reports_the_expected_metrics(): + metrics = make_trainer().step() + assert set(metrics) == {"step", "loss", "lr", "grad_norm", "tokens", "tokens_per_sec"} + assert metrics["step"] == 1 + assert metrics["tokens"] == BATCH * SEQ + assert metrics["tokens_per_sec"] > 0 + + +def test_first_loss_is_near_uniform_entropy(): + loss = make_trainer().step()["loss"] + assert abs(loss - math.log(VOCAB)) < 0.3 + + +def test_step_actually_updates_the_weights(): + trainer = make_trainer() + before = trainer.model.embed_tokens.weight.detach().clone() + trainer.step() + assert not torch.equal(before, trainer.model.embed_tokens.weight) + + +def test_gradients_are_cleared_between_steps(): + """A missing zero_grad accumulates forever and silently inflates the step.""" + trainer = make_trainer() + trainer.step() + grads = [p.grad for p in trainer.model.parameters() if p.grad is not None] + assert all(g.count_nonzero() == 0 for g in grads) or not grads + + +# --------------------------------------------------------------------------- +# The exit test +# --------------------------------------------------------------------------- + + +def test_overfits_a_single_batch(): + """Repeat one batch until the model memorises it. Loss must collapse. + + This is the stage-4 exit criterion: a flat curve here means something + upstream (mask, label shift, gradient path, schedule) is broken. + """ + trainer = make_trainer(steps=200, lr=3e-3, warmup_steps=10) + history = trainer.train(200) + + first, last = history[0]["loss"], history[-1]["loss"] + assert first > 3.0, f"expected to start near ln(32)={math.log(32):.2f}, got {first:.2f}" + assert last < 0.2, f"failed to memorise one batch: {first:.2f} -> {last:.2f}" + + +def test_loss_is_monotonic_enough_to_be_real(): + """Averaged over windows the curve must fall, not just end low by luck.""" + history = make_trainer(steps=120, lr=3e-3, warmup_steps=5).train(120) + losses = [h["loss"] for h in history] + early = sum(losses[:20]) / 20 + late = sum(losses[-20:]) / 20 + assert late < early / 2 + + +# --------------------------------------------------------------------------- +# Gradient accumulation +# --------------------------------------------------------------------------- + + +def test_accumulating_k_micro_batches_matches_one_big_batch(): + """Two micro-batches of 4 must give the same gradient as one batch of 8. + + If the per-micro-batch loss is not divided by grad_accum_steps, the summed + gradient is K times too large -- which looks like training with a K-times + larger learning rate and is easy to mistake for a tuning problem. + """ + a, b = make_batch(1), make_batch(2) + big = {k: torch.cat([a[k], b[k]], dim=0) for k in a} + + accum = make_trainer([a, b], grad_accum_steps=2, grad_clip=0.0) + single = make_trainer([big], grad_accum_steps=1, grad_clip=0.0) + + accum.model.load_state_dict(single.model.state_dict()) + + for trainer in (accum, single): + trainer.model.train() + for _ in range(trainer.grad_accum_steps): + batch = trainer.engine.next_batch(trainer.device) + (trainer.loss_on(batch) / trainer.grad_accum_steps).backward() + + for (name, p_accum), (_, p_single) in zip( + accum.model.named_parameters(), single.model.named_parameters(), strict=True + ): + torch.testing.assert_close( + p_accum.grad, p_single.grad, rtol=1e-4, atol=1e-6, msg=f"gradient differs: {name}" + ) + + +def test_accumulation_consumes_one_batch_per_micro_step(): + trainer = make_trainer([make_batch(1), make_batch(2), make_batch(3)], grad_accum_steps=3) + trainer.step() + assert trainer.engine.served == 3 + + +# --------------------------------------------------------------------------- +# Gradient clipping and schedule +# --------------------------------------------------------------------------- + + +def test_gradient_clipping_bounds_the_update(): + """With a tiny threshold the post-clip norm must sit at the threshold.""" + trainer = make_trainer(grad_clip=1e-4, lr=1e-2) + trainer.model.train() + batch = trainer.engine.next_batch(trainer.device) + trainer.loss_on(batch).backward() + torch.nn.utils.clip_grad_norm_(trainer.model.parameters(), 1e-4) + + total = torch.norm( + torch.stack([p.grad.norm() for p in trainer.model.parameters() if p.grad is not None]) + ) + assert total <= 1e-4 * 1.01 + + +def test_reported_grad_norm_is_the_pre_clip_value(): + """Reporting the post-clip norm would pin it at the threshold and hide spikes.""" + metrics = make_trainer(grad_clip=1e-6).step() + assert metrics["grad_norm"] > 1e-6 + + +def test_learning_rate_follows_the_warmup_then_decay_schedule(): + trainer = make_trainer(steps=50, lr=1e-2, min_lr=1e-3, warmup_steps=10) + lrs = [trainer.step()["lr"] for _ in range(50)] + assert lrs[0] == pytest.approx(1e-3) # first warmup step + assert lrs[9] == pytest.approx(1e-2) # peak at end of warmup + assert lrs[-1] < lrs[9] # decaying afterwards + + +def test_amp_is_disabled_on_cpu(): + """fp16 on CPU is emulated and slower than fp32 -- don't pretend otherwise.""" + assert make_trainer(amp=True).amp_dtype is None + + +# --------------------------------------------------------------------------- +# Checkpointing +# --------------------------------------------------------------------------- + + +def test_checkpoint_round_trips_the_weights(tmp_path): + trainer = make_trainer() + trainer.train(3) + path = tmp_path / "ckpt.pt" + trainer.save(str(path)) + + fresh = make_trainer() + assert not torch.equal(fresh.model.embed_tokens.weight, trainer.model.embed_tokens.weight) + fresh.load(str(path)) + for (name, a), (_, b) in zip( + fresh.model.named_parameters(), trainer.model.named_parameters(), strict=True + ): + torch.testing.assert_close(a, b, msg=f"parameter differs after resume: {name}") + + +def test_resume_restores_the_step_counter(tmp_path): + trainer = make_trainer() + trainer.train(5) + path = tmp_path / "ckpt.pt" + trainer.save(str(path)) + + fresh = make_trainer() + fresh.load(str(path)) + assert fresh.step_count == 5 + + +def test_resume_restores_optimizer_moments(tmp_path): + """Dropping Adam's moments causes a visible loss spike on the next step.""" + trainer = make_trainer() + trainer.train(5) + path = tmp_path / "ckpt.pt" + trainer.save(str(path)) + + fresh = make_trainer() + fresh.load(str(path)) + + original = [ + trainer.optimizer.state[p]["exp_avg"] + for g in trainer.optimizer.param_groups + for p in g["params"] + if p in trainer.optimizer.state + ] + restored = [ + fresh.optimizer.state[p]["exp_avg"] + for g in fresh.optimizer.param_groups + for p in g["params"] + if p in fresh.optimizer.state + ] + assert original and len(original) == len(restored) + for a, b in zip(original, restored, strict=True): + torch.testing.assert_close(a, b) + + +def test_resumed_run_continues_identically(tmp_path): + """The step after a resume must match the step that would have happened.""" + uninterrupted = make_trainer() + uninterrupted.train(4) + path = tmp_path / "ckpt.pt" + uninterrupted.save(str(path)) + expected = uninterrupted.step()["loss"] + + resumed = make_trainer() + resumed.load(str(path)) + assert resumed.step()["loss"] == pytest.approx(expected, rel=1e-5) + + +def test_loading_into_a_different_architecture_is_rejected(tmp_path): + """A partial load yields a half-random model that still trains.""" + trainer = make_trainer() + path = tmp_path / "ckpt.pt" + trainer.save(str(path)) + + cfg = make_cfg() + cfg["model"]["n_layers"] = 4 + mismatched = CausalLM.from_config(cfg) + with pytest.raises(RuntimeError): + load_checkpoint(str(path), model=mismatched) + + +def test_save_checkpoint_creates_missing_directories(tmp_path): + trainer = make_trainer() + path = tmp_path / "nested" / "deeper" / "ckpt.pt" + save_checkpoint(path, model=trainer.model, step=0) + assert path.exists() diff --git a/tokenizer/__init__.py b/tokenizer/__init__.py new file mode 100644 index 0000000..78cc3eb --- /dev/null +++ b/tokenizer/__init__.py @@ -0,0 +1,6 @@ +"""Byte-level BPE tokenizer module.""" + +from tokenizer.bpe import BPE +from tokenizer.pretokenize import GPT2_PATTERN, pretokenize, tiles_exactly + +__all__ = ["BPE", "GPT2_PATTERN", "pretokenize", "tiles_exactly"] diff --git a/tokenizer/bpe.py b/tokenizer/bpe.py new file mode 100644 index 0000000..faa14bb --- /dev/null +++ b/tokenizer/bpe.py @@ -0,0 +1,346 @@ +"""Byte-level Byte Pair Encoding (BPE) tokenizer built from scratch. + +Design notes (see docs/TDD.md ยง1, TDR-002, TDR-020): + +- **Byte-level.** Text is encoded to UTF-8 bytes, so any text โ€” emoji, rare + scripts โ€” is tokenizable with no "unknown" token. +- **Vocabulary.** Ids 0..k-1 are reserved special tokens; then the 256 base byte + tokens; then learned merges. +- **Pre-tokenized.** Text is split into pieces (words, digit runs, symbol runs) + before merging, and merges never cross a piece boundary (TDR-020). This is + what makes both training and encoding tractable โ€” see the complexity note + below โ€” and it stops the vocabulary being spent on cross-word artefacts. + +Why pre-tokenization is not optional (measured, not assumed): + +Merging is quadratic in the length of the sequence it runs over. Run over a whole +document, `encode()` measured 4.0x slower per 2x input โ€” 50 KB/s at 1 KB falling +to 2.9 KB/s at 16 KB โ€” which makes a 500 MB corpus unusable. Restricted to +pieces of ~1-20 bytes, with identical pieces cached, the same corpus encodes at +~11 MB/s. That is the difference between a two-day preprocessing step and a +40-second one. + +Training uses the same idea from the other direction: pieces are counted once +into a frequency table, and each merge only revisits the pieces that actually +contained the merged pair, rather than rescanning the corpus. + +API contract (see docs/BUILD_ORDER.md "Contracts"): + BPE(vocab_size, special_tokens) + .train(texts) + .encode(text) -> list[int] + .decode(ids) -> str + .save(path) / .load(path) + .vocab_size property +""" + +from __future__ import annotations + +import json +import re +from collections import Counter +from collections.abc import Iterable + +from tokenizer.pretokenize import GPT2_PATTERN, pretokenize + +SAVE_FORMAT_VERSION = 2 + + +class BPE: + """Byte-level Byte Pair Encoding tokenizer built from scratch.""" + + def __init__( + self, + vocab_size: int = 4096, + special_tokens: list[str] | None = None, + ) -> None: + if vocab_size < 256: + raise ValueError("vocab_size must be >= 256 (bytes)") + + self.vocab_size = vocab_size + self.special_tokens = list(special_tokens) if special_tokens else [] + + # Predefined byte<->unicode rendering table (GPT-2 style). + self.bs, self.sb = self._bytes_to_unicode_map() + + # encoder: token_key (str) -> id + # decoder: id -> token_key (str) + # For bytes, token_key is the rendered unicode string. + # For special tokens, token_key is the literal special string itself. + self.encoder: dict[str, int] = {} + self.decoder: dict[int, str] = {} + + # Special tokens first, occupying the lowest ids (atomic, never merged). + for i, sp in enumerate(self.special_tokens): + self.encoder[sp] = i + self.decoder[i] = sp + self.num_special = len(self.special_tokens) + + # Add the 256 base byte tokens. + next_id = self.num_special + for b in range(256): + rendered = self.bs[b] + self.encoder[rendered] = next_id + self.decoder[next_id] = rendered + next_id += 1 + self._base_size = next_id + + # Byte decomposition for every token id (token_id -> list[byte value]). + self._byte_seqs: dict[int, list[int]] = {self.num_special + b: [b] for b in range(256)} + + # BPE merges: list of ((a, b), merged_bytes, merged_key), rank-ordered. + self.merges: list[tuple] = [] + self.merge_rank: dict[tuple, int] = {} + # (a, b) -> merged id. Keeping the id here means the encode hot loop + # never re-renders a token key or hashes a long string. + self.merge_to_id: dict[tuple, int] = {} + + self._special_keys = set(self.special_tokens) + self._special_re: re.Pattern[str] | None = None + self._rebuild_special_re() + # piece -> ids. Natural text repeats pieces heavily, so after warmup most + # of a corpus costs one dict lookup. + self._piece_cache: dict[str, list[int]] = {} + + # ------------------------------------------------------------------ + # bytes <-> unicode rendering (GPT-2 style) + # ------------------------------------------------------------------ + @staticmethod + def _bytes_to_unicode_map(): + bs = ( + list(range(ord("!"), ord("~") + 1)) + + list(range(ord("ยก"), ord("ยฌ") + 1)) + + list(range(ord("ยฎ"), ord("รฟ") + 1)) + ) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8 + n) + n += 1 + b_to_u = {b: chr(c) for b, c in zip(bs, cs, strict=True)} + u_to_b = {c: b for b, c in b_to_u.items()} + return b_to_u, u_to_b + + def _render_bytes(self, byte_tokens: list[int]) -> str: + """Render a list of byte values (0-255) to the unicode token key string.""" + return "".join(self.bs[b] for b in byte_tokens) + + def _rebuild_special_re(self) -> None: + """Alternation over special tokens, longest first so prefixes lose.""" + if not self.special_tokens: + self._special_re = None + return + ordered = sorted(self.special_tokens, key=len, reverse=True) + self._special_re = re.compile("(" + "|".join(re.escape(s) for s in ordered) + ")") + + # ------------------------------------------------------------------ + # Training + # ------------------------------------------------------------------ + def train(self, texts: Iterable[str]) -> None: + """Learn BPE merges from a corpus. + + Counts pre-tokenized pieces into a frequency table, then repeatedly + merges the most frequent adjacent pair. An index from pair -> pieces + containing it means each merge only rewrites the pieces it affects, + instead of rescanning the whole corpus โ€” the difference between + O(merges x corpus) and O(merges x affected pieces). + """ + freqs: Counter[str] = Counter() + for text in texts: + freqs.update(pretokenize(text)) + + words: list[list[int]] = [ + [self.num_special + b for b in piece.encode("utf-8")] for piece in freqs + ] + weights: list[int] = list(freqs.values()) + + pair_counts: dict[tuple, int] = {} + pair_where: dict[tuple, set[int]] = {} + + def distinct_pairs(word: list[int]) -> set[tuple]: + return {(word[j], word[j + 1]) for j in range(len(word) - 1)} + + def add(index: int) -> None: + word, weight = words[index], weights[index] + for j in range(len(word) - 1): + pair = (word[j], word[j + 1]) + pair_counts[pair] = pair_counts.get(pair, 0) + weight + for pair in distinct_pairs(word): + pair_where.setdefault(pair, set()).add(index) + + def remove(index: int) -> None: + word, weight = words[index], weights[index] + for j in range(len(word) - 1): + pair = (word[j], word[j + 1]) + pair_counts[pair] = pair_counts.get(pair, 0) - weight + # Discard once per *distinct* pair: a word holding the same pair + # twice must not be dropped from the index on the first decrement. + for pair in distinct_pairs(word): + holders = pair_where.get(pair) + if holders is not None: + holders.discard(index) + if not holders: + pair_where.pop(pair, None) + if pair_counts.get(pair, 0) <= 0: + pair_counts.pop(pair, None) + + for index in range(len(words)): + add(index) + + while len(self.encoder) < self.vocab_size and pair_counts: + # Most frequent pair; tie-break deterministically by tuple value. + best = max(pair_counts, key=lambda p: (pair_counts[p], p)) + affected = list(pair_where.get(best, ())) + if not affected: + break + + a, b = best + merged_bytes = self._byte_seqs[a] + self._byte_seqs[b] + merged_key = self._render_bytes(merged_bytes) + merged_id = len(self.encoder) + + self.encoder[merged_key] = merged_id + self.decoder[merged_id] = merged_key + self._byte_seqs[merged_id] = merged_bytes + self.merges.append((best, merged_bytes, merged_key)) + self.merge_rank[best] = len(self.merges) - 1 + self.merge_to_id[best] = merged_id + + for index in affected: + remove(index) + for index in affected: + words[index] = self._apply_merge(words[index], best, merged_id) + for index in affected: + add(index) + + self.vocab_size = len(self.encoder) + self._piece_cache.clear() + + @staticmethod + def _apply_merge(word: list[int], pair: tuple, merged_id: int) -> list[int]: + """Replace every occurrence of ``pair`` in ``word`` with ``merged_id``.""" + out: list[int] = [] + i = 0 + n = len(word) + while i < n: + if i + 1 < n and (word[i], word[i + 1]) == pair: + out.append(merged_id) + i += 2 + else: + out.append(word[i]) + i += 1 + return out + + # ------------------------------------------------------------------ + # Encoding + # ------------------------------------------------------------------ + def _encode_piece(self, piece: str) -> list[int]: + """Merge within one pre-tokenized piece, applying all occurrences per pass.""" + ids = [self.num_special + b for b in piece.encode("utf-8")] + while len(ids) >= 2: + best_pair = None + best_rank = len(self.merges) + for i in range(len(ids) - 1): + rank = self.merge_rank.get((ids[i], ids[i + 1])) + if rank is not None and rank < best_rank: + best_rank = rank + best_pair = (ids[i], ids[i + 1]) + if best_pair is None: + break + ids = self._apply_merge(ids, best_pair, self.merge_to_id[best_pair]) + return ids + + def encode(self, text: str) -> list[int]: + """Encode a UTF-8 string to a list of token ids.""" + if not text: + return [] + + spans = self._special_re.split(text) if self._special_re is not None else [text] + + tokens: list[int] = [] + for span in spans: + if not span: + continue + if span in self._special_keys: + tokens.append(self.encoder[span]) + continue + for piece in pretokenize(span): + cached = self._piece_cache.get(piece) + if cached is None: + cached = self._piece_cache[piece] = self._encode_piece(piece) + tokens.extend(cached) + return tokens + + # ------------------------------------------------------------------ + # Introspection + # ------------------------------------------------------------------ + def token_bytes(self, idx: int) -> bytes: + """Return the raw byte sequence backing a token id.""" + key = self.decoder[idx] + if key in self._special_keys: + return key.encode("utf-8") + return bytes(self.sb[u] for u in key) + + @property + def pattern(self) -> re.Pattern[str]: + return GPT2_PATTERN + + # ------------------------------------------------------------------ + # Decoding + # ------------------------------------------------------------------ + def decode(self, ids: list[int]) -> str: + """Decode a list of token ids back to a UTF-8 string.""" + byte_parts: list[int] = [] + for i in ids: + key = self.decoder[i] + if key in self._special_keys: + byte_parts.extend(key.encode("utf-8")) + else: + for u in key: + byte_parts.append(self.sb[u]) + return bytes(byte_parts).decode("utf-8", errors="replace") + + # ------------------------------------------------------------------ + # Serialization + # ------------------------------------------------------------------ + def save(self, path: str) -> None: + """Serialize vocabulary + merges to a JSON file.""" + data = { + "version": SAVE_FORMAT_VERSION, + "vocab_size": self.vocab_size, + "special_tokens": self.special_tokens, + "encoder": dict(self.encoder), + "merges": [ + {"a": a, "b": b, "bytes": mb, "key": key} for (a, b), mb, key in self.merges + ], + } + with open(path, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, sort_keys=True) + + def load(self, path: str) -> None: + """Load vocabulary + merges from a JSON file.""" + with open(path, encoding="utf-8") as f: + data = json.load(f) + self.vocab_size = data["vocab_size"] + self.special_tokens = list(data["special_tokens"]) + self.encoder = {k: int(v) for k, v in data["encoder"].items()} + self.decoder = {int(v): k for k, v in data["encoder"].items()} + self.bs, self.sb = self._bytes_to_unicode_map() + self.num_special = len(self.special_tokens) + self._special_keys = set(self.special_tokens) + self._rebuild_special_re() + self._piece_cache.clear() + + self._byte_seqs = {self.num_special + b: [b] for b in range(256)} + + self.merges = [] + self.merge_rank = {} + self.merge_to_id = {} + for rank, m in enumerate(data["merges"]): + a, b, merged_bytes, merged_key = m["a"], m["b"], m["bytes"], m["key"] + merged_id = self.encoder[merged_key] + self.merges.append(((a, b), merged_bytes, merged_key)) + self._byte_seqs[merged_id] = merged_bytes + self.merge_rank[(a, b)] = rank + self.merge_to_id[(a, b)] = merged_id diff --git a/tokenizer/pretokenize.py b/tokenizer/pretokenize.py new file mode 100644 index 0000000..18ea15c --- /dev/null +++ b/tokenizer/pretokenize.py @@ -0,0 +1,46 @@ +"""Pre-tokenization: split text into pieces before BPE merges (TDR-020). + +BPE merges are applied *within* a piece and never across piece boundaries. That +serves two purposes: + +1. **Quality.** It stops the tokenizer learning merges that straddle word + boundaries (`"the c"`), which waste vocabulary on artefacts of word order. +2. **Speed.** Merging is quadratic in the length of the sequence it runs over. + Applied to a whole document that is fatal; applied to pieces of ~1-20 bytes + the quadratic term is a small constant, and identical pieces can be cached. + +The pattern is GPT-2's, with one deviation: GPT-2 uses `\\p{L}` / `\\p{N}`, which +the standard library's `re` does not support (it needs the third-party `regex` +module). The unicode-aware equivalents here are `[^\\W\\d]` for letters and `\\d` +for digits, which keeps the dependency list honest at the cost of grouping +underscore with letters. + +**The pattern must tile the input exactly** โ€” the concatenation of all matches +has to reproduce the input character for character. A pattern with a gap makes +`encode()` silently drop text, which round-trip tests only catch if they happen +to use the dropped character. `test_pattern_tiles_every_input` asserts it +directly instead. +""" + +from __future__ import annotations + +import re + +# Order matters: contractions first, then letter/digit/symbol runs each allowed +# one leading space, then whitespace. The `\s+(?!\S)` branch keeps a trailing +# run of whitespace together while letting a single space attach to the next +# word rather than becoming its own token. +GPT2_PATTERN = re.compile( + r"""'(?:s|t|re|ve|m|ll|d)| ?[^\W\d]+| ?\d+| ?[^\s\w]+|\s+(?!\S)|\s+""", + re.UNICODE, +) + + +def pretokenize(text: str, pattern: re.Pattern[str] = GPT2_PATTERN) -> list[str]: + """Split ``text`` into pieces that BPE may merge within, but not across.""" + return pattern.findall(text) + + +def tiles_exactly(text: str, pattern: re.Pattern[str] = GPT2_PATTERN) -> bool: + """True if the pattern reproduces ``text`` exactly โ€” no dropped characters.""" + return "".join(pattern.findall(text)) == text diff --git a/training/__init__.py b/training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/training/checkpoint.py b/training/checkpoint.py new file mode 100644 index 0000000..c4e7a0c --- /dev/null +++ b/training/checkpoint.py @@ -0,0 +1,85 @@ +"""Checkpoint save/load. + +What is saved: model weights, optimizer state, step counter, the config the run +was launched with, and RNG state. + +**What is deliberately not saved: the data loader's position.** Resuming +therefore restarts the corpus from the beginning of a shuffled epoch rather +than the exact window the run died on. For a shuffled multi-epoch run that is +harmless; for a single-pass run over a large corpus it means some windows are +seen twice and others not at all. Recording it here rather than implying +resumption is exact. + +Optimizer state matters more than it looks. Dropping it resets Adam's first and +second moments to zero, so the first step after a resume takes a full +bias-corrected stride in whatever direction the current gradient points โ€” +usually a visible loss spike. ``test_resume_restores_optimizer_moments`` pins it. +""" + +from __future__ import annotations + +import os +import random +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch import nn + + +def save_checkpoint( + path: str | os.PathLike[str], + *, + model: nn.Module, + optimizer: torch.optim.Optimizer | None = None, + step: int = 0, + cfg: dict[str, Any] | None = None, + extra: dict[str, Any] | None = None, +) -> Path: + """Write a checkpoint. Returns the path written.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + payload: dict[str, Any] = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict() if optimizer is not None else None, + "step": step, + "cfg": cfg, + "rng": { + "torch": torch.get_rng_state(), + "python": random.getstate(), + "numpy": np.random.get_state(), + }, + "extra": extra or {}, + } + torch.save(payload, path) + return path + + +def load_checkpoint( + path: str | os.PathLike[str], + *, + model: nn.Module, + optimizer: torch.optim.Optimizer | None = None, + map_location: str | torch.device | None = None, + restore_rng: bool = True, +) -> dict[str, Any]: + """Restore into ``model`` (and ``optimizer``). Returns the payload.""" + # weights_only=False: the payload carries RNG state and the config dict, not + # just tensors. Only ever load checkpoints this project wrote. + payload: dict[str, Any] = torch.load(path, map_location=map_location, weights_only=False) + + # strict=True on purpose: a silently partial load produces a model that is + # half-trained and half-random, which trains without ever looking broken. + model.load_state_dict(payload["model"], strict=True) + + if optimizer is not None and payload.get("optimizer") is not None: + optimizer.load_state_dict(payload["optimizer"]) + + if restore_rng and payload.get("rng"): + rng = payload["rng"] + torch.set_rng_state(rng["torch"]) + random.setstate(rng["python"]) + np.random.set_state(rng["numpy"]) + + return payload diff --git a/training/device.py b/training/device.py new file mode 100644 index 0000000..d524ffc --- /dev/null +++ b/training/device.py @@ -0,0 +1,45 @@ +"""Device selection (TDR-019). + +Preference order is ``mps -> cuda -> cpu``: Apple Silicon is the hardware this +project is developed and measured on, CUDA is supported but untested here, and +CPU is the always-correct fallback. + +MPS is not a drop-in CUDA. Two differences that leak into the layers above: + +- There is no ``torch.cuda.max_memory_allocated`` equivalent, so peak-memory + metrics have to come from ``tracemalloc``/RSS instead and are not directly + comparable to CUDA numbers. +- ``GradScaler`` is CUDA-only, and bf16 does not need loss scaling anyway. AMP + on MPS therefore means ``autocast(bfloat16)`` with no scaler. +""" + +from __future__ import annotations + +import torch + + +def resolve_device(preferred: str | None = None) -> torch.device: + """Return the device to run on. ``preferred`` overrides auto-selection.""" + if preferred and preferred != "auto": + return torch.device(preferred) + if torch.backends.mps.is_available(): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +def autocast_dtype(device: torch.device, requested: str = "bfloat16") -> torch.dtype | None: + """AMP dtype for ``device``, or ``None`` when mixed precision should be off. + + fp16 on CPU is emulated and slower than fp32, so CPU always returns ``None`` + rather than pretending AMP helps. + """ + if device.type == "cpu": + return None + return {"bfloat16": torch.bfloat16, "float16": torch.float16}.get(requested) + + +def supports_grad_scaler(device: torch.device) -> bool: + """``GradScaler`` is a CUDA fp16 mechanism; bf16 and MPS do not use it.""" + return device.type == "cuda" diff --git a/training/optimizer.py b/training/optimizer.py new file mode 100644 index 0000000..6807762 --- /dev/null +++ b/training/optimizer.py @@ -0,0 +1,124 @@ +"""Hand-written AdamW (TDR-009). + +Implemented directly rather than imported, per the project's no-external- +implementations rule. It is written to be *bit-comparable* with +``torch.optim.AdamW`` given the same hyperparameters โ€” that equivalence is the +only thing that makes "custom optimizer" a defensible claim, and +``test_matches_torch_adamw_step_for_step`` enforces it. + +The update, in order: + + p <- p * (1 - lr * weight_decay) # decoupled, applied to the weight + m <- b1*m + (1-b1)*g + v <- b2*v + (1-b2)*gยฒ + p <- p - (lr / (1-b1^t)) * m / (sqrt(v)/sqrt(1-b2^t) + eps) + +The decay is applied to the *parameter*, not folded into the gradient. Folding +it in (plain L2) means Adam's per-parameter normalisation rescales the decay +too, so the effective decay ends up depending on the gradient magnitude โ€” which +is the whole reason AdamW exists. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + + +class AdamW(torch.optim.Optimizer): + """Adam with decoupled weight decay.""" + + def __init__( + self, + params: Iterable[Any], + lr: float = 1e-3, + betas: tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 1e-2, + ) -> None: + if lr < 0.0: + raise ValueError(f"lr must be >= 0, got {lr}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"beta1 must be in [0, 1), got {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"beta2 must be in [0, 1), got {betas[1]}") + if eps < 0.0: + raise ValueError(f"eps must be >= 0, got {eps}") + if weight_decay < 0.0: + raise ValueError(f"weight_decay must be >= 0, got {weight_decay}") + super().__init__( + params, {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay} + ) + + @torch.no_grad() + def step(self, closure: Any = None) -> Any: # type: ignore[override] + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + lr = group["lr"] + beta1, beta2 = group["betas"] + eps = group["eps"] + weight_decay = group["weight_decay"] + + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + state = self.state[p] + if not state: + state["step"] = 0 + state["exp_avg"] = torch.zeros_like(p) + state["exp_avg_sq"] = torch.zeros_like(p) + + state["step"] += 1 + t = state["step"] + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + + if weight_decay != 0.0: + p.mul_(1.0 - lr * weight_decay) + + exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) + + bias_correction1 = 1.0 - beta1**t + bias_correction2 = 1.0 - beta2**t + + denom = (exp_avg_sq.sqrt() / bias_correction2**0.5).add_(eps) + p.addcdiv_(exp_avg, denom, value=-lr / bias_correction1) + + return loss + + +def build_param_groups( + module: nn.Module, weight_decay: float +) -> tuple[dict[str, Any], dict[str, Any]]: + """Split parameters into decayed (matrices) and undecayed (1-D) groups. + + Weight decay is a prior toward small *weights*. Applying it to a 1-D + RMSNorm gain or a bias pulls a learned scale toward zero for no + regularisation benefit, so those are excluded โ€” standard practice, and + cheap to get wrong silently. + + Tied parameters appear under more than one name but are one tensor; + de-duplicating by identity keeps them from being decayed twice per step. + """ + seen: set[int] = set() + decay: list[nn.Parameter] = [] + no_decay: list[nn.Parameter] = [] + + for param in module.parameters(): + if not param.requires_grad or id(param) in seen: + continue + seen.add(id(param)) + (decay if param.ndim >= 2 else no_decay).append(param) + + return ( + {"params": decay, "weight_decay": weight_decay}, + {"params": no_decay, "weight_decay": 0.0}, + ) diff --git a/training/scheduler.py b/training/scheduler.py new file mode 100644 index 0000000..bc8694a --- /dev/null +++ b/training/scheduler.py @@ -0,0 +1,67 @@ +"""Cosine learning-rate decay with linear warmup (TDR-009). + + step < warmup : base_lr * (step + 1) / warmup โ€” linear ramp + otherwise : min_lr + ยฝ(base_lr โˆ’ min_lr)(1 + cos(ฯ€ ยท progress)) + +Warmup uses ``(step + 1) / warmup`` rather than ``step / warmup`` so the very +first step actually moves; a zero learning rate on step 0 is a wasted step. + +The two pieces meet exactly at ``step == warmup``: the ramp's last value and the +cosine's value at ``progress = 0`` are both ``base_lr``, so there is no +discontinuity to show up as a loss spike. +""" + +from __future__ import annotations + +import math + +import torch + + +class CosineScheduler: + """Stateless schedule โ€” the LR is a pure function of the step number. + + Stateless on purpose: resuming from a checkpoint only needs the step count, + with no scheduler state to save, restore, or get out of sync with the + optimizer. + """ + + def __init__(self, base_lr: float, min_lr: float, warmup_steps: int, total_steps: int) -> None: + if min_lr > base_lr: + raise ValueError(f"min_lr ({min_lr}) must not exceed base_lr ({base_lr})") + if warmup_steps < 0: + raise ValueError(f"warmup_steps must be >= 0, got {warmup_steps}") + if total_steps < warmup_steps: + raise ValueError( + f"total_steps ({total_steps}) must be >= warmup_steps ({warmup_steps})" + ) + self.base_lr = base_lr + self.min_lr = min_lr + self.warmup_steps = warmup_steps + self.total_steps = total_steps + + def lr_at(self, step: int) -> float: + if step < self.warmup_steps: + return self.base_lr * (step + 1) / self.warmup_steps + + decay_steps = max(1, self.total_steps - self.warmup_steps) + progress = min(1.0, (step - self.warmup_steps) / decay_steps) + cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) + return self.min_lr + (self.base_lr - self.min_lr) * cosine + + def apply(self, optimizer: torch.optim.Optimizer, step: int) -> float: + """Write the rate for ``step`` onto every param group. Returns it.""" + lr = self.lr_at(step) + for group in optimizer.param_groups: + group["lr"] = lr + return lr + + @classmethod + def from_config(cls, cfg: dict) -> CosineScheduler: + t = cfg["training"] + return cls( + base_lr=float(t["lr"]), + min_lr=float(t.get("min_lr", float(t["lr"]) * 0.1)), + warmup_steps=int(t.get("warmup_steps", 0)), + total_steps=int(t["steps"]), + ) diff --git a/training/train.py b/training/train.py new file mode 100644 index 0000000..99c3dbf --- /dev/null +++ b/training/train.py @@ -0,0 +1,190 @@ +"""Training entry point. + +Usage: + python -m training.train --config configs/default.yaml + python -m training.train --config configs/default.yaml training.lr=1e-4 model.n_layers=12 + +This is a thin orchestration layer: it loads the config, wires together the +dataset / model / optimizer, and runs the training loop. The heavy lifting is +implemented in the respective modules as they are built out (see docs/TDD.md +and docs/BUILD_ORDER.md for the dependency-aware plan). +""" + +import argparse +import os +import random +from pathlib import Path +from typing import Any + +import torch +import yaml + +from dataset import DataEngine +from model import CausalLM +from training.device import autocast_dtype, resolve_device +from training.trainer import Trainer + + +def load_config(path: str) -> dict[str, Any]: + with open(path, encoding="utf-8") as f: + data = yaml.safe_load(f) + return dict(data) if data else {} + + +def set_deep(cfg: dict[str, Any], dotted_key: str, value: Any) -> None: + keys = dotted_key.split(".") + node = cfg + for k in keys[:-1]: + node = node.setdefault(k, {}) + node[keys[-1]] = value + + +def coerce(raw: str) -> Any: + """Parse an override value with the same rules that parsed the YAML file. + + Reusing the YAML scalar parser (rather than a hand-rolled int/float/bool + ladder) means ``lr=3e-4`` is a float, ``amp=true`` is a bool, and + ``out_dir=checkpoints`` is a string, all exactly as they would be if written + into the config file -- so an override can never produce a differently typed + value than the key it replaces. + """ + try: + value = yaml.safe_load(raw) + except yaml.YAMLError: + return raw + + # PyYAML implements YAML 1.1, whose float grammar demands a decimal point + # in scientific notation: "3.0e-4" parses, "3e-4" comes back as a string. + # Learning rates are almost always written the short way, so rescue it -- + # otherwise `training.lr=3e-4` sets the LR to the *string* "3e-4". + if isinstance(value, str): + try: + return float(value) + except ValueError: + return value + return value + + +def apply_overrides(cfg: dict[str, Any], overrides: list[str]) -> dict[str, Any]: + for item in overrides: + key, sep, raw = item.partition("=") + if not sep: + raise ValueError(f"override {item!r} is not of the form dotted.key=value") + set_deep(cfg, key, coerce(raw)) + return cfg + + +def build_val_engine(cfg: dict[str, Any]) -> DataEngine | None: + """A second engine over ``dataset.val_path``, or None if none is configured. + + Held-out evaluation is opt-in rather than automatic: the synthetic source has + no meaningful validation split (it is uniform noise by construction), and + silently evaluating against training data would report a number that looks + like generalisation and is not. + """ + val_path = cfg.get("dataset", {}).get("val_path") + if not val_path: + return None + val_cfg = { + **cfg, + "dataset": {**cfg["dataset"], "source": "file", "path": val_path, "shuffle": False}, + } + return DataEngine.from_config(val_cfg, vocab_size=cfg["model"]["vocab_size"]) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Train the language model") + parser.add_argument("--config", default="configs/default.yaml") + parser.add_argument("-v", "--verbose", action="store_true", help="echo the full config") + parser.add_argument("overrides", nargs="*", help="dotted.key=value overrides") + args = parser.parse_args() + + cfg = apply_overrides(load_config(args.config), args.overrides) + + seed = int(cfg.get("seed", 42)) + random.seed(seed) + torch.manual_seed(seed) + + device = resolve_device(cfg.get("device")) + amp_dtype = autocast_dtype(device, cfg["training"].get("amp_dtype", "bfloat16")) + + out_dir = Path(cfg["training"]["out_dir"]) + os.makedirs(out_dir, exist_ok=True) + + if args.verbose: + print("Loaded config from", args.config) + print(yaml.safe_dump(cfg, sort_keys=False)) + + model = CausalLM.from_config(cfg) + engine = DataEngine.from_config(cfg, vocab_size=cfg["model"]["vocab_size"]) + trainer = Trainer(model, engine, cfg, device=device) + val_engine = build_val_engine(cfg) + + steps = int(cfg["training"]["steps"]) + log_every = max(1, int(cfg["training"].get("log_every", 50))) + ckpt_every = int(cfg["training"].get("ckpt_every", 0)) + eval_every = int(cfg["training"].get("eval_every", 0)) + eval_batches = int(cfg["training"].get("eval_batches", 20)) + patience = int(cfg["training"].get("early_stop_patience", 0)) + + print( + f"device {device} | amp {amp_dtype or 'off (fp32)'} | " + f"{model.num_parameters() / 1e6:.2f}M params | {steps} steps" + + ("" if val_engine else " | no validation split configured") + ) + + last_evaluated = -1 + + def report_validation(step: int) -> None: + nonlocal last_evaluated + # Guard against the final report duplicating an eval_every hit that just + # landed on the same step. + if val_engine is None or step == last_evaluated: + return + last_evaluated = step + val = trainer.evaluate(val_engine, batches=eval_batches) + marker = "" + if val["val_improved"]: + # Save on improvement, not on a step interval: a periodic checkpoint + # will miss the optimum and leave only over-trained weights on disk. + trainer.save(str(out_dir / "best.pt")) + marker = " <- best, saved" + print( + f"step {step:>6}/{steps} val_loss {val['val_loss']:.4f} " + f"val_ppl {val['val_perplexity']:.2f}{marker}" + ) + + for _ in range(steps): + metrics = trainer.step() + n = trainer.step_count + if n % log_every == 0 or n == steps: + print( + f"step {n:>6}/{steps} loss {metrics['loss']:.4f} " + f"lr {metrics['lr']:.2e} grad_norm {metrics['grad_norm']:.3f} " + f"{metrics['tokens_per_sec']:,.0f} tok/s" + ) + if eval_every and n % eval_every == 0: + report_validation(n) + if ckpt_every and n % ckpt_every == 0: + trainer.save(str(out_dir / f"step_{n}.pt")) + if trainer.should_stop_early(patience): + print( + f"early stop at step {n}: {trainer.evals_since_improvement} evaluations " + f"without improvement (best val_loss {trainer.best_val_loss:.4f} " + f"at step {trainer.best_step})" + ) + break + + report_validation(trainer.step_count) + if trainer.best_step >= 0: + print( + f"best val_loss {trainer.best_val_loss:.4f} at step {trainer.best_step} " + f"-> {out_dir / 'best.pt'}" + ) + final = out_dir / "final.pt" + trainer.save(str(final)) + print(f"saved {final}") + + +if __name__ == "__main__": + main() diff --git a/training/trainer.py b/training/trainer.py new file mode 100644 index 0000000..fe1aa48 --- /dev/null +++ b/training/trainer.py @@ -0,0 +1,250 @@ +"""The training loop (TDD ยง5). + +One ``step()`` is one *optimizer* step, which may span several micro-batches: + + for _ in range(grad_accum_steps): + loss = cross_entropy(model(input_ids), labels) / grad_accum_steps + loss.backward() + grad_norm = clip_grad_norm_(params, grad_clip) + scheduler.apply(optimizer, step) + optimizer.step(); optimizer.zero_grad() + +Dividing the micro-batch loss by ``grad_accum_steps`` is what makes accumulation +equivalent to a single larger batch rather than to a larger *learning rate* โ€” +gradients sum across the backward passes, so without the division the effective +step is K times too big. + +Mixed precision follows the device (TDR-019): ``GradScaler`` exists only for +CUDA fp16. bf16 has fp32's exponent range and needs no loss scaling, and MPS has +no scaler at all, so on those paths autocast runs bare. +""" + +from __future__ import annotations + +import math +import time +from collections.abc import Callable +from contextlib import nullcontext +from typing import Any, Protocol + +import torch +import torch.nn.functional as F +from torch import nn + +from training.checkpoint import load_checkpoint, save_checkpoint +from training.device import autocast_dtype, resolve_device, supports_grad_scaler +from training.optimizer import AdamW, build_param_groups +from training.scheduler import CosineScheduler + + +class BatchSource(Protocol): + """Anything that can hand over the next batch โ€” DataEngine satisfies this.""" + + def next_batch(self, device: torch.device | str | None = ...) -> dict[str, torch.Tensor]: ... + + +class Trainer: + """Owns the model, optimizer, schedule and step counter for one run.""" + + def __init__( + self, + model: nn.Module, + engine: BatchSource, + cfg: dict[str, Any], + *, + device: torch.device | None = None, + ) -> None: + t = cfg.get("training", {}) + self.cfg = cfg + self.model = model + self.engine = engine + self.device = device if device is not None else resolve_device(cfg.get("device")) + self.model.to(self.device) + + self.grad_accum_steps = max(1, int(t.get("grad_accum_steps", 1))) + self.grad_clip = float(t.get("grad_clip", 0.0)) + self.total_steps = int(t.get("steps", 0)) + self.log_every = int(t.get("log_every", 50)) + self.step_count = 0 + # Best-validation tracking, so a run can keep its best weights rather + # than only its most recent ones. + self.best_val_loss = float("inf") + self.best_step = -1 + self.evals_since_improvement = 0 + + decay, no_decay = build_param_groups(model, float(t.get("weight_decay", 0.0))) + self.optimizer = AdamW( + [decay, no_decay], + lr=float(t.get("lr", 3e-4)), + betas=(float(t.get("beta1", 0.9)), float(t.get("beta2", 0.95))), + eps=float(t.get("eps", 1e-8)), + ) + self.scheduler = CosineScheduler( + base_lr=float(t.get("lr", 3e-4)), + min_lr=float(t.get("min_lr", float(t.get("lr", 3e-4)) * 0.1)), + warmup_steps=int(t.get("warmup_steps", 0)), + total_steps=max(1, self.total_steps), + ) + + self.amp_dtype = ( + autocast_dtype(self.device, str(t.get("amp_dtype", "bfloat16"))) + if t.get("amp", True) + else None + ) + use_scaler = supports_grad_scaler(self.device) and self.amp_dtype is torch.float16 + self.scaler = torch.amp.GradScaler(enabled=use_scaler) + + # ------------------------------------------------------------------ + # Loss + # ------------------------------------------------------------------ + def _autocast(self) -> Any: + if self.amp_dtype is None: + return nullcontext() + return torch.autocast(device_type=self.device.type, dtype=self.amp_dtype) + + def loss_on(self, batch: dict[str, torch.Tensor]) -> torch.Tensor: + """Next-token cross-entropy. Lives here, not in the model (BUILD_ORDER).""" + input_ids = batch["input_ids"].to(self.device) + labels = batch["labels"].to(self.device) + with self._autocast(): + logits = self.model(input_ids) + # float() before cross-entropy: under bf16 autocast the softmax + # denominator is summed over the whole vocabulary, where reduced + # mantissa shows up directly in the loss. + return F.cross_entropy(logits.float().reshape(-1, logits.shape[-1]), labels.reshape(-1)) + + # ------------------------------------------------------------------ + # One optimizer step + # ------------------------------------------------------------------ + def step(self) -> dict[str, float]: + self.model.train() + started = time.perf_counter() + total_loss = 0.0 + tokens = 0 + + for _ in range(self.grad_accum_steps): + batch = self.engine.next_batch(self.device) + loss = self.loss_on(batch) / self.grad_accum_steps + self.scaler.scale(loss).backward() + total_loss += loss.item() + tokens += batch["input_ids"].numel() + + if self.grad_clip > 0: + # Unscale first, or the clip threshold is applied to scaled + # gradients and effectively does nothing. + self.scaler.unscale_(self.optimizer) + grad_norm = float( + torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip) + ) + else: + grad_norm = float("nan") + + lr = self.scheduler.apply(self.optimizer, self.step_count) + self.scaler.step(self.optimizer) + self.scaler.update() + self.optimizer.zero_grad(set_to_none=True) + self.step_count += 1 + + elapsed = max(time.perf_counter() - started, 1e-9) + return { + "step": float(self.step_count), + "loss": total_loss, + "lr": lr, + "grad_norm": grad_norm, + "tokens": float(tokens), + "tokens_per_sec": tokens / elapsed, + } + + # ------------------------------------------------------------------ + # Loop + # ------------------------------------------------------------------ + def train( + self, + steps: int | None = None, + *, + on_log: Callable[[dict[str, float]], None] | None = None, + ) -> list[dict[str, float]]: + history = [] + for _ in range(steps if steps is not None else self.total_steps): + metrics = self.step() + history.append(metrics) + if on_log and self.step_count % max(1, self.log_every) == 0: + on_log(metrics) + return history + + # ------------------------------------------------------------------ + # Evaluation + # ------------------------------------------------------------------ + @torch.no_grad() + def evaluate(self, engine: BatchSource, *, batches: int = 20) -> dict[str, float]: + """Mean loss and perplexity over ``batches`` from a held-out source. + + Averages the *loss* and exponentiates once at the end. Averaging + perplexities instead would be wrong โ€” perplexity is exp of a mean, and + the mean of exponentials is not the exponential of the mean (Jensen), so + per-batch averaging reports a number that is always too high. + + Also tracks the best validation loss seen and how many evaluations have + passed without improvement, which is what makes best-checkpoint saving + and early stopping possible. Without it a run happily trains past its own + optimum and keeps only the *worst* model: the first TinyShakespeare run + bottomed out at validation perplexity 98.9 around step 500 and finished at + 304.6, and because checkpoints were periodic rather than + best-tracked, the good weights were never written to disk. + + Restores the previous train/eval mode so this can be called mid-loop + without silently leaving dropout off for the rest of training. + """ + was_training = self.model.training + self.model.eval() + try: + total = 0.0 + for _ in range(batches): + total += self.loss_on(engine.next_batch(self.device)).item() + mean = total / max(1, batches) + + improved = mean < self.best_val_loss + if improved: + self.best_val_loss = mean + self.best_step = self.step_count + self.evals_since_improvement = 0 + else: + self.evals_since_improvement += 1 + + return { + "val_loss": mean, + "val_perplexity": math.exp(min(mean, 80.0)), + "val_improved": float(improved), + "best_val_loss": self.best_val_loss, + "best_step": float(self.best_step), + "evals_since_improvement": float(self.evals_since_improvement), + } + finally: + self.model.train(was_training) + + def should_stop_early(self, patience: int) -> bool: + """True once ``patience`` consecutive evaluations have failed to improve. + + ``patience <= 0`` disables the check. Patience is counted in + *evaluations*, not steps, so it means the same thing regardless of how + often the loop evaluates. + """ + return patience > 0 and self.evals_since_improvement >= patience + + # ------------------------------------------------------------------ + # Checkpointing + # ------------------------------------------------------------------ + def save(self, path: str) -> None: + save_checkpoint( + path, + model=self.model, + optimizer=self.optimizer, + step=self.step_count, + cfg=self.cfg, + ) + + def load(self, path: str) -> None: + payload = load_checkpoint( + path, model=self.model, optimizer=self.optimizer, map_location=self.device + ) + self.step_count = int(payload.get("step", 0))