v2: 363M hero run (Muon hybrid, WSD, Liger, SmolLM 75/15/10 mix)

.gitignore

@@ -0,0 +1,19 @@
+# caches
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ipynb_checkpoints/
+
+# data + model artifacts (regenerated; never commit)
+data/
+checkpoints/
+results/
+*.bin
+*.pt
+*.jsonl
+wandb/
+
+# env
+.venv/
+venv/
+.env

README.md

@@ -0,0 +1,132 @@
+# Matilda-Mini
+
+A sub-400M-parameter language model **trained from scratch** — and, more to the
+point, the **training infrastructure** around it: distributed-ready training loop,
+crash-safe checkpoint/resume, fault tolerance, observability, and a verifiable
+data pipeline. Built from first principles in PyTorch, no training frameworks.
+
+> Not a fine-tune. Not a wrapper. Random init → a working LM, trained by code in
+> this repo. The model is standard-modern; the **systems work is the point**.
+
+## Two configs, one stack
+
+| Config | Params | Tokens | Optimizer | Data | Schedule | Purpose |
+|---|---|---|---|---|---|---|
+| `configs/base_124m.json` | 114M | ~3B | AdamW | FineWeb-Edu | warmup+cosine | v1 — original portfolio run |
+| `configs/base_350m.json` | **363M** | **~15B** | **Muon+AdamW hybrid** | **SmolLM-corpus 75/15/10** | **WSD (MiniCPM)** | **v2 — hero run, deep+thin shape, Liger kernels** |
+
+The v2 hero run is designed around three claims: (1) **Muon scales to 350M**
+(only validated at 124M publicly + our v1 ablation); (2) **deep+thin × Muon
+interaction** (32L × 960d, ~50% deeper than Pythia-410M at comparable width);
+(3) **Liger Kernel in a custom nn.Module** unlocks BS ~50% larger on A100 40GB,
+which is the cost story. Target comparison: **token-matched vs Pythia-410M's
+published checkpoints**, ~$55-60 on A100 SXM4 vs Pythia's $1000+ to reach the
+same checkpoint.
+
+## Why this exists
+
+This is a portfolio project for an LLM **training-infrastructure** role. The
+interesting problems in training large models aren't the architecture (well
+understood) — they're the systems: making multi-day runs reliable, resumable,
+observable, and fast on the hardware you have. So this repo is deliberately
+weighted toward operational excellence over architectural novelty.
+
+## Architecture (`src/matilda/model.py`)
+
+A modern dense decoder-only transformer — the same recipe as Llama/Qwen-class
+models. Shape is a runtime knob:
+
+| Component | v1 (124M) | v2 (363M, hero) |
+|-----------|-----------|-----------------|
+| Layers × d_model × n_heads | 12 × 768 × 12 | **32 × 960 × 12** |
+| KV heads (GQA) | 4 | 4 |
+| Head dim | 64 | **80** |
+| Seq length | 1024 | **2048** |
+| Positions | RoPE | RoPE (optional Liger fused) |
+| Normalization | RMSNorm (fp32 reduction) | RMSNorm (optional Liger fused) |
+| MLP | SwiGLU (8/3 sizing) | SwiGLU (8/3 sizing) |
+| QK-Norm | on | on |
+| Embedding tying | on | on |
+| Loss | CE | CE **+ z-loss (1e-4)** |
+| Vocab projection | `nn.Linear` + `F.cross_entropy` | **Liger fused linear+CE** (skips logit materialization) |
+| Optimizer | AdamW | **Hybrid Muon (2D) + AdamW (embed/norm)** |
+| LR schedule | warmup + cosine | **warmup + 80% stable + 20% linear decay (WSD)** |
+
+## Training infrastructure (the actual deliverable)
+
+| Capability | Where | What it does |
+|-----------|-------|--------------|
+| **Bit-for-bit resume** | `checkpoint.py` | atomic writes; saves model+opt+sched+step+**RNG+dataloader position**; a killed run resumes to a loss curve identical to the uninterrupted one (`< 1e-6`, tested) |
+| **Fault tolerance** | `train.py` | NaN/Inf guard (skip+log+abort-after-N); SIGTERM → checkpoint-and-exit for spot-instance death |
+| **Observability** | `monitor.py` | MFU (incl. attention FLOPs), tokens/s, rolling step-time (catches throttling), grad-norm, peak GPU mem → always-on `metrics.jsonl` + optional W&B |
+| **Throughput** | `train.py` | bf16 autocast, Flash-SDPA, `torch.compile`, fused AdamW, TF32, pinned/non-blocking H2D, grad-accum with DDP `no_sync` |
+| **Data pipeline** | `data.py`, `scripts/prepare_data.py` | streams FineWeb-Edu → tokenizes → `uint16` shards with **SHA-256 manifest**; mmap'd, resumable `BinStream` |
+| **Optimizers** | `optim.py` | AdamW (correct param-group decay) + **Muon** (Newton-Schulz orthogonalization, hybrid with AdamW) |
+| **Reproducibility** | `train.py` | full config + **git SHA** logged per run; deterministic seeding |
+
+## Results
+
+**Validated (RTX 3090):** 30/30 tests pass on GPU, smoke + bit-for-bit resume
+clean, **53.4% MFU** at batch_size=24 with `torch.compile` (BS≥28 OOMs on the
+vocab projection — the expected memory hotspot).
+
+**Training run + ablations:** pending the A100 run. The ablation harness
+(`scripts/ablate.py`) emits `docs/ABLATIONS.md` — a controlled comparison, one
+change per row:
+
+| Variant | What it isolates |
+|---------|------------------|
+| baseline | full modern stack |
+| no_qk_norm | QK-Norm's stability contribution |
+| mha / mqa | GQA ratio vs full multi-head / multi-query |
+| muon | Muon vs AdamW convergence |
+
+Target (124M, ~3B tokens, vs Pythia-160M): HellaSwag ~30-35%, ARC-easy ~40-45%,
+PIQA ~60%.
+
+## Quickstart
+
+```bash
+pip install -r requirements.txt          # GPU: install torch from cu124 first (see runbook)
+pytest tests/ -q                          # 35 tests: correctness, resume, NaN guard, data integrity, WSD, z-loss
+
+# train (synthetic dry run, no data needed)
+python run.py --config configs/calibration.json --dry-run \
+  --set model.d_model=128 model.n_layers=2 train.total_steps=20 train.device=cpu train.compile=false
+
+# v1 run (124M / FineWeb-Edu / ~3B tokens)
+python scripts/prepare_data.py --out-dir data/fwedu --target-tokens 3000000000
+python run.py --config configs/base_124m.json --data-dir data/fwedu
+
+# v2 hero run (363M / SmolLM 75/15/10 / ~15B tokens) — A100 + liger-kernel required
+pip install liger-kernel
+python scripts/prepare_smollm_data.py --out-dir data/smollm_mix --target-tokens 15000000000
+python run.py --config configs/base_350m.json --data-dir data/smollm_mix
+```
+
+Full GPU procedure (validate → calibrate → ablate → train → eval) is in
+[`docs/INSTANCE_RUNBOOK.md`](docs/INSTANCE_RUNBOOK.md).
+
+## Repository layout
+
+```
+src/matilda/    config, model, optim, checkpoint, monitor, data, train
+scripts/        prepare_data.py (FineWeb-Edu), prepare_smollm_data.py (SmolLM 75/15/10 mix),
+                ablate.py (experiments), launch_vast.sh
+configs/        calibration.json (MFU tuning), base_124m.json (v1), base_350m.json (v2 hero)
+tests/          35 tests — model, checkpoint, train loop, data, optim, ablation, run, WSD, z-loss
+docs/           INSTANCE_RUNBOOK.md (operating manual)
+run.py          training entrypoint (--config + --set overrides)
+```
+
+## Testing
+
+35 tests run on CPU in ~2 min. Highlights: overfit-single-batch (the model can
+learn), causal-mask-no-leak (no future-token leakage), bit-for-bit resume,
+NaN-skip-then-recover, shard checksum corruption detection, Muon overfit, WSD
+three-phase shape, z-loss equals CE + lse² when on. BASE_350M shape test skips
+on CPU dev boxes without `liger-kernel`.
+
+```bash
+pytest tests/ -q
+```

configs/base_124m.json

@@ -0,0 +1,32 @@
+{
+  "_comment": "The long run. ~3B tokens = total_steps*batch*grad_accum*seq_len = 5814*24*21*1024 ~= 3.0B. These bs/accum are the VALIDATED 3090 ceiling (BS=24 @ 53.4% MFU; BS>=28 OOMs on the vocab projection). On A100, raise batch_size to its ceiling (expect 64-128) and re-derive grad_accum/total_steps to keep tokens/step ~0.5M and tokens ~3B (see runbook section 6). lr 6e-4; drop to 3e-4 if grad_norm spikes.",
+  "model": {
+    "vocab_size": 50257,
+    "max_seq_len": 1024,
+    "d_model": 768,
+    "n_layers": 12,
+    "n_heads": 12,
+    "n_kv_heads": 4,
+    "qk_norm": true
+  },
+  "train": {
+    "total_steps": 5814,
+    "warmup_steps": 300,
+    "grad_accum": 21,
+    "batch_size": 24,
+    "seq_len": 1024,
+    "lr": 6e-4,
+    "weight_decay": 0.1,
+    "grad_clip": 1.0,
+    "min_lr_ratio": 0.1,
+    "log_every": 10,
+    "ckpt_every": 500,
+    "keep_last": 3,
+    "device": "cuda",
+    "dtype": "bfloat16",
+    "compile": true,
+    "max_skips": 20,
+    "ckpt_dir": "checkpoints/base_124m",
+    "wandb_project": "matilda-mini"
+  }
+}

configs/base_350m.json

@@ -0,0 +1,41 @@
+{
+  "_comment": "350M hero run, 15B tokens (~2x Chinchilla). Tokens/step ~= 1.05M (32*16*2048); 15000 steps * 1.05M ~= 15.7B tokens. ETA on A100 SXM4 40GB at the target 50% MFU: ~55-60h, ~$55-60 at $1.017/hr. Re-derive batch_size/grad_accum on the instance: with Liger fused linear+CE we expect to push micro-batch to 32-48; if calibration says 48, drop grad_accum to ~11 to keep tokens/step ~= 1M. lr 3e-4 (AdamW arm); muon_lr 0.02 (validated in v1 ablation). Schedule = WSD with 80%-stable, 20%-linear-decay; min_lr_ratio 0.1 so the tail isn't dead. z_loss 1e-4 for logit stability at scale. use_liger=true is mandatory for the hero throughput claim — fails loudly if liger-kernel not installed.",
+  "model": {
+    "vocab_size": 50257,
+    "max_seq_len": 2048,
+    "d_model": 960,
+    "n_layers": 32,
+    "n_heads": 12,
+    "n_kv_heads": 4,
+    "mlp_ratio": 2.6666666666666665,
+    "mlp_multiple_of": 256,
+    "tie_weights": true,
+    "qk_norm": true,
+    "z_loss_coef": 1e-4,
+    "use_liger": true
+  },
+  "train": {
+    "total_steps": 15000,
+    "warmup_steps": 2000,
+    "grad_accum": 16,
+    "batch_size": 32,
+    "seq_len": 2048,
+    "lr": 3e-4,
+    "weight_decay": 0.1,
+    "grad_clip": 1.0,
+    "min_lr_ratio": 0.1,
+    "optimizer": "muon",
+    "muon_lr": 0.02,
+    "lr_schedule": "wsd",
+    "wsd_stable_share": 0.8,
+    "log_every": 10,
+    "ckpt_every": 1000,
+    "keep_last": 20,
+    "device": "cuda",
+    "dtype": "bfloat16",
+    "compile": true,
+    "max_skips": 20,
+    "ckpt_dir": "checkpoints/base_350m",
+    "wandb_project": "matilda-mini"
+  }
+}

configs/calibration.json

@@ -0,0 +1,26 @@
+{
+  "_comment": "15-min A100 smoke test. Goal: maximize MFU before the long run. Sweep batch_size with --set train.batch_size=N and read mfu_avg from the log. Few steps, frequent logging, compile on.",
+  "model": {
+    "vocab_size": 50257,
+    "max_seq_len": 1024,
+    "d_model": 768,
+    "n_layers": 12,
+    "n_heads": 12,
+    "n_kv_heads": 4
+  },
+  "train": {
+    "total_steps": 60,
+    "warmup_steps": 10,
+    "grad_accum": 1,
+    "batch_size": 24,
+    "seq_len": 1024,
+    "lr": 3e-4,
+    "log_every": 5,
+    "ckpt_every": 60,
+    "keep_last": 1,
+    "device": "cuda",
+    "dtype": "bfloat16",
+    "compile": true,
+    "ckpt_dir": "checkpoints/calib"
+  }
+}

conftest.py

@@ -0,0 +1,4 @@
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent / "src"))

docs/A100_KICKOFF_PROMPT.md

@@ -0,0 +1,110 @@
+# A100 Run — Kickoff Prompts (paste into Claude Code on the instance)
+
+Two prompts. **Prompt 1** validates + calibrates (cheap, ~$0.50) and stops.
+**Prompt 2** runs the paid training + ablations + eval. Run Prompt 1 first,
+read the calibration result, fill Prompt 2's batch numbers, then run it.
+
+GPU: **A100 40GB recommended** (124M model — 80GB is overkill; the memory
+hotspot is the vocab-projection logits, not weights). 3090 also works.
+Attention is PyTorch SDPA (built-in Flash 2) — no `flash-attn` package.
+
+---
+
+## PROMPT 1 — validate + calibrate, then STOP
+
+```
+You're on a remote GPU instance. Validate and calibrate a from-scratch LLM
+training project. Do NOT start paid training or ablations yet — stop and report.
+
+1. git clone https://huggingface.co/prometheus04/matilda-mini && cd matilda-mini
+2. Read docs/INSTANCE_RUNBOOK.md fully; follow it and respect every "Gate".
+3. Execute runbook sections 2-6:
+   - §2: nvidia-smi (report GPU + VRAM). Install torch from the cu124 index FIRST,
+     then `pip install -r requirements.txt`. Confirm CUDA + bf16 available.
+   - §3: `pytest tests/ -q` — gate is "30 passed". If a GPU-only failure appears,
+     STOP and show the full traceback.
+   - §4: tokenize ~20M tokens FineWeb-Edu, verify checksum, run 100 steps — loss
+     should fall from ~10.8 toward ~7.
+   - §5: rerun to 150 steps — must print "[resume] ... at step 100" and continue
+     with no loss spike.
+   - §6: sweep batch_size with configs/calibration.json; report batch_size -> MFU.
+4. STOP and report: (1) GPU+VRAM, (2) pytest result, (3) smoke loss + resume,
+   (4) the calibration table and the best batch_size.
+Use tmux for long commands. Never commit data/ or checkpoints/.
+```
+
+---
+
+## Between prompts — pick batch/accum/total_steps
+
+From the calibration sweep, take the `batch_size` (BS) with the best MFU that
+fits, then keep tokens/step ≈ 524288 and total tokens ≈ 3.0B:
+
+```
+grad_accum  = round(524288 / (BS * 1024))
+total_steps = round(3.0e9 / (BS * grad_accum * 1024))
+```
+
+Pre-computed safe defaults (use if you don't want to wait on the sweep):
+
+| GPU | BS | grad_accum | total_steps | tokens |
+|-----|----|-----------|-------------|--------|
+| A100 40GB | 64 | 8 | 5722 | 3.00B |
+| A100 80GB | 128 | 4 | 5722 | 3.00B |
+| 3090 (validated) | 24 | 21 | 5814 | 3.00B |
+
+---
+
+## PROMPT 2 — the paid run (fill BS / ACCUM / STEPS, then paste)
+
+```
+Continue in the matilda-mini directory on the instance. Calibration is done.
+Run the full training + ablations + eval. This is the paid run.
+
+Set these from calibration (A100 40GB defaults shown): BS=64, ACCUM=8, STEPS=5722.
+
+1. Tokenize the full dataset if not already present:
+   python scripts/prepare_data.py --out-dir data/fwedu --target-tokens 3000000000
+   (verify the manifest checksum after.)
+
+2. Launch the long run inside tmux so an SSH drop can't kill it:
+   tmux new -s train
+   python run.py --config configs/base_124m.json --data-dir data/fwedu \
+     --set train.batch_size=64 train.grad_accum=8 train.total_steps=5722
+   Detach with Ctrl-b d. Reattach with `tmux attach -t train`. If it ever dies,
+   re-run the same command — it auto-resumes from the latest checkpoint.
+   Watch checkpoints/base_124m/metrics.jsonl: loss falling, grad_norm ~0.2-1.5,
+   MFU flat, few/no nan_skip events. If grad_norm spikes / many NaN skips, lower
+   lr to 3e-4 (--set train.lr=3e-4) and resume.
+
+3. After training completes, run the architecture ablations (~$2):
+   python scripts/ablate.py --data-dir data/fwedu --tokens 300000000
+   This writes docs/ABLATIONS.md (the results table).
+
+4. Evaluate the final model:
+   pip install -q lm-eval
+   (export the checkpoint to a HF-style dir if needed, then) lm_eval --model hf \
+     --model_args pretrained=<dir>,dtype=bfloat16 \
+     --tasks hellaswag,arc_easy,piqa,winogrande --batch_size 32
+
+5. Save artifacts OFF the instance (it's ephemeral): the final ckpt_*.pt,
+   checkpoints/base_124m/metrics.jsonl, docs/ABLATIONS.md, and the eval JSON.
+   Upload to HF or scp them out. Report the final loss, MFU, and eval numbers.
+```
+
+---
+
+## Reference
+
+**Expected (124M, ~3B tokens):** final loss ~2.7-3.0; MFU ~40-50% (A100, compile
+on); HellaSwag ~30-35%, ARC-easy ~40-45%, PIQA ~60% (vs Pythia-160M: ~30/40/62).
+
+**Cost:** A100 40GB long run ~$4-8; ablations ~$2; calibration ~$0.50.
+
+**Troubleshooting:** OOM → lower BS, raise grad_accum to keep tokens/step.
+Low MFU (<25%) → confirm compile=true + bf16 engaged. SSH drop → tmux + re-run
+(auto-resume). git_sha "unknown" → only if cloned non-git (HF clone IS git, fine).
+
+**After the run:** paste the final loss / MFU / eval numbers + docs/ABLATIONS.md
+back to laptop-Claude to fill the README "Results" section, then flip the repo
+public as the finished portfolio piece.

docs/DEEPSEEK_REVIEW_BUNDLE.md

@@ -0,0 +1,1520 @@
+# Review Prompt for DeepSeek
+
+> Paste everything below into DeepSeek. DeepSeek cannot read your disk, so after
+> the prompt, paste the contents of the files it asks for (or attach them). The
+> file map and paths are included so you both share a mental model.
+
+---
+
+## ROLE
+
+You are a senior LLM training-infrastructure engineer reviewing a from-scratch
+small language model project. Be rigorous and skeptical. I want correctness bugs,
+reliability gaps, and anything that would embarrass me in a technical interview —
+not encouragement. Prioritize findings by severity (CRITICAL / HIGH / MEDIUM / LOW).
+
+## WHY THIS PROJECT EXISTS (the job I'm applying for)
+
+I'm building this as a portfolio piece for an **AI Training Infrastructure Engineer**
+role at **Maincode** (Melbourne, Australia). Maincode trains "Matilda", the first
+LLM built and trained from scratch in Australia. Key facts about the role:
+
+- It is an **infrastructure** role, NOT a research/architecture role. The work is:
+  distributed training pipelines, data ingestion/preprocessing, experiment
+  management, checkpointing, reproducibility, monitoring/observability, debugging
+  long-running (days–weeks) training jobs, optimizing throughput (compute/memory/
+  data), and diagnosing failures that appear hours into a run.
+- Explicitly NOT about: wrapping external APIs, prompt engineering, user-facing apps.
+- Stack: Python, PyTorch/JAX, reliable infra for large compute, debugging
+  distributed systems and long jobs.
+- They value engineers who understand how systems behave over long runtimes and
+  who prefer understanding internals over relying on abstraction.
+
+So the project is **deliberately weighted toward operational excellence**
+(checkpoint/resume, fault tolerance, reproducibility, MFU/observability) over
+architectural breadth. The architecture is modern but standard; the infra is the
+flex. Please judge it through that lens.
+
+## ORIGINAL PLAN & KEY DECISIONS (locked)
+
+- **Goal:** train a sub-200M-param modern transformer from scratch, end-to-end,
+  with a clean ablation table, on a tight budget.
+- **Hardware:** rent ONE A100 80GB on Vast.ai for the single long paid run.
+- **Long paid run = modernized DENSE ~114M, maximize MFU.** Rationale: MoE on a
+  single GPU collapses MFU to ~20% (needs expert parallelism across devices);
+  dense hits 45–50%. "DeepSeek-inspired" here means the efficiency stack
+  (Flash/SDPA, torch.compile, fused optimizer, bf16, Muon, μP, WSD), NOT sparsity.
+  A100 has no FP8 (bf16 ceiling; FP8 would need H100).
+- **MoE is built only as a CHEAP SHORT ablation** (~500M tokens, ~$0.60) with an
+  honest write-up of the single-GPU MFU drop. Not the headline run.
+- **Tokenizer:** GPT-2 50257 (tiktoken "gpt2") — fits uint16, comparable to Pythia
+  baselines. (We caught and avoided a bug from our source primer that paired
+  cl100k ~100k vocab with uint16, which overflows 65535.)
+- **Dataset:** FineWeb-Edu sample-10BT primary (best small-scale benchmark
+  movement + Pythia comparability). Loader is dataset-agnostic so DCLM /
+  Nemotron-CC are drop-in data-ablation rows.
+- **Cost discipline:** only the ONE long run costs real money (~$5–15). Ablations
+  ~$0.60 each. ALL correctness validated FREE on Colab/CPU first; a ~$0.20 A100
+  smoke test maximizes MFU before committing to the long run.
+
+### Phase plan
+1. ✅ Model + sanity tests
+2. ✅ Infra harness (checkpoint/resume, NaN-guard, monitor/MFU, optimizer, loop)
+3. ✅ Data pipeline (FineWeb→uint16 shards + checksum, mmap BinStream)
+4. ⏳ NEXT: Colab GPU correctness pass + $0.20 A100 MFU calibration
+5. Long dense run (~3B tokens)
+6. Architecture ablations (RoPE/RMSNorm/SwiGLU/GQA/QK-norm on/off)
+7. Muon vs AdamW
+8. MoE ablation
+9. Eval (lm-eval-harness: HellaSwag / ARC-easy / PIQA vs Pythia-160M)
+
+## WHAT IS BUILT (Phases 1–3 complete; 18/18 tests pass on CPU)
+
+Repository root on my machine:
+`C:\Users\Swajay\Downloads\trade\matilda-mini\`
+
+File map (line counts):
+
+```
+src/matilda/
+  config.py      (65)   frozen dataclass; DEV_TINY + BASE_124M (114M total / 75.5M non-embed)
+  model.py      (199)   dense decoder: RoPE, RMSNorm, SwiGLU, GQA, QK-Norm, weight tying, residual-scaled init
+  optim.py       (48)   AdamW with param groups (no WD on norms/biases/embeddings); cosine+warmup schedule
+  checkpoint.py  (95)   atomic write (tmp->os.replace), rotation, saves model+opt+sched+step+RNG+dataloader pos
+  monitor.py     (70)   MFU, tokens/s, rolling step-time window, per-GPU peak-TFLOPS table
+  data.py       (182)   ShardWriter (+SHA-256 manifest), verify_manifest, SyntheticStream, BinStream (mmap, resumable)
+  train.py      (236)   loop: bf16 autocast, grad-accum + DDP no_sync, grad-clip, NaN/Inf guard, SIGTERM->checkpoint, auto-resume
+  __init__.py     (4)
+scripts/
+  prepare_data.py (73)  stream HF dataset (default FineWeb-Edu sample-10BT) -> tokenize gpt2 -> uint16 shards + verified manifest
+tests/            (18 tests total, all green on CPU/torch 2.7.1)
+  test_model.py   (86)  forward shapes, weight-tying, init-loss≈log(V), CAUSAL-MASK-NO-FUTURE-LEAK, GQA head counts, OVERFIT-SINGLE-BATCH
+  test_checkpoint.py(116) BIT-FOR-BIT RESUME (post-resume losses match uninterrupted run <1e-6), rotation/latest
+  test_train.py   (79)  loop runs+checkpoints, loop-level resume, NaN-abort, NaN-skip-then-recover, MFU sanity
+  test_data.py    (78)  shard round-trip, CHECKSUM CORRUPTION DETECTION, BinStream next-token shift, deterministic resume, multi-shard
+conftest.py, pytest.ini, requirements.txt
+```
+
+Verified facts:
+- Overfit-one-batch: loss → <0.1 on a fixed batch (model can learn).
+- Causal mask: perturbing token t leaves all logits at positions < t bit-identical.
+- Resume: interrupt at the halfway point, rebuild fresh model/opt/sched/stream,
+  restore, continue → per-step losses match an uninterrupted run to <1e-6.
+- Data: corrupting one byte of a shard is caught by verify_manifest.
+- BinStream → Trainer integration: loss drops on structured data with train.py unchanged.
+
+## WHAT I WANT FROM YOU (review checklist)
+
+Please go through these and report concrete findings with file/function references:
+
+1. **Transformer correctness.** RoPE (half-rotation convention, cos/sin build,
+   applied to Q/K only, before/after QK-norm ordering), GQA repeat_interleave vs
+   repeat semantics, SwiGLU 2/3 sizing, RMSNorm fp32 reduction, weight tying,
+   residual-projection init scaling 1/sqrt(2*n_layers). Any subtle bug?
+2. **Numerical stability for a long bf16 run.** Is QK-norm enough? Should I add
+   z-loss / logit soft-cap / attention-logit cap? Embedding init scale?
+3. **Checkpoint/resume completeness.** Anything that influences the next step that
+   I'm NOT saving (RNG, data position, scaler, AMP state)? Atomicity on Vast spot kill.
+4. **NaN/Inf guard design.** Is skip-batch correct, or should I roll back to last
+   checkpoint? Should the guard also catch inf grad-norm before clip (it does)?
+   Is consecutive-skip abort the right policy?
+5. **Throughput / MFU.** Is the 6N FLOPs/token approximation appropriate, or should
+   I include attention FLOPs (2*L*T*d term) at seq_len 1024? Will I actually hit
+   45-50% MFU on an A100 with this code, or what's missing (e.g., no torch.compile
+   tuning, no fused/foreach, dataloader not overlapped with compute, H2D copies not
+   pinned/non_blocking)?
+6. **DDP correctness.** no_sync usage on non-final micro-steps is wired but DDP
+   wrapping isn't applied (single-GPU). If I later wrap in DDP, what breaks
+   (checkpoint should save model.module, num_params unwrap, sampler sharding,
+   rank-0-only logging/checkpointing)?
+7. **Data pipeline.** Random-window sampling vs sequential cursor — implications for
+   epoch coverage and resume. Cross-document contamination from packing without
+   attention masking at seq_len 1024 — does it matter at this scale?
+8. **Chinchilla / token budget.** For ~114M total (~75M non-embedding), what token
+   count is right for a portfolio run (compute-optimal ~20x vs over-train 50-100x)?
+9. **Anything that would make a Maincode interviewer wince.** Missing observability,
+   missing reproducibility (git SHA logging?), config hygiene, test gaps.
+10. **Highest-ROI additions** before I spend money, ranked, given the infra-role framing.
+
+For each finding give: severity, file/location, the problem, and the concrete fix.
+
+---
+
+SOURCE FILES FOLLOW. I will paste them below in this order: config.py, model.py,
+optim.py, checkpoint.py, monitor.py, data.py, train.py, then the four test files.
+
+
+===== FILE: src/matilda/config.py =====
+```python
+"""Model + run configuration.
+
+One frozen dataclass is the single source of truth for a run. It gets logged
+verbatim (alongside the git SHA) so any run is exactly reproducible. Nothing in
+the training stack reads a hyperparameter from anywhere else.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field, asdict
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    # --- vocab / sequence ---
+    vocab_size: int = 50257          # GPT-2 (tiktoken "gpt2"); fits uint16
+    max_seq_len: int = 1024
+
+    # --- transformer shape ---
+    d_model: int = 768
+    n_layers: int = 12
+    n_heads: int = 12
+    n_kv_heads: int = 4              # GQA: n_heads must be divisible by this
+    mlp_ratio: float = 8 / 3         # SwiGLU keeps params ~= 4x dense MLP
+    mlp_multiple_of: int = 256       # round hidden dim up to this for kernels
+
+    # --- numerics / regularization ---
+    norm_eps: float = 1e-6
+    rope_theta: float = 10000.0
+    init_std: float = 0.02
+    dropout: float = 0.0             # pretraining: keep at 0
+
+    # --- structural switches ---
+    tie_weights: bool = True         # share embedding <-> lm_head
+    qk_norm: bool = True             # RMSNorm on Q,K before attention
+
+    @property
+    def head_dim(self) -> int:
+        assert self.d_model % self.n_heads == 0, "d_model must divide n_heads"
+        return self.d_model // self.n_heads
+
+    def __post_init__(self) -> None:
+        assert self.n_heads % self.n_kv_heads == 0, \
+            "n_heads must be divisible by n_kv_heads (GQA grouping)"
+        assert self.head_dim % 2 == 0, "head_dim must be even for RoPE"
+
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        d["head_dim"] = self.head_dim
+        return d
+
+
+# Tiny config for free correctness validation on Colab T4 / RTX 3060.
+# Real vocab kept so tokenizer wiring is exercised; dims shrunk so it runs fast.
+DEV_TINY = ModelConfig(
+    vocab_size=50257,
+    max_seq_len=256,
+    d_model=128,
+    n_layers=2,
+    n_heads=4,
+    n_kv_heads=2,
+)
+
+# ~124M params (GPT-2 base footprint). The intended A100 long-run config.
+BASE_124M = ModelConfig()
+```
+
+===== FILE: src/matilda/model.py =====
+```python
+"""Modernized dense decoder-only transformer.
+
+Block recipe (pre-norm): x = x + attn(rmsnorm(x)); x = x + swiglu(rmsnorm(x)).
+Modernizations baked in from the start (validated free on Colab before any
+paid run): RoPE, RMSNorm, SwiGLU, GQA, QK-Norm, weight tying, residual-scaled
+init. This is the architecture the long A100 run uses.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import ModelConfig
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm with the reduction done in fp32 for mixed-precision safety."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
+        return (x * rms).to(dtype) * self.weight
+
+
+def build_rope_cache(head_dim: int, max_seq_len: int, theta: float,
+                     device=None, dtype=torch.float32):
+    """Precompute (cos, sin) of shape (max_seq_len, head_dim).
+
+    Half-rotation (Llama) convention: freqs are computed for head_dim/2 pairs
+    and concatenated with themselves so they line up with rotate_half.
+    """
+    i = torch.arange(0, head_dim, 2, device=device, dtype=torch.float32)
+    inv_freq = 1.0 / (theta ** (i / head_dim))             # (head_dim/2,)
+    t = torch.arange(max_seq_len, device=device, dtype=torch.float32)
+    freqs = torch.outer(t, inv_freq)                        # (T, head_dim/2)
+    emb = torch.cat([freqs, freqs], dim=-1)                 # (T, head_dim)
+    return emb.cos().to(dtype), emb.sin().to(dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    # x: (B, n_heads, T, head_dim); cos/sin: (T, head_dim) -> broadcast
+    cos = cos[None, None, :, :]
+    sin = sin[None, None, :, :]
+    return (x * cos) + (_rotate_half(x) * sin)
+
+
+class Attention(nn.Module):
+    """Causal grouped-query attention with optional QK-Norm and RoPE."""
+
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.n_heads = cfg.n_heads
+        self.n_kv_heads = cfg.n_kv_heads
+        self.head_dim = cfg.head_dim
+        self.n_rep = cfg.n_heads // cfg.n_kv_heads
+
+        self.wq = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(cfg.n_heads * self.head_dim, cfg.d_model, bias=False)
+
+        self.qk_norm = cfg.qk_norm
+        if cfg.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, cfg.norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, cfg.norm_eps)
+        self.dropout = cfg.dropout
+
+    def forward(self, x, cos, sin):
+        B, T, _ = x.shape
+        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
+        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
+        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        # (B, n_heads, T, head_dim)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+
+        # expand KV heads to match Q heads (GQA). repeat_interleave on head dim.
+        if self.n_rep > 1:
+            k = k.repeat_interleave(self.n_rep, dim=1)
+            v = v.repeat_interleave(self.n_rep, dim=1)
+
+        out = F.scaled_dot_product_attention(
+            q, k, v, is_causal=True,
+            dropout_p=self.dropout if self.training else 0.0,
+        )
+        out = out.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.wo(out)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        hidden = int(cfg.mlp_ratio * cfg.d_model)
+        m = cfg.mlp_multiple_of
+        hidden = ((hidden + m - 1) // m) * m
+        self.gate = nn.Linear(cfg.d_model, hidden, bias=False)
+        self.up = nn.Linear(cfg.d_model, hidden, bias=False)
+        self.down = nn.Linear(hidden, cfg.d_model, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.norm1 = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.attn = Attention(cfg)
+        self.norm2 = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.mlp = SwiGLU(cfg)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
+        self.norm_f = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        if cfg.tie_weights:
+            self.lm_head.weight = self.embed.weight
+
+        cos, sin = build_rope_cache(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+        self.apply(self._init_weights)
+        # Residual-projection scaling (GPT-2 trick): keep residual-stream growth
+        # bounded with depth by shrinking the layers that write back into it.
+        scale = 1.0 / math.sqrt(2 * cfg.n_layers)
+        for name, p in self.named_parameters():
+            if name.endswith("wo.weight") or name.endswith("down.weight"):
+                with torch.no_grad():
+                    p.mul_(scale)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.cfg.init_std)
+            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=self.cfg.init_std)
+
+    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
+        B, T = idx.shape
+        assert T <= self.cfg.max_seq_len, "sequence longer than rope cache"
+        x = self.embed(idx)
+        cos = self.rope_cos[:T]
+        sin = self.rope_sin[:T]
+        for block in self.blocks:
+            x = block(x, cos, sin)
+        x = self.norm_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=-1,
+            )
+        return logits, loss
+
+    def num_params(self, non_embedding: bool = True) -> int:
+        n = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            # tied: lm_head shares embed, so subtract the one embedding table
+            n -= self.embed.weight.numel()
+        return n
+```
+
+===== FILE: src/matilda/optim.py =====
+```python
+"""Optimizer + LR schedule construction.
+
+Two details that materially affect quality and that most tutorials get wrong:
+  1. No weight decay on 1-D params (norms, biases) or the embedding table.
+     Decaying norm/embedding weights quietly hurts.
+  2. The LR schedule includes linear warmup; starting at peak LR on random
+     weights diverges.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch.optim.lr_scheduler import LambdaLR
+
+
+def build_adamw(model, lr=3e-4, weight_decay=0.1,
+                betas=(0.9, 0.95), eps=1e-8) -> torch.optim.AdamW:
+    decay, no_decay = [], []
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        # 2-D matmul weights get decay; norms/biases (1-D) and embeddings don't.
+        if p.ndim < 2 or "embed" in name:
+            no_decay.append(p)
+        else:
+            decay.append(p)
+    groups = [
+        {"params": decay, "weight_decay": weight_decay},
+        {"params": no_decay, "weight_decay": 0.0},
+    ]
+    fused = torch.cuda.is_available()
+    return torch.optim.AdamW(groups, lr=lr, betas=betas, eps=eps, fused=fused)
+
+
+def cosine_warmup_scheduler(optimizer, warmup_steps, total_steps, min_lr_ratio=0.1):
+    """Linear warmup then cosine decay to min_lr_ratio * peak."""
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return (step + 1) / max(1, warmup_steps)
+        if step >= total_steps:
+            return min_lr_ratio
+        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+        return min_lr_ratio + (1 - min_lr_ratio) * cosine
+
+    return LambdaLR(optimizer, lr_lambda)
+```
+
+===== FILE: src/matilda/checkpoint.py =====
+```python
+"""Crash-safe checkpointing.
+
+A resumed run must continue *identically*, not just "without an obvious spike".
+That requires saving everything that influences the next step:
+  model, optimizer, scheduler, step, AND all RNG states AND the dataloader
+  position. Dropping the RNG or data position is the classic cause of a loss
+  blip on resume (you re-see data / re-sample dropout differently).
+
+Writes are atomic (write tmp -> os.replace) so an instance dying mid-save can
+never leave a half-written checkpoint that fails to load.
+"""
+
+from __future__ import annotations
+
+import os
+import glob
+import random
+from dataclasses import asdict
+
+import numpy as np
+import torch
+
+
+def _rng_state() -> dict:
+    state = {
+        "python": random.getstate(),
+        "numpy": np.random.get_state(),
+        "torch": torch.get_rng_state(),
+    }
+    if torch.cuda.is_available():
+        state["cuda"] = torch.cuda.get_rng_state_all()
+    return state
+
+
+def _set_rng_state(state: dict) -> None:
+    random.setstate(state["python"])
+    np.random.set_state(state["numpy"])
+    torch.set_rng_state(state["torch"])
+    if "cuda" in state and torch.cuda.is_available():
+        torch.cuda.set_rng_state_all(state["cuda"])
+
+
+def save_checkpoint(path, *, model, optimizer, scheduler, step,
+                    config=None, data_state=None) -> None:
+    """Atomically write a complete checkpoint to `path`."""
+    payload = {
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "scheduler": scheduler.state_dict() if scheduler is not None else None,
+        "step": step,
+        "rng": _rng_state(),
+        "data_state": data_state,
+        "config": asdict(config) if hasattr(config, "__dataclass_fields__") else config,
+    }
+    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+    tmp = f"{path}.tmp.{os.getpid()}"
+    torch.save(payload, tmp)
+    os.replace(tmp, path)  # atomic on POSIX and Windows
+
+
+def load_checkpoint(path, *, model, optimizer=None, scheduler=None,
+                    map_location="cpu", restore_rng=True) -> dict:
+    """Restore in-place. Returns the raw payload (for step, data_state, config)."""
+    ck = torch.load(path, map_location=map_location, weights_only=False)
+    model.load_state_dict(ck["model"])
+    if optimizer is not None and ck.get("optimizer") is not None:
+        optimizer.load_state_dict(ck["optimizer"])
+    if scheduler is not None and ck.get("scheduler") is not None:
+        scheduler.load_state_dict(ck["scheduler"])
+    if restore_rng and ck.get("rng") is not None:
+        _set_rng_state(ck["rng"])
+    return ck
+
+
+def latest_checkpoint(directory) -> str | None:
+    """Highest-step checkpoint matching ckpt_*.pt, or None."""
+    paths = glob.glob(os.path.join(directory, "ckpt_*.pt"))
+    if not paths:
+        return None
+    return max(paths, key=lambda p: int(os.path.basename(p)[5:-3]))
+
+
+def rotate_checkpoints(directory, keep_last: int, protect: set[str] | None = None) -> None:
+    """Delete oldest ckpt_*.pt beyond keep_last. `protect` = basenames to keep."""
+    protect = protect or set()
+    paths = sorted(
+        glob.glob(os.path.join(directory, "ckpt_*.pt")),
+        key=lambda p: int(os.path.basename(p)[5:-3]),
+    )
+    deletable = [p for p in paths if os.path.basename(p) not in protect]
+    for p in deletable[:-keep_last] if keep_last > 0 else deletable:
+        try:
+            os.remove(p)
+        except OSError:
+            pass
+```
+
+===== FILE: src/matilda/monitor.py =====
+```python
+"""Throughput / utilization observability.
+
+MFU (Model FLOPs Utilization) is the headline number a training engineer lives
+in: what fraction of the GPU's theoretical bf16 throughput you actually extract.
+We also track a rolling step-time window so a *degradation* mid-run (thermal
+throttling, a noisy Vast neighbour) is visible, not just the instantaneous rate.
+"""
+
+from __future__ import annotations
+
+import time
+from collections import deque
+
+# Peak bf16 *dense* TFLOPs (no 2:4 sparsity). Substring-matched on device name.
+PEAK_TFLOPS_BF16 = {
+    "A100": 312.0,
+    "H100": 494.0,
+    "H800": 494.0,
+    "4090": 165.0,
+    "3090": 71.0,
+    "3060": 51.0,
+    "T4": 65.0,    # fp16; T4 has no native bf16 (Turing) -> emulated, ignore MFU
+    "A10": 125.0,
+    "L4": 121.0,
+}
+
+
+def peak_tflops(device_name: str, default: float = 312.0) -> float:
+    for key, val in PEAK_TFLOPS_BF16.items():
+        if key in device_name:
+            return val
+    return default
+
+
+def mfu(n_params_active: int, tokens_per_step: int, dt_seconds: float,
+        peak_flops_per_sec: float) -> float:
+    """6N flops/token (forward+backward, matmul-dominated approximation)."""
+    if dt_seconds <= 0:
+        return 0.0
+    achieved = 6 * n_params_active * tokens_per_step / dt_seconds
+    return achieved / peak_flops_per_sec
+
+
+class Throughput:
+    """Rolling step-time tracker; reports tokens/sec and MFU."""
+
+    def __init__(self, n_params_active, tokens_per_step, peak_flops_per_sec,
+                 window=50):
+        self.n = n_params_active
+        self.tps = tokens_per_step
+        self.peak = peak_flops_per_sec
+        self.times = deque(maxlen=window)
+        self._t0 = None
+
+    def tick(self) -> dict | None:
+        now = time.perf_counter()
+        if self._t0 is None:
+            self._t0 = now
+            return None
+        dt = now - self._t0
+        self._t0 = now
+        self.times.append(dt)
+        avg = sum(self.times) / len(self.times)
+        return {
+            "dt_s": dt,
+            "dt_avg_s": avg,
+            "tokens_per_s": self.tps / dt,
+            "mfu": mfu(self.n, self.tps, dt, self.peak),
+            "mfu_avg": mfu(self.n, self.tps, avg, self.peak),
+        }
+```
+
+===== FILE: src/matilda/data.py =====
+```python
+"""Data streams + tokenized-shard tooling.
+
+All streams expose the same interface so the training loop never changes when
+we swap synthetic data for real tokenized FineWeb shards:
+
+    x, y = stream.next()            # (B, T) input, (B, T) targets
+    state = stream.state_dict()     # resumable position / RNG
+    stream.load_state_dict(state)
+
+On-disk format: token ids as a flat `uint16` `.bin` (valid for vocab <= 65535,
+which GPT-2's 50257 satisfies). A `manifest.json` records per-shard token counts
+and SHA-256 so the prepared data is verifiable and reproducible.
+"""
+
+from __future__ import annotations
+
+import os
+import json
+import glob
+import hashlib
+
+import numpy as np
+import torch
+
+DTYPE = np.uint16
+
+
+# --------------------------------------------------------------------------
+# shard writing / verification (used by scripts/prepare_data.py, unit-tested)
+# --------------------------------------------------------------------------
+def sha256_file(path, chunk=1 << 20) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for block in iter(lambda: f.read(chunk), b""):
+            h.update(block)
+    return h.hexdigest()
+
+
+class ShardWriter:
+    """Accumulate token ids and flush fixed-size `.bin` shards + a manifest."""
+
+    def __init__(self, out_dir, shard_tokens=100_000_000, prefix="shard"):
+        self.out_dir = out_dir
+        self.shard_tokens = shard_tokens
+        self.prefix = prefix
+        os.makedirs(out_dir, exist_ok=True)
+        self._buf: list[np.ndarray] = []
+        self._buf_len = 0
+        self._shard_idx = 0
+        self.total_tokens = 0
+        self.manifest_entries: list[dict] = []
+
+    def add(self, tokens) -> None:
+        arr = np.asarray(tokens, dtype=DTYPE)
+        self._buf.append(arr)
+        self._buf_len += arr.size
+        self.total_tokens += arr.size
+        while self._buf_len >= self.shard_tokens:
+            self._flush(self.shard_tokens)
+
+    def _flush(self, n) -> None:
+        flat = np.concatenate(self._buf) if len(self._buf) > 1 else self._buf[0]
+        out, rest = flat[:n], flat[n:]
+        path = os.path.join(self.out_dir, f"{self.prefix}_{self._shard_idx:05d}.bin")
+        out.tofile(path)
+        self.manifest_entries.append({
+            "file": os.path.basename(path),
+            "tokens": int(out.size),
+            "sha256": sha256_file(path),
+        })
+        self._shard_idx += 1
+        self._buf = [rest] if rest.size else []
+        self._buf_len = rest.size
+
+    def close(self, meta: dict | None = None) -> dict:
+        if self._buf_len > 0:
+            self._flush(self._buf_len)        # final partial shard
+        manifest = {
+            "total_tokens": int(self.total_tokens),
+            "shards": self.manifest_entries,
+            **(meta or {}),
+        }
+        with open(os.path.join(self.out_dir, "manifest.json"), "w") as f:
+            json.dump(manifest, f, indent=2)
+        return manifest
+
+
+def verify_manifest(out_dir) -> bool:
+    """Re-checksum every shard against the manifest. Raises on mismatch."""
+    with open(os.path.join(out_dir, "manifest.json")) as f:
+        manifest = json.load(f)
+    for entry in manifest["shards"]:
+        path = os.path.join(out_dir, entry["file"])
+        actual = sha256_file(path)
+        if actual != entry["sha256"]:
+            raise ValueError(f"checksum mismatch for {entry['file']}")
+        if os.path.getsize(path) != entry["tokens"] * 2:  # uint16 = 2 bytes
+            raise ValueError(f"size mismatch for {entry['file']}")
+    return True
+
+
+def shard_paths(out_dir) -> list[str]:
+    with open(os.path.join(out_dir, "manifest.json")) as f:
+        manifest = json.load(f)
+    return [os.path.join(out_dir, e["file"]) for e in manifest["shards"]]
+
+
+# --------------------------------------------------------------------------
+# streams
+# --------------------------------------------------------------------------
+class SyntheticStream:
+    """Deterministic random tokens. For dev/CI and loop testing on CPU."""
+
+    def __init__(self, vocab_size, batch_size, seq_len, seed=0, device="cpu"):
+        self.vocab_size = vocab_size
+        self.B = batch_size
+        self.T = seq_len
+        self.device = device
+        self.gen = torch.Generator().manual_seed(seed)
+        self.pos = 0
+
+    def next(self):
+        seq = torch.randint(0, self.vocab_size, (self.B, self.T + 1),
+                            generator=self.gen)
+        x = seq[:, :-1].contiguous().to(self.device)
+        y = seq[:, 1:].contiguous().to(self.device)
+        self.pos += 1
+        return x, y
+
+    def state_dict(self):
+        return {"pos": self.pos, "gen": self.gen.get_state()}
+
+    def load_state_dict(self, s):
+        self.pos = s["pos"]
+        self.gen.set_state(s["gen"])
+
+
+class BinStream:
+    """Packed windows sampled from memory-mapped uint16 token shards.
+
+    Each batch element is a random T+1 window from a shard chosen with
+    probability proportional to its size; x/y are the next-token shift. Random
+    sampling (rather than a sequential cursor) means epochs are implicit and
+    resume only needs the RNG state.
+    """
+
+    def __init__(self, bin_paths, batch_size, seq_len, seed=0, device="cpu"):
+        assert bin_paths, "no shards provided"
+        self.paths = list(bin_paths)
+        self.B = batch_size
+        self.T = seq_len
+        self.device = device
+        self.arrays = [np.memmap(p, dtype=DTYPE, mode="r") for p in self.paths]
+        self.sizes = torch.tensor([float(a.size) for a in self.arrays])
+        assert (self.sizes > seq_len + 1).all(), "a shard is smaller than seq_len+1"
+        self.weights = self.sizes / self.sizes.sum()
+        self.gen = torch.Generator().manual_seed(seed)
+        self.pos = 0
+
+    def next(self):
+        shard_ids = torch.multinomial(self.weights, self.B, replacement=True,
+                                      generator=self.gen)
+        xb = np.empty((self.B, self.T), dtype=np.int64)
+        yb = np.empty((self.B, self.T), dtype=np.int64)
+        for i, sid in enumerate(shard_ids.tolist()):
+            arr = self.arrays[sid]
+            hi = arr.size - (self.T + 1)
+            start = int(torch.randint(0, hi + 1, (1,), generator=self.gen).item())
+            window = arr[start:start + self.T + 1].astype(np.int64)
+            xb[i] = window[:-1]
+            yb[i] = window[1:]
+        self.pos += 1
+        x = torch.from_numpy(xb).to(self.device)
+        y = torch.from_numpy(yb).to(self.device)
+        return x, y
+
+    def state_dict(self):
+        return {"pos": self.pos, "gen": self.gen.get_state()}
+
+    def load_state_dict(self, s):
+        self.pos = s["pos"]
+        self.gen.set_state(s["gen"])
+```
+
+===== FILE: src/matilda/train.py =====
+```python
+"""Training loop with the reliability guards a long run needs.
+
+Designed so an overnight Vast run survives the things that actually kill runs:
+  - one bad batch (NaN/Inf loss or grad) -> skip it, log loudly, keep going
+  - spot-instance SIGTERM -> checkpoint before exiting
+  - process death -> resume bit-for-bit from latest checkpoint on restart
+
+Throughput (MFU, tokens/s, step-time) is logged every `log_every` steps. The
+data source is any object with .next()/.state_dict()/.load_state_dict() (see
+data.py), so swapping synthetic data for real FineWeb shards changes nothing.
+"""
+
+from __future__ import annotations
+
+import os
+import math
+import time
+import signal
+from dataclasses import dataclass, field
+
+import torch
+
+from .model import Transformer
+from .config import ModelConfig
+from .optim import build_adamw, cosine_warmup_scheduler
+from .monitor import Throughput, peak_tflops
+from .checkpoint import (
+    save_checkpoint, load_checkpoint, latest_checkpoint, rotate_checkpoints,
+)
+
+
+@dataclass
+class TrainConfig:
+    total_steps: int = 1000
+    warmup_steps: int = 100
+    grad_accum: int = 1
+    lr: float = 3e-4
+    weight_decay: float = 0.1
+    grad_clip: float = 1.0
+    min_lr_ratio: float = 0.1
+
+    batch_size: int = 8
+    seq_len: int = 1024
+
+    log_every: int = 10
+    ckpt_every: int = 500
+    keep_last: int = 3
+    ckpt_dir: str = "checkpoints"
+
+    device: str = "cuda"
+    dtype: str = "bfloat16"            # bf16 on Ampere+, fp32 fallback on CPU
+    compile: bool = False
+    seed: int = 1234
+    max_skips: int = 20               # abort if too many bad batches in a row
+    wandb_project: str | None = None
+
+
+class Trainer:
+    def __init__(self, model_cfg: ModelConfig, train_cfg: TrainConfig, stream):
+        self.mcfg = model_cfg
+        self.cfg = train_cfg
+        self.stream = stream
+        self.step = 0
+        self.consecutive_skips = 0
+        self._interrupted = False
+
+        torch.manual_seed(train_cfg.seed)
+        self.device = torch.device(
+            train_cfg.device if torch.cuda.is_available()
+            or train_cfg.device == "cpu" else "cpu")
+
+        self.model = Transformer(model_cfg).to(self.device)
+        if train_cfg.compile:
+            self.model = torch.compile(self.model)
+        self.opt = build_adamw(self.model, lr=train_cfg.lr,
+                               weight_decay=train_cfg.weight_decay)
+        self.sched = cosine_warmup_scheduler(
+            self.opt, train_cfg.warmup_steps, train_cfg.total_steps,
+            train_cfg.min_lr_ratio)
+
+        # bf16 only where supported; CPU/T4 fall back to fp32 for correctness.
+        want_bf16 = (train_cfg.dtype == "bfloat16"
+                     and self.device.type == "cuda"
+                     and torch.cuda.is_bf16_supported())
+        self.amp_dtype = torch.bfloat16 if want_bf16 else None
+
+        tokens_per_step = (train_cfg.batch_size * train_cfg.seq_len
+                           * train_cfg.grad_accum)
+        dev_name = (torch.cuda.get_device_name(self.device)
+                    if self.device.type == "cuda" else "cpu")
+        self.monitor = Throughput(
+            n_params_active=self._active_params(),
+            tokens_per_step=tokens_per_step,
+            peak_flops_per_sec=peak_tflops(dev_name) * 1e12,
+        )
+        self.wandb = self._init_wandb()
+        self._install_signal_handlers()
+
+    # --- helpers -----------------------------------------------------------
+    def _active_params(self):
+        m = self.model._orig_mod if hasattr(self.model, "_orig_mod") else self.model
+        return m.num_params(non_embedding=True)
+
+    def _init_wandb(self):
+        if not self.cfg.wandb_project:
+            return None
+        try:
+            import wandb
+            wandb.init(project=self.cfg.wandb_project,
+                       config={**self.mcfg.to_dict(), **vars(self.cfg)})
+            return wandb
+        except Exception as e:
+            print(f"[warn] wandb disabled: {e}")
+            return None
+
+    def _install_signal_handlers(self):
+        def handler(signum, frame):
+            print(f"[signal] {signum} received -> checkpoint and exit")
+            self._interrupted = True
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            try:
+                signal.signal(sig, handler)
+            except (ValueError, OSError):
+                pass  # not in main thread (e.g. under pytest) -> skip
+
+    def _autocast(self):
+        if self.amp_dtype is not None:
+            return torch.autocast(device_type="cuda", dtype=self.amp_dtype)
+        return torch.autocast(device_type="cpu", enabled=False)
+
+    # --- checkpoint --------------------------------------------------------
+    def _ckpt_path(self):
+        return os.path.join(self.cfg.ckpt_dir, f"ckpt_{self.step}.pt")
+
+    def save(self):
+        save_checkpoint(self._ckpt_path(), model=self.model, optimizer=self.opt,
+                        scheduler=self.sched, step=self.step, config=self.mcfg,
+                        data_state=self.stream.state_dict())
+        rotate_checkpoints(self.cfg.ckpt_dir, self.cfg.keep_last)
+
+    def maybe_resume(self):
+        path = latest_checkpoint(self.cfg.ckpt_dir)
+        if path is None:
+            return False
+        ck = load_checkpoint(path, model=self.model, optimizer=self.opt,
+                             scheduler=self.sched,
+                             map_location=self.device)
+        self.step = ck["step"]
+        if ck.get("data_state") is not None:
+            self.stream.load_state_dict(ck["data_state"])
+        print(f"[resume] from {path} at step {self.step}")
+        return True
+
+    # --- one optimizer step (grad accum + guards) --------------------------
+    def _step(self):
+        self.opt.zero_grad(set_to_none=True)
+        total_loss = 0.0
+        for micro in range(self.cfg.grad_accum):
+            x, y = self.stream.next()
+            x, y = x.to(self.device), y.to(self.device)
+            sync = (micro == self.cfg.grad_accum - 1)
+            ctx = (self.model.no_sync()
+                   if (not sync and hasattr(self.model, "no_sync"))
+                   else _nullcontext())
+            with ctx, self._autocast():
+                _, loss = self.model(x, y)
+                loss = loss / self.cfg.grad_accum
+            if not torch.isfinite(loss):
+                return None  # bad micro-batch -> abort this step
+            loss.backward()
+            total_loss += loss.item()
+
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            self.model.parameters(), self.cfg.grad_clip)
+        if not torch.isfinite(grad_norm):
+            return None  # non-finite grad -> skip update
+
+        self.opt.step()
+        self.sched.step()
+        return total_loss, grad_norm.item()
+
+    # --- main loop ---------------------------------------------------------
+    def train(self):
+        os.makedirs(self.cfg.ckpt_dir, exist_ok=True)
+        self.maybe_resume()
+        self.model.train()
+        while self.step < self.cfg.total_steps:
+            result = self._step()
+            if result is None:
+                self.consecutive_skips += 1
+                print(f"[nan-guard] step {self.step} skipped "
+                      f"({self.consecutive_skips}/{self.cfg.max_skips})")
+                self.opt.zero_grad(set_to_none=True)
+                if self.consecutive_skips >= self.cfg.max_skips:
+                    raise RuntimeError("too many non-finite steps; aborting run")
+                continue
+            self.consecutive_skips = 0
+            loss, grad_norm = result
+            self.step += 1
+
+            stats = self.monitor.tick()
+            if self.step % self.cfg.log_every == 0:
+                self._log(loss, grad_norm, stats)
+            if self.step % self.cfg.ckpt_every == 0:
+                self.save()
+            if self._interrupted:
+                self.save()
+                print(f"[exit] checkpointed at step {self.step}")
+                break
+        else:
+            self.save()  # final checkpoint on clean completion
+        return self.step
+
+    def _log(self, loss, grad_norm, stats):
+        lr = self.sched.get_last_lr()[0]
+        msg = (f"step {self.step:>6} | loss {loss:7.4f} | lr {lr:.2e} "
+               f"| gnorm {grad_norm:5.2f}")
+        if stats:
+            msg += (f" | {stats['tokens_per_s']:8.0f} tok/s "
+                    f"| mfu {stats['mfu_avg']*100:4.1f}%")
+        print(msg)
+        if self.wandb:
+            row = {"loss": loss, "lr": lr, "grad_norm": grad_norm}
+            if stats:
+                row.update({"tokens_per_s": stats["tokens_per_s"],
+                            "mfu": stats["mfu_avg"],
+                            "step_time_s": stats["dt_avg_s"]})
+            self.wandb.log(row, step=self.step)
+
+
+class _nullcontext:
+    def __enter__(self):
+        return None
+
+    def __exit__(self, *exc):
+        return False
+```
+
+===== FILE: scripts/prepare_data.py =====
+```python
+"""Tokenize a HuggingFace text dataset into verifiable uint16 shards.
+
+Default: FineWeb-Edu sample-10BT (best small-scale benchmark movement, max
+comparability to Pythia/SmolLM). Dataset-agnostic by design so DCLM /
+Nemotron-CC become drop-in data-ablation rows:
+
+    python scripts/prepare_data.py --target-tokens 3_000_000_000 \
+        --dataset HuggingFaceFW/fineweb-edu --name sample-10BT --out-dir data/fwedu
+
+    # ablation rows (same flags, different source):
+    --dataset mlfoundations/dclm-baseline-1.0 --out-dir data/dclm
+    --dataset nvidia/Nemotron-CC            --out-dir data/nemotron
+
+Streams the source (no full download), so disk holds only the tokenized output.
+"""
+
+from __future__ import annotations
+
+import sys
+import argparse
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+from matilda.data import ShardWriter, verify_manifest  # noqa: E402
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu")
+    ap.add_argument("--name", default="sample-10BT")
+    ap.add_argument("--split", default="train")
+    ap.add_argument("--text-key", default="text")
+    ap.add_argument("--tokenizer", default="gpt2")
+    ap.add_argument("--target-tokens", type=int, default=3_000_000_000)
+    ap.add_argument("--shard-tokens", type=int, default=100_000_000)
+    ap.add_argument("--out-dir", default="data/fwedu")
+    args = ap.parse_args()
+
+    import tiktoken
+    from datasets import load_dataset
+
+    enc = tiktoken.get_encoding(args.tokenizer)
+    eot = enc.eot_token
+    assert enc.n_vocab <= 65535, "vocab > uint16; use a smaller tokenizer"
+
+    ds = load_dataset(args.dataset, name=args.name, split=args.split,
+                      streaming=True)
+    writer = ShardWriter(args.out_dir, shard_tokens=args.shard_tokens)
+
+    n_docs = 0
+    for doc in ds:
+        ids = enc.encode_ordinary(doc[args.text_key])
+        ids.append(eot)                       # document boundary
+        writer.add(ids)
+        n_docs += 1
+        if n_docs % 1000 == 0:
+            print(f"\rdocs={n_docs:,} tokens={writer.total_tokens:,}", end="")
+        if writer.total_tokens >= args.target_tokens:
+            break
+
+    manifest = writer.close(meta={
+        "dataset": args.dataset, "name": args.name, "split": args.split,
+        "tokenizer": args.tokenizer, "eot_token": eot, "n_docs": n_docs,
+    })
+    print(f"\nwrote {manifest['total_tokens']:,} tokens in "
+          f"{len(manifest['shards'])} shards -> {args.out_dir}")
+    verify_manifest(args.out_dir)
+    print("manifest verified (checksums + sizes OK)")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+===== FILE: tests/test_model.py =====
+```python
+"""Sanity tests that gate the paid run. All must be green on Colab first.
+
+  - test_forward_shapes / test_weight_tying / test_param_count: wiring is correct
+  - test_causal_mask: no information leaks from future tokens (the bug that
+    silently inflates eval and is invisible in the loss curve)
+  - test_overfit_single_batch: the model can actually learn (loss -> ~0 on one
+    fixed batch). The cheapest, highest-signal correctness check in ML.
+"""
+
+import torch
+import pytest
+
+from matilda import Transformer, ModelConfig, DEV_TINY
+
+
+def _tiny():
+    return Transformer(DEV_TINY).eval()
+
+
+def test_forward_shapes():
+    model = _tiny()
+    B, T = 2, 16
+    idx = torch.randint(0, DEV_TINY.vocab_size, (B, T))
+    logits, loss = model(idx)
+    assert logits.shape == (B, T, DEV_TINY.vocab_size)
+    assert loss is None
+    targets = torch.randint(0, DEV_TINY.vocab_size, (B, T))
+    _, loss = model(idx, targets)
+    assert loss is not None and loss.ndim == 0
+
+
+def test_weight_tying():
+    model = _tiny()
+    assert model.lm_head.weight.data_ptr() == model.embed.weight.data_ptr()
+
+
+def test_loss_at_init_is_near_uniform():
+    # untrained model should be ~ -log(1/V) = log(V)
+    model = _tiny()
+    idx = torch.randint(0, DEV_TINY.vocab_size, (4, 32))
+    tgt = torch.randint(0, DEV_TINY.vocab_size, (4, 32))
+    _, loss = model(idx, tgt)
+    expected = torch.log(torch.tensor(float(DEV_TINY.vocab_size)))
+    assert abs(loss.item() - expected.item()) < 1.0
+
+
+def test_causal_mask_no_future_leak():
+    """Changing token at position t must not alter logits at positions < t."""
+    model = _tiny()
+    torch.manual_seed(0)
+    idx = torch.randint(0, DEV_TINY.vocab_size, (1, 24))
+    with torch.no_grad():
+        base, _ = model(idx)
+        idx2 = idx.clone()
+        idx2[0, -1] = (idx2[0, -1] + 1) % DEV_TINY.vocab_size  # perturb last token
+        perturbed, _ = model(idx2)
+    # all positions except the last must be identical
+    assert torch.allclose(base[:, :-1], perturbed[:, :-1], atol=1e-5)
+    assert not torch.allclose(base[:, -1], perturbed[:, -1], atol=1e-5)
+
+
+def test_gqa_kv_head_counts():
+    model = _tiny()
+    attn = model.blocks[0].attn
+    assert attn.wk.out_features == DEV_TINY.n_kv_heads * DEV_TINY.head_dim
+    assert attn.wq.out_features == DEV_TINY.n_heads * DEV_TINY.head_dim
+
+
+@pytest.mark.slow
+def test_overfit_single_batch():
+    """The model must drive loss toward zero on one fixed batch."""
+    cfg = ModelConfig(vocab_size=256, max_seq_len=64, d_model=128,
+                      n_layers=2, n_heads=4, n_kv_heads=2)
+    model = Transformer(cfg).train()
+    torch.manual_seed(0)
+    idx = torch.randint(0, cfg.vocab_size, (4, 32))
+    tgt = torch.randint(0, cfg.vocab_size, (4, 32))
+    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
+    losses = []
+    for _ in range(300):
+        _, loss = model(idx, tgt)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step()
+        losses.append(loss.item())
+    assert losses[-1] < 0.1, f"failed to overfit; final loss={losses[-1]:.3f}"
+```
+
+===== FILE: tests/test_checkpoint.py =====
+```python
+"""Resume must continue *identically* to an uninterrupted run.
+
+Strategy: run a short reference training. Run it again but checkpoint halfway,
+throw the objects away, build *fresh* model/optimizer/scheduler/data-stream,
+restore from the checkpoint, and continue. The post-resume per-step losses must
+match the reference run's second half exactly. If any piece of state is missing
+(opt moments, scheduler, RNG, data position) the curves diverge and this fails.
+"""
+
+import os
+import random
+
+import numpy as np
+import torch
+
+from matilda import Transformer, ModelConfig
+from matilda.optim import build_adamw, cosine_warmup_scheduler
+from matilda.checkpoint import (
+    save_checkpoint, load_checkpoint, latest_checkpoint, rotate_checkpoints,
+)
+
+CFG = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64,
+                  n_layers=2, n_heads=4, n_kv_heads=2)
+TOTAL, WARMUP, HALF = 12, 2, 6
+
+
+class SyntheticStream:
+    """Resumable deterministic batch source (own generator => independent RNG)."""
+
+    def __init__(self, seed, B=4, T=16):
+        self.B, self.T = B, T
+        self.gen = torch.Generator().manual_seed(seed)
+        self.pos = 0
+
+    def next(self):
+        x = torch.randint(0, CFG.vocab_size, (self.B, self.T), generator=self.gen)
+        y = torch.randint(0, CFG.vocab_size, (self.B, self.T), generator=self.gen)
+        self.pos += 1
+        return x, y
+
+    def state_dict(self):
+        return {"pos": self.pos, "gen": self.gen.get_state()}
+
+    def load_state_dict(self, s):
+        self.pos = s["pos"]
+        self.gen.set_state(s["gen"])
+
+
+def _seed_all(seed=1234):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def _build():
+    model = Transformer(CFG).train()
+    opt = build_adamw(model, lr=1e-3)
+    sched = cosine_warmup_scheduler(opt, WARMUP, TOTAL)
+    return model, opt, sched
+
+
+def _train_steps(model, opt, sched, stream, n):
+    losses = []
+    for _ in range(n):
+        x, y = stream.next()
+        _, loss = model(x, y)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step()
+        sched.step()
+        losses.append(loss.item())
+    return losses
+
+
+def test_resume_is_bit_for_bit(tmp_path):
+    # --- reference: uninterrupted run ---
+    _seed_all()
+    m, o, s = _build()
+    ref_stream = SyntheticStream(seed=42)
+    ref_losses = _train_steps(m, o, s, ref_stream, TOTAL)
+
+    # --- interrupted run: train HALF, checkpoint, drop everything ---
+    _seed_all()
+    m, o, s = _build()
+    stream = SyntheticStream(seed=42)
+    first_half = _train_steps(m, o, s, stream, HALF)
+    assert first_half == ref_losses[:HALF], "training is not deterministic"
+
+    ckpt = os.path.join(tmp_path, f"ckpt_{HALF}.pt")
+    save_checkpoint(ckpt, model=m, optimizer=o, scheduler=s, step=HALF,
+                    config=CFG, data_state=stream.state_dict())
+
+    # --- fresh objects, restore, continue ---
+    m2, o2, s2 = _build()                    # fresh random init + zeroed moments
+    stream2 = SyntheticStream(seed=999)      # deliberately wrong seed...
+    ck = load_checkpoint(ckpt, model=m2, optimizer=o2, scheduler=s2)
+    stream2.load_state_dict(ck["data_state"])  # ...corrected by restore
+    assert ck["step"] == HALF
+
+    resumed = _train_steps(m2, o2, s2, stream2, TOTAL - HALF)
+
+    for i, (a, b) in enumerate(zip(resumed, ref_losses[HALF:])):
+        assert abs(a - b) < 1e-6, f"resume diverged at step {HALF+i}: {a} vs {b}"
+
+
+def test_latest_and_rotation(tmp_path):
+    for step in (10, 20, 30, 40):
+        p = os.path.join(tmp_path, f"ckpt_{step}.pt")
+        torch.save({"step": step}, p)
+    assert latest_checkpoint(tmp_path).endswith("ckpt_40.pt")
+
+    rotate_checkpoints(tmp_path, keep_last=2, protect={"ckpt_10.pt"})
+    remaining = sorted(os.path.basename(p) for p in
+                       __import__("glob").glob(os.path.join(tmp_path, "ckpt_*.pt")))
+    # keep last 2 (30, 40) + protected 10; drop 20
+    assert remaining == ["ckpt_10.pt", "ckpt_30.pt", "ckpt_40.pt"]
+```
+
+===== FILE: tests/test_train.py =====
+```python
+"""Training-loop reliability guards (all CPU-testable)."""
+
+import os
+import glob
+
+import torch
+import pytest
+
+from matilda import ModelConfig
+from matilda.data import SyntheticStream
+from matilda.train import Trainer, TrainConfig
+from matilda.monitor import mfu, peak_tflops
+
+MCFG = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64,
+                  n_layers=2, n_heads=4, n_kv_heads=2)
+
+
+def _make(tmp_path, **over):
+    kw = dict(total_steps=20, warmup_steps=2, batch_size=4, seq_len=16,
+              log_every=5, ckpt_every=10, keep_last=2, device="cpu",
+              dtype="float32", ckpt_dir=str(tmp_path))
+    kw.update(over)
+    tc = TrainConfig(**kw)
+    stream = SyntheticStream(MCFG.vocab_size, tc.batch_size, tc.seq_len, seed=0)
+    return Trainer(MCFG, tc, stream)
+
+
+def test_loop_runs_and_checkpoints(tmp_path):
+    t = _make(tmp_path)
+    final = t.train()
+    assert final == 20
+    # final checkpoint on clean completion + rotation kept <= keep_last
+    ckpts = glob.glob(os.path.join(tmp_path, "ckpt_*.pt"))
+    assert len(ckpts) <= 2
+    assert os.path.exists(os.path.join(tmp_path, "ckpt_20.pt"))
+
+
+def test_loop_resume_continues(tmp_path):
+    # run only 10 steps, leaving a checkpoint at step 10
+    t1 = _make(tmp_path, total_steps=10, ckpt_every=10)
+    assert t1.train() == 10
+    # a fresh trainer in the same dir must resume from step 10, not restart
+    t2 = _make(tmp_path, total_steps=15, ckpt_every=10)
+    assert t2.maybe_resume() is True
+    assert t2.step == 10
+    assert t2.train() == 15
+
+
+def test_nan_guard_aborts_after_max_skips(tmp_path):
+    t = _make(tmp_path, total_steps=50, max_skips=3)
+    nan = torch.tensor(float("nan"))
+    t.model.forward = lambda x, targets=None: (None, nan)  # shadow bound method
+    with pytest.raises(RuntimeError, match="non-finite"):
+        t.train()
+    assert t.step == 0  # never advanced past a bad batch
+
+
+def test_nan_guard_skips_then_recovers(tmp_path):
+    t = _make(tmp_path, total_steps=5, max_skips=10)
+    real_forward = t.model.forward
+    calls = {"n": 0}
+
+    def flaky(x, targets=None):
+        calls["n"] += 1
+        if calls["n"] <= 2:                     # first two micro-batches are bad
+            return None, torch.tensor(float("inf"))
+        return real_forward(x, targets)
+
+    t.model.forward = flaky
+    assert t.train() == 5                        # recovered and finished
+    assert t.consecutive_skips == 0
+
+
+def test_mfu_sanity():
+    # 100M params, 500k tokens/step, 2.0s, A100 -> ~0.48 MFU
+    val = mfu(100_000_000, 500_000, 2.0, peak_tflops("A100") * 1e12)
+    assert 0.0 < val < 1.0
+    assert peak_tflops("A100") == 312.0
+    assert peak_tflops("totally-unknown-gpu") == 312.0  # default
+```
+
+===== FILE: tests/test_data.py =====
+```python
+"""Data pipeline: shard writing/verification + BinStream correctness."""
+
+import os
+import json
+
+import numpy as np
+import pytest
+import torch
+
+from matilda.data import (
+    ShardWriter, verify_manifest, shard_paths, BinStream, DTYPE,
+)
+
+
+def test_shardwriter_roundtrip_and_manifest(tmp_path):
+    w = ShardWriter(str(tmp_path), shard_tokens=100)
+    # 250 tokens -> two full shards (100) + one partial (50)
+    w.add(list(range(0, 120)))
+    w.add(list(range(120, 250)))
+    manifest = w.close(meta={"tokenizer": "gpt2"})
+
+    assert manifest["total_tokens"] == 250
+    assert [s["tokens"] for s in manifest["shards"]] == [100, 100, 50]
+    assert manifest["tokenizer"] == "gpt2"
+    assert verify_manifest(str(tmp_path)) is True
+
+    # reconstructed token stream matches the input exactly
+    rebuilt = np.concatenate(
+        [np.fromfile(p, dtype=DTYPE) for p in shard_paths(str(tmp_path))])
+    assert rebuilt.tolist() == list(range(250))
+
+
+def test_verify_detects_corruption(tmp_path):
+    w = ShardWriter(str(tmp_path), shard_tokens=100)
+    w.add(list(range(150)))
+    w.close()
+    # corrupt the first shard's bytes without changing its size
+    p = shard_paths(str(tmp_path))[0]
+    data = bytearray(open(p, "rb").read())
+    data[0] ^= 0xFF
+    open(p, "wb").write(data)
+    with pytest.raises(ValueError, match="checksum mismatch"):
+        verify_manifest(str(tmp_path))
+
+
+def _write_ramp_bin(path, n):
+    # values 0..n-1 so a correct next-token window always has y == x + 1
+    np.arange(n, dtype=DTYPE).tofile(path)
+
+
+def test_binstream_shift_is_next_token(tmp_path):
+    p = os.path.join(tmp_path, "ramp.bin")
+    _write_ramp_bin(p, 5000)
+    s = BinStream([p], batch_size=8, seq_len=32, seed=0)
+    x, y = s.next()
+    assert x.shape == (8, 32) and y.shape == (8, 32)
+    assert torch.equal(y, x + 1)              # packed next-token target
+
+
+def test_binstream_resume_is_deterministic(tmp_path):
+    p = os.path.join(tmp_path, "ramp.bin")
+    _write_ramp_bin(p, 5000)
+    s = BinStream([p], batch_size=4, seq_len=16, seed=7)
+    snap = s.state_dict()
+    x1, y1 = s.next()
+    s.load_state_dict(snap)
+    x2, y2 = s.next()
+    assert torch.equal(x1, x2) and torch.equal(y1, y2)
+
+
+def test_binstream_multishard_weighting(tmp_path):
+    p1 = os.path.join(tmp_path, "a.bin")
+    p2 = os.path.join(tmp_path, "b.bin")
+    _write_ramp_bin(p1, 2000)
+    _write_ramp_bin(p2, 2000)
+    s = BinStream([p1, p2], batch_size=4, seq_len=8, seed=1)
+    x, y = s.next()
+    assert torch.equal(y, x + 1)              # still correct across shards
+```

docs/DEEPSEEK_REVIEW_PROMPT.md

@@ -0,0 +1,144 @@
+# Review Prompt for DeepSeek
+
+> Paste everything below into DeepSeek. DeepSeek cannot read your disk, so after
+> the prompt, paste the contents of the files it asks for (or attach them). The
+> file map and paths are included so you both share a mental model.
+
+---
+
+## ROLE
+
+You are a senior LLM training-infrastructure engineer reviewing a from-scratch
+small language model project. Be rigorous and skeptical. I want correctness bugs,
+reliability gaps, and anything that would embarrass me in a technical interview —
+not encouragement. Prioritize findings by severity (CRITICAL / HIGH / MEDIUM / LOW).
+
+## WHY THIS PROJECT EXISTS (the job I'm applying for)
+
+I'm building this as a portfolio piece for an **AI Training Infrastructure Engineer**
+role at **Maincode** (Melbourne, Australia). Maincode trains "Matilda", the first
+LLM built and trained from scratch in Australia. Key facts about the role:
+
+- It is an **infrastructure** role, NOT a research/architecture role. The work is:
+  distributed training pipelines, data ingestion/preprocessing, experiment
+  management, checkpointing, reproducibility, monitoring/observability, debugging
+  long-running (days–weeks) training jobs, optimizing throughput (compute/memory/
+  data), and diagnosing failures that appear hours into a run.
+- Explicitly NOT about: wrapping external APIs, prompt engineering, user-facing apps.
+- Stack: Python, PyTorch/JAX, reliable infra for large compute, debugging
+  distributed systems and long jobs.
+- They value engineers who understand how systems behave over long runtimes and
+  who prefer understanding internals over relying on abstraction.
+
+So the project is **deliberately weighted toward operational excellence**
+(checkpoint/resume, fault tolerance, reproducibility, MFU/observability) over
+architectural breadth. The architecture is modern but standard; the infra is the
+flex. Please judge it through that lens.
+
+## ORIGINAL PLAN & KEY DECISIONS (locked)
+
+- **Goal:** train a sub-200M-param modern transformer from scratch, end-to-end,
+  with a clean ablation table, on a tight budget.
+- **Hardware:** rent ONE A100 80GB on Vast.ai for the single long paid run.
+- **Long paid run = modernized DENSE ~114M, maximize MFU.** Rationale: MoE on a
+  single GPU collapses MFU to ~20% (needs expert parallelism across devices);
+  dense hits 45–50%. "DeepSeek-inspired" here means the efficiency stack
+  (Flash/SDPA, torch.compile, fused optimizer, bf16, Muon, μP, WSD), NOT sparsity.
+  A100 has no FP8 (bf16 ceiling; FP8 would need H100).
+- **MoE is built only as a CHEAP SHORT ablation** (~500M tokens, ~$0.60) with an
+  honest write-up of the single-GPU MFU drop. Not the headline run.
+- **Tokenizer:** GPT-2 50257 (tiktoken "gpt2") — fits uint16, comparable to Pythia
+  baselines. (We caught and avoided a bug from our source primer that paired
+  cl100k ~100k vocab with uint16, which overflows 65535.)
+- **Dataset:** FineWeb-Edu sample-10BT primary (best small-scale benchmark
+  movement + Pythia comparability). Loader is dataset-agnostic so DCLM /
+  Nemotron-CC are drop-in data-ablation rows.
+- **Cost discipline:** only the ONE long run costs real money (~$5–15). Ablations
+  ~$0.60 each. ALL correctness validated FREE on Colab/CPU first; a ~$0.20 A100
+  smoke test maximizes MFU before committing to the long run.
+
+### Phase plan
+1. ✅ Model + sanity tests
+2. ✅ Infra harness (checkpoint/resume, NaN-guard, monitor/MFU, optimizer, loop)
+3. ✅ Data pipeline (FineWeb→uint16 shards + checksum, mmap BinStream)
+4. ⏳ NEXT: Colab GPU correctness pass + $0.20 A100 MFU calibration
+5. Long dense run (~3B tokens)
+6. Architecture ablations (RoPE/RMSNorm/SwiGLU/GQA/QK-norm on/off)
+7. Muon vs AdamW
+8. MoE ablation
+9. Eval (lm-eval-harness: HellaSwag / ARC-easy / PIQA vs Pythia-160M)
+
+## WHAT IS BUILT (Phases 1–3 complete; 18/18 tests pass on CPU)
+
+Repository root on my machine:
+`C:\Users\Swajay\Downloads\trade\matilda-mini\`
+
+File map (line counts):
+
+```
+src/matilda/
+  config.py      (65)   frozen dataclass; DEV_TINY + BASE_124M (114M total / 75.5M non-embed)
+  model.py      (199)   dense decoder: RoPE, RMSNorm, SwiGLU, GQA, QK-Norm, weight tying, residual-scaled init
+  optim.py       (48)   AdamW with param groups (no WD on norms/biases/embeddings); cosine+warmup schedule
+  checkpoint.py  (95)   atomic write (tmp->os.replace), rotation, saves model+opt+sched+step+RNG+dataloader pos
+  monitor.py     (70)   MFU, tokens/s, rolling step-time window, per-GPU peak-TFLOPS table
+  data.py       (182)   ShardWriter (+SHA-256 manifest), verify_manifest, SyntheticStream, BinStream (mmap, resumable)
+  train.py      (236)   loop: bf16 autocast, grad-accum + DDP no_sync, grad-clip, NaN/Inf guard, SIGTERM->checkpoint, auto-resume
+  __init__.py     (4)
+scripts/
+  prepare_data.py (73)  stream HF dataset (default FineWeb-Edu sample-10BT) -> tokenize gpt2 -> uint16 shards + verified manifest
+tests/            (18 tests total, all green on CPU/torch 2.7.1)
+  test_model.py   (86)  forward shapes, weight-tying, init-loss≈log(V), CAUSAL-MASK-NO-FUTURE-LEAK, GQA head counts, OVERFIT-SINGLE-BATCH
+  test_checkpoint.py(116) BIT-FOR-BIT RESUME (post-resume losses match uninterrupted run <1e-6), rotation/latest
+  test_train.py   (79)  loop runs+checkpoints, loop-level resume, NaN-abort, NaN-skip-then-recover, MFU sanity
+  test_data.py    (78)  shard round-trip, CHECKSUM CORRUPTION DETECTION, BinStream next-token shift, deterministic resume, multi-shard
+conftest.py, pytest.ini, requirements.txt
+```
+
+Verified facts:
+- Overfit-one-batch: loss → <0.1 on a fixed batch (model can learn).
+- Causal mask: perturbing token t leaves all logits at positions < t bit-identical.
+- Resume: interrupt at the halfway point, rebuild fresh model/opt/sched/stream,
+  restore, continue → per-step losses match an uninterrupted run to <1e-6.
+- Data: corrupting one byte of a shard is caught by verify_manifest.
+- BinStream → Trainer integration: loss drops on structured data with train.py unchanged.
+
+## WHAT I WANT FROM YOU (review checklist)
+
+Please go through these and report concrete findings with file/function references:
+
+1. **Transformer correctness.** RoPE (half-rotation convention, cos/sin build,
+   applied to Q/K only, before/after QK-norm ordering), GQA repeat_interleave vs
+   repeat semantics, SwiGLU 2/3 sizing, RMSNorm fp32 reduction, weight tying,
+   residual-projection init scaling 1/sqrt(2*n_layers). Any subtle bug?
+2. **Numerical stability for a long bf16 run.** Is QK-norm enough? Should I add
+   z-loss / logit soft-cap / attention-logit cap? Embedding init scale?
+3. **Checkpoint/resume completeness.** Anything that influences the next step that
+   I'm NOT saving (RNG, data position, scaler, AMP state)? Atomicity on Vast spot kill.
+4. **NaN/Inf guard design.** Is skip-batch correct, or should I roll back to last
+   checkpoint? Should the guard also catch inf grad-norm before clip (it does)?
+   Is consecutive-skip abort the right policy?
+5. **Throughput / MFU.** Is the 6N FLOPs/token approximation appropriate, or should
+   I include attention FLOPs (2*L*T*d term) at seq_len 1024? Will I actually hit
+   45-50% MFU on an A100 with this code, or what's missing (e.g., no torch.compile
+   tuning, no fused/foreach, dataloader not overlapped with compute, H2D copies not
+   pinned/non_blocking)?
+6. **DDP correctness.** no_sync usage on non-final micro-steps is wired but DDP
+   wrapping isn't applied (single-GPU). If I later wrap in DDP, what breaks
+   (checkpoint should save model.module, num_params unwrap, sampler sharding,
+   rank-0-only logging/checkpointing)?
+7. **Data pipeline.** Random-window sampling vs sequential cursor — implications for
+   epoch coverage and resume. Cross-document contamination from packing without
+   attention masking at seq_len 1024 — does it matter at this scale?
+8. **Chinchilla / token budget.** For ~114M total (~75M non-embedding), what token
+   count is right for a portfolio run (compute-optimal ~20x vs over-train 50-100x)?
+9. **Anything that would make a Maincode interviewer wince.** Missing observability,
+   missing reproducibility (git SHA logging?), config hygiene, test gaps.
+10. **Highest-ROI additions** before I spend money, ranked, given the infra-role framing.
+
+For each finding give: severity, file/location, the problem, and the concrete fix.
+
+---
+
+SOURCE FILES FOLLOW. I will paste them below in this order: config.py, model.py,
+optim.py, checkpoint.py, monitor.py, data.py, train.py, then the four test files.

docs/INSTANCE_RUNBOOK.md

@@ -0,0 +1,342 @@
+# Instance Runbook — Matilda-Mini
+
+> This is the operating manual for running Matilda-Mini on a remote GPU instance
+> (RTX 3090 for validation; A100 for the long run). Follow top to bottom. Each
+> step has a **gate** (what success looks like) — do not proceed past a failed gate.
+
+## 0. Context (read first)
+
+- **What:** a sub-200M modern dense transformer trained from scratch + production
+  training infra. Portfolio piece for a Maincode training-infrastructure role.
+- **Goal of THIS session:** validate the stack on a real CUDA GPU, then (optionally)
+  calibrate MFU and run training.
+- **Model (BASE_124M):** 114M params total / 75.5M non-embedding. d_model=768,
+  12 layers, 12 heads / 4 KV heads (GQA), seq_len=1024, GPT-2 50k vocab.
+- **Already proven on CPU (23 tests):** correctness, bit-for-bit resume, NaN guard,
+  data checksums. The unknowns the GPU adds: bf16 path, fused AdamW, Flash-SDPA,
+  torch.compile. That's what we validate here.
+- **Key files:** `run.py` (entrypoint), `configs/{calibration,base_124m}.json`,
+  `scripts/prepare_data.py`, `src/matilda/*`, `tests/*`.
+
+## 1. Get the code onto the instance
+
+Use whichever delivery method was set up (see "Code delivery" at bottom):
+
+```bash
+# Option A — public GitHub:
+git clone https://github.com/swajayresources/matilda-mini.git
+cd matilda-mini
+
+# Option B — public HF repo:
+git clone https://huggingface.co/prometheus04/matilda-mini
+cd matilda-mini
+
+# Option C — scp'd zip:
+unzip -oq matilda-mini.zip -d matilda-mini && cd matilda-mini
+```
+
+**Gate:** `ls src/matilda/model.py` exists.
+
+## 2. Environment setup
+
+```bash
+python --version            # expect 3.10+; if 'python3' use that
+nvidia-smi                  # confirm the GPU + driver; note GPU name + VRAM
+# Install torch FIRST from cu124 — the default pip wheel may be a cu130 build
+# that cannot init CUDA on current drivers (confirmed on the 3090 session).
+pip install -q torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
+pip install -q -r requirements.txt
+```
+
+Verify torch sees the GPU and bf16:
+
+```bash
+python -c "import torch; print('torch', torch.__version__, '| cuda', torch.cuda.is_available(), '|', torch.cuda.get_device_name(0)); print('bf16 native:', torch.cuda.is_bf16_supported())"
+```
+
+**Gate:** `cuda True`, GPU name printed, `bf16 native: True` (Ampere/3090/A100).
+If torch has no CUDA, install the CUDA wheel: `pip install torch --index-url https://download.pytorch.org/whl/cu121`.
+
+## 3. Test suite (the correctness gate)
+
+```bash
+python -m pytest tests/ -q
+```
+
+**Gate: `23 passed`.** If something fails here that passed on CPU, it's a real
+GPU-only bug — STOP, capture the full traceback, fix before spending on training.
+Likely suspects: `fused=True` AdamW (optim.py), `pin_memory` in BinStream (data.py).
+
+## 4. Tiny real-data smoke (end-to-end on real tokens)
+
+```bash
+# tokenize ~20M tokens of FineWeb-Edu (1-3 min). No HF token needed (public dataset).
+python scripts/prepare_data.py --out-dir data/smoke --target-tokens 20000000 --shard-tokens 20000000
+python -c "import sys; sys.path.insert(0,'src'); from matilda.data import verify_manifest; print('verified:', verify_manifest('data/smoke'))"
+
+# short run; loss should fall from ~10.8 toward ~7
+python run.py --config configs/calibration.json --data-dir data/smoke \
+  --set train.total_steps=100 train.warmup_steps=10 train.batch_size=8 \
+        train.compile=false train.ckpt_dir=checkpoints/smoke
+```
+
+**Gate:** `verified: True`, and loss visibly **decreasing** over the 100 steps
+(not flat, not NaN). Flat loss = data/label bug; NaN spam = stability bug.
+
+## 5. Resume check on GPU
+
+```bash
+# rerun with a higher step count; must print "[resume] ... at step 100" then continue
+python run.py --config configs/calibration.json --data-dir data/smoke \
+  --set train.total_steps=150 train.warmup_steps=10 train.batch_size=8 \
+        train.compile=false train.ckpt_dir=checkpoints/smoke
+```
+
+**Gate:** logs show `[resume] from checkpoints/smoke/ckpt_100.pt at step 100`,
+then steps 100→150 with **no loss spike** at the seam.
+
+> ✅ If steps 3–5 pass, the stack is validated on real hardware. Decide GPU for
+> the long run:
+> - **3090 (24GB):** ~$4 / ~17h — already validated, fully sufficient.
+> - **A100 40GB (recommended):** ~4h, cheaper than 80GB. A 124M model needs
+>   nowhere near 80GB — the memory hotspot is the vocab-projection logits, not
+>   weights. 40GB fits BS ~48-64; grad_accum reaches the token target regardless.
+> - **A100 80GB:** overkill here; only marginally fewer accum steps.
+>
+> Attention uses PyTorch's built-in SDPA (Flash Attention 2 kernel) — **no
+> `flash-attn` package, no compilation, nothing to version-match.** It can't break.
+
+## 6. MFU calibration (do this on the SAME GPU you'll run long on)
+
+Goal: maximize MFU by finding the largest batch that fits + benefits from compile.
+Run the calibration config, sweeping batch_size; read `mfu_avg` from the log tail.
+
+```bash
+for BS in 16 24 32 48; do
+  echo "=== batch_size=$BS ==="
+  python run.py --config configs/calibration.json --data-dir data/smoke \
+    --set train.batch_size=$BS train.ckpt_dir=checkpoints/calib_$BS 2>&1 | tail -5
+done
+```
+
+Notes:
+- `calibration.json` has `compile=true` — the first ~20 steps are slow (compile +
+  cuDNN autotune); the monitor excludes warmup ticks, so trust the later `mfu`.
+- If a batch size OOMs, that's your ceiling — back off one. The OOM hotspot is the
+  `lm_head` vocab projection (50257-wide), not attention.
+- Pick the batch with the best steady `mfu_avg`. **Record it.**
+
+**Validated result (RTX 3090, this stack):** BS=24 → **53.4% MFU** with compile;
+BS>=28 OOMs. `base_124m.json` already ships these (bs=24, grad_accum=21,
+total_steps=5814 ≈ 3.0B tokens).
+
+**A100 re-derivation:** the A100 has far more VRAM — raise `batch_size` to its
+ceiling (expect ~64 on 40GB, ~128 on 80GB), then keep tokens/step ≈ 0.5M and
+total tokens ≈ 3B by recomputing the other two:
+
+```
+tokens_per_step = batch_size * grad_accum * seq_len   # target ~524288
+grad_accum   = round(524288 / (batch_size * 1024))
+total_steps  = round(3.0e9 / (batch_size * grad_accum * 1024))
+```
+
+| GPU | batch_size | grad_accum | total_steps | ~tokens |
+|-----|-----------|-----------|-------------|---------|
+| 3090 (validated) | 24 | 21 | 5814 | 3.00B |
+| A100 40GB (est.) | 64 | 8  | 5722 | 3.00B |
+| A100 80GB (est.) | 128 | 4 | 5722 | 3.00B |
+
+Apply with `--set` (no file edit needed), e.g. on A100 80GB:
+`--set train.batch_size=128 train.grad_accum=4 train.total_steps=5722`
+
+## 7. The long run
+
+Run under tmux so a dropped SSH connection doesn't kill it:
+
+```bash
+tmux new -s train
+# inside tmux:
+export REMOTE="s3://<bucket>/matilda"   # optional: checkpoint backup target
+bash scripts/launch_vast.sh configs/base_124m.json
+# detach: Ctrl-b then d ;  reattach later: tmux attach -t train
+```
+
+Or directly with nohup:
+```bash
+nohup python run.py --config configs/base_124m.json --data-dir data/fwedu > train.log 2>&1 &
+tail -f train.log
+```
+
+**Monitor:** `metrics.jsonl` in the ckpt dir has every logged step (loss, lr,
+grad_norm, mfu, tokens_per_s, gpu_mem_peak_gb). Watch for: loss steadily falling,
+grad_norm stable (~0.2-1.5), mfu flat (a drop = throttling/noisy neighbor),
+nan_skip events (a few are fine; many = lower lr).
+
+**If it dies / SSH drops:** reconnect, `tmux attach` (or just re-run the same
+launch command) — it auto-resumes from the latest checkpoint, bit-for-bit.
+
+## 7b. Architecture ablations (cheap, high portfolio signal)
+
+Each variant trains on the same token budget; one thing changes per row. ~300M
+tokens each (~4 variants ≈ $2 on a 3090). Emits `docs/ABLATIONS.md` + JSON.
+
+```bash
+python scripts/ablate.py --data-dir data/fwedu --tokens 300000000
+# subset: python scripts/ablate.py --data-dir data/fwedu --only baseline mha
+```
+Variants: `baseline` (full modern), `no_qk_norm` (+softcap), `mha` (no GQA),
+`mqa` (extreme GQA). Read `docs/ABLATIONS.md` — that table is the README centerpiece.
+Save it OFF the instance (section 9). Add more rows by editing the `VARIANTS`
+list in `scripts/ablate.py`.
+
+## 8. Evaluation (after the run)
+
+```bash
+pip install -q lm-eval
+# convert/point lm-eval at the checkpoint, then:
+lm_eval --model hf --model_args pretrained=<exported_dir>,dtype=bfloat16 \
+  --tasks hellaswag,arc_easy,piqa,winogrande --batch_size 32
+```
+Target (124M, ~3B tokens): HellaSwag ~30-35%, ARC-easy ~40-45%, PIQA ~60%.
+Compare to Pythia-160M (HellaSwag ~30, ARC-e ~40, PIQA ~62).
+
+## 9. Save artifacts before killing the instance
+
+The instance is ephemeral — pull results off it:
+```bash
+# the final checkpoint + the full metric history
+ls -la checkpoints/base_124m/
+# copy these OFF the box: latest ckpt_*.pt, metrics.jsonl, train.log
+# either: aws s3 cp ... / huggingface-cli upload ... / scp back to laptop
+```
+Minimum to keep for the README table: `metrics.jsonl` (the loss/MFU curves) and
+the final eval numbers.
+
+## Gotchas & decisions
+
+- **OOM:** lower `train.batch_size`, raise `grad_accum` to keep tokens/step constant.
+- **NaN spam:** lower `lr` (6e-4 → 3e-4). The guard skips bad batches but many
+  skips = unstable. The qk_norm=OFF ablation needs `attn_logit_softcap=20`.
+- **Low MFU (<25%):** confirm `compile=true`, bf16 engaged (Ampere+), data not
+  starving (it's pinned/non_blocking already).
+- **SSH drop:** always run training in tmux/nohup. Resume is automatic.
+- **git_sha "unknown":** only if the code wasn't cloned via git (zip path). Fine.
+- **Do NOT** commit `data/`, `checkpoints/`, `*.bin`, `*.pt` anywhere.
+
+## Code delivery (how the code got here)
+
+- [x] **HF public: `huggingface.co/prometheus04/matilda-mini`** ← v1 (124M):
+      `git clone https://huggingface.co/prometheus04/matilda-mini && cd matilda-mini`
+- [x] **HF public: `huggingface.co/prometheus04/matilda-mini-v2`** ← v2 (363M hero):
+      `git clone https://huggingface.co/prometheus04/matilda-mini-v2 && cd matilda-mini-v2`
+- [ ] GitHub private: `github.com/swajayresources/matilda-mini` (needs auth)
+- [ ] zip via scp
+
+---
+
+# v2 (363M hero run) — operating manual
+
+The v2 path supersedes v1 for the actual paid run. Lay out the same as the v1
+sections above, but with these substitutions.
+
+## v2.1 Hardware
+
+A100 SXM4 40GB at ≤ $1.05/hr (Vast.ai or equivalent). Confirm spot/interruptible
+is available — bit-for-bit resume is already proven, so spot is the right call.
+
+## v2.2 Environment
+
+```bash
+# v1 setup is unchanged (torch cu124 first, then requirements.txt)
+pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
+pip install -r requirements.txt
+
+# v2-specific: liger-kernel is required by BASE_350M (use_liger=true).
+# The model will fail loudly at Transformer __init__ if this is missing.
+pip install liger-kernel
+python -c "from liger_kernel.transformers.rms_norm import LigerRMSNorm; print('liger OK')"
+```
+
+## v2.3 Gate the run on the existing test suite
+
+```bash
+pytest tests/ -q                          # 35 passed expected (1 was skipped on CPU)
+# On GPU+liger the BASE_350M shape test should also pass — total 36 passed.
+```
+
+## v2.4 Tokenize the SmolLM mix
+
+```bash
+# 15B tokens, 75% FineWeb-Edu-dedup / 15% Cosmopedia v2 / 10% Python-Edu.
+# Token-based mixing (not doc-based) — see prepare_smollm_data.py for the
+# argmin-deficit picking that converges to exact ratios.
+python scripts/prepare_smollm_data.py \
+    --out-dir data/smollm_mix \
+    --target-tokens 15000000000
+```
+
+Disk cost: ~30GB output (15B × 2 bytes/token for uint16). Run-time on the
+GPU instance: ~1-2 hours of streaming + tokenization. Save the resulting
+`data/smollm_mix/` to durable storage if you'll re-launch.
+
+## v2.5 Calibrate batch size on A100 with Liger
+
+`base_350m.json` ships with `batch_size=32`, `grad_accum=16` (tokens/step ≈ 1M).
+Liger fused linear+CE removes the (B·T, vocab) logit memory hotspot that capped
+v1 at BS=24 on the 3090, so micro-batch on A100 40GB should land higher. Sweep
+to find the ceiling:
+
+```bash
+# 50-step smoke; watch peak GPU mem and MFU
+python run.py --config configs/base_350m.json --data-dir data/smollm_mix \
+    --set train.total_steps=50 train.batch_size=32 train.grad_accum=2 \
+          train.ckpt_dir=checkpoints/calib_b32
+
+python run.py --config configs/base_350m.json --data-dir data/smollm_mix \
+    --set train.total_steps=50 train.batch_size=48 train.grad_accum=2 \
+          train.ckpt_dir=checkpoints/calib_b48
+```
+
+Pick the largest BS that doesn't OOM, then re-derive `grad_accum` to keep
+`batch × accum × seq_len ≈ 1M` (e.g. BS=48 → grad_accum=11).
+
+## v2.6 Launch the hero run
+
+```bash
+tmux new -s matilda
+bash scripts/launch_vast.sh configs/base_350m.json
+# If you found BS=48 in calibration:
+#   --set train.batch_size=48 train.grad_accum=11
+# (override via the run.py CLI; launch_vast.sh forwards CONFIG only.)
+```
+
+ETA at the target ~50% MFU: 55-60h. Total cost ~$55-60 at $1.017/hr. Checkpoints
+land every 1000 steps (~1B tokens) in `checkpoints/base_350m/`, keeping the last 20.
+
+## v2.7 Eval (token-matched vs Pythia-410M)
+
+For each saved checkpoint, run lm-eval-harness against Pythia-410M's published
+checkpoints at the same token count. The plot for the writeup is loss/HellaSwag/
+ARC-e/LAMBADA vs tokens, with our curve overlaid on Pythia-410M's.
+
+```bash
+pip install lm-eval
+lm-eval --model hf --model_args pretrained=./checkpoints/base_350m/ckpt_15000.pt \
+        --tasks hellaswag,arc_easy,piqa,lambada_openai \
+        --batch_size 32 --device cuda
+```
+
+(Note: lm-eval expects a HF-loadable model; you may need a quick adapter that
+loads our nn.Module + tokenizer. Plan a 1-2h budget for the adapter.)
+
+## v2.8 Gotchas specific to v2
+
+- **liger-kernel missing:** Transformer init raises ImportError with the install
+  hint. Don't try to silently fall back — the cost story depends on Liger.
+- **head_dim 80 + SDPA:** valid but ~5% slower than 64. We accept this for the
+  deep+thin shape signal. Don't "fix" it by switching n_heads.
+- **z_loss spam:** with z_loss_coef=1e-4 the contribution should be <<CE.
+  If z-loss dominates loss, lower the coefficient or check logit magnitude.
+- **WSD decay knee:** the LR holds flat for 80% of training then drops sharply.
+  Eyeballing the loss curve, expect the "second descent" only in the last 20%.
+  This is normal and is the WSD selling point — earlier checkpoints stay useful.

notebooks/colab_validate.ipynb

@@ -0,0 +1,138 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Matilda-Mini — Colab GPU validation (Phase 4, part 1: free)\n",
+    "\n",
+    "Goal: confirm the code is correct on a **real CUDA GPU** before spending money on Vast.\n",
+    "This catches anything CPU hid (bf16 paths, fused AdamW, SDPA Flash backend).\n",
+    "\n",
+    "**Note on T4:** Colab's free T4 is Turing and has *no native bf16* (it emulates, slowly).\n",
+    "Validate *correctness* here; do **not** trust T4 speed/MFU numbers. MFU tuning happens\n",
+    "in the $0.20 A100 calibration."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "print('cuda:', torch.cuda.is_available())\n",
+    "print('device:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'cpu')\n",
+    "print('bf16 supported:', torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get the code\n",
+    "Either `git clone` your pushed repo, or upload+unzip the `matilda-mini` folder, then `cd` in."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Option A: from GitHub (after you push)\n",
+    "# !git clone https://github.com/<you>/matilda-mini.git\n",
+    "# %cd matilda-mini\n",
+    "\n",
+    "# Option B: upload a zip via the Files panel, then:\n",
+    "# !unzip -q matilda-mini.zip && cd matilda-mini\n",
+    "\n",
+    "!pip install -q -r requirements.txt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Run the test suite on GPU\n",
+    "Expect **23 passed**. If anything fails here that passed on CPU, it's a real GPU-only bug."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!python -m pytest tests/ -q"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Tiny real-data smoke: tokenize a little FineWeb-Edu and watch loss fall\n",
+    "50M tokens is plenty to confirm the end-to-end path. Loss should drop from ~10.8 toward ~7 quickly."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!python scripts/prepare_data.py --out-dir data/smoke --target-tokens 50000000 --shard-tokens 50000000\n",
+    "!python -c \"import sys; sys.path.insert(0,'src'); from matilda.data import verify_manifest; print('verified:', verify_manifest('data/smoke'))\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Short real run on the smoke shard. bf16 only engages on Ampere+ (A100/L4), not T4.\n",
+    "!python run.py --config configs/calibration.json --data-dir data/smoke \\\n",
+    "  --set train.total_steps=100 train.warmup_steps=10 train.batch_size=8 \\\n",
+    "        train.compile=false train.ckpt_dir=checkpoints/smoke"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Confirm resume works on GPU\n",
+    "Re-run the same command; it should print `[resume] from ... at step 100` and continue,\n",
+    "not restart. Bump `total_steps` to 150 to see it progress."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!python run.py --config configs/calibration.json --data-dir data/smoke \\\n",
+    "  --set train.total_steps=150 train.warmup_steps=10 train.batch_size=8 \\\n",
+    "        train.compile=false train.ckpt_dir=checkpoints/smoke"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If 23 tests pass, loss falls on real data, and resume continues cleanly — the code is\n",
+    "validated. Next: rent the A100 and run `configs/calibration.json` to tune batch size for\n",
+    "max MFU (with `compile=true`), then launch the long run."
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {"provenance": []},
+  "kernelspec": {"display_name": "Python 3", "name": "python3"},
+  "language_info": {"name": "python"}
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+markers =
+    slow: longer training-based sanity checks (overfit one batch)

requirements.txt

@@ -0,0 +1,13 @@
+# GPU: install torch FIRST from the cu124 index (the default wheel may be a
+# cu130 build that fails CUDA init on current drivers):
+#   pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
+torch>=2.4,<3.0
+numpy
+tiktoken
+datasets
+wandb
+pytest
+
+# v2 (350M hero) only. liger-kernel needs CUDA Triton, so install LAST and
+# only on the GPU instance. CPU dev boxes leave use_liger=False and skip it.
+liger-kernel>=0.4.0; platform_system != "Windows"

run.py

@@ -0,0 +1,83 @@
+"""Launch entrypoint: build configs + data stream and train.
+
+    python run.py --config configs/base_124m.json --data-dir data/fwedu
+    python run.py --config configs/calibration.json --data-dir data/fwedu
+    python run.py --config configs/calibration.json --dry-run   # synthetic, no data
+
+A config JSON has two objects, "model" and "train", whose keys map directly onto
+ModelConfig / TrainConfig fields. CLI --set k=v applies last-mile overrides
+(e.g. --set train.batch_size=48) so calibration sweeps don't need new files.
+"""
+
+from __future__ import annotations
+
+import sys
+import json
+import argparse
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent / "src"))
+
+from matilda.config import ModelConfig                    # noqa: E402
+from matilda.train import Trainer, TrainConfig            # noqa: E402
+from matilda.data import SyntheticStream, BinStream, shard_paths  # noqa: E402
+
+
+def _coerce(value: str):
+    for cast in (int, float):
+        try:
+            return cast(value)
+        except ValueError:
+            pass
+    if value.lower() in ("true", "false"):
+        return value.lower() == "true"
+    return value
+
+
+def apply_overrides(cfg: dict, overrides: list[str]) -> dict:
+    """--set train.batch_size=48 -> cfg['train']['batch_size'] = 48"""
+    for item in overrides:
+        path, _, raw = item.partition("=")
+        section, _, key = path.partition(".")
+        cfg.setdefault(section, {})[key] = _coerce(raw)
+    return cfg
+
+
+def build(config: dict):
+    mcfg = ModelConfig(**config.get("model", {}))
+    tcfg = TrainConfig(**config.get("train", {}))
+    return mcfg, tcfg
+
+
+def build_stream(mcfg, tcfg, data_dir: str | None, dry_run: bool):
+    if dry_run or not data_dir:
+        print("[data] synthetic stream (dry run, no real data)")
+        return SyntheticStream(mcfg.vocab_size, tcfg.batch_size, tcfg.seq_len,
+                               seed=tcfg.seed, device=tcfg.device)
+    paths = shard_paths(data_dir)
+    print(f"[data] {len(paths)} shards from {data_dir}")
+    return BinStream(paths, tcfg.batch_size, tcfg.seq_len, seed=tcfg.seed,
+                     device=tcfg.device)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", required=True)
+    ap.add_argument("--data-dir", default=None)
+    ap.add_argument("--dry-run", action="store_true")
+    ap.add_argument("--set", nargs="*", default=[], dest="overrides")
+    args = ap.parse_args()
+
+    config = json.loads(Path(args.config).read_text())
+    config = apply_overrides(config, args.overrides)
+    mcfg, tcfg = build(config)
+    stream = build_stream(mcfg, tcfg, args.data_dir, args.dry_run)
+
+    print(f"[run] model={mcfg.d_model}d/{mcfg.n_layers}L "
+          f"steps={tcfg.total_steps} bs={tcfg.batch_size}x{tcfg.grad_accum} "
+          f"seq={tcfg.seq_len} compile={tcfg.compile} ckpt={tcfg.ckpt_dir}")
+    Trainer(mcfg, tcfg, stream).train()
+
+
+if __name__ == "__main__":
+    main()

scripts/ablate.py

@@ -0,0 +1,160 @@
+"""Architecture ablation harness.
+
+Trains each variant for the SAME fixed token budget on the SAME data, then emits
+a markdown results table + JSON. The point of an ablation is a controlled
+comparison: one thing changes per row, everything else held fixed.
+
+    python scripts/ablate.py --data-dir data/fwedu --tokens 300000000
+    python scripts/ablate.py --dry-run            # tiny synthetic, for CI/sanity
+    python scripts/ablate.py --data-dir data/fwedu --only baseline no_qk_norm
+
+Each variant only overrides what it's testing; the rest inherits from BASE below.
+Results land in results/ablations/<name>/ (checkpoints + metrics.jsonl) and a
+combined docs/ABLATIONS.md + results/ablations.json.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import json
+import argparse
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "src"))
+
+from matilda.config import ModelConfig                          # noqa: E402
+from matilda.train import Trainer, TrainConfig                  # noqa: E402
+from matilda.model import Transformer                           # noqa: E402
+from matilda.data import SyntheticStream, BinStream, shard_paths  # noqa: E402
+
+# Base config every variant inherits. Real architecture, short token budget.
+BASE_MODEL = dict(vocab_size=50257, max_seq_len=1024, d_model=768,
+                  n_layers=12, n_heads=12, n_kv_heads=4)
+import torch  # noqa: E402
+
+_HAS_CUDA = torch.cuda.is_available()
+BASE_TRAIN = dict(seq_len=1024, batch_size=24, grad_accum=8, warmup_steps=80,
+                  lr=6e-4, log_every=20, ckpt_every=100000,
+                  device="cuda" if _HAS_CUDA else "cpu",
+                  dtype="bfloat16" if _HAS_CUDA else "float32",
+                  compile=_HAS_CUDA)
+
+# One change per row. Inherits BASE; only the listed keys differ.
+VARIANTS = [
+    {"name": "baseline",   "model": {}},
+    {"name": "no_qk_norm", "model": {"qk_norm": False, "attn_logit_softcap": 20.0}},
+    {"name": "mha",        "model": {"n_kv_heads": 12}},   # full multi-head (no GQA)
+    {"name": "mqa",        "model": {"n_kv_heads": 1}},    # multi-query (extreme GQA)
+    {"name": "muon",       "model": {}, "train": {"optimizer": "muon"}},
+]
+
+
+def steps_for_tokens(tokens, tcfg: TrainConfig) -> int:
+    return max(1, round(tokens / (tcfg.batch_size * tcfg.grad_accum * tcfg.seq_len)))
+
+
+def final_metrics(ckpt_dir, tail=5) -> dict:
+    """Average the last `tail` logged steps from metrics.jsonl."""
+    path = os.path.join(ckpt_dir, "metrics.jsonl")
+    try:
+        with open(path) as f:
+            rows = [json.loads(l) for l in f if l.strip()]
+    except FileNotFoundError:
+        return {"loss": None, "mfu": None}     # crashed before logging a step
+    steps = [r for r in rows if r.get("event") == "step"]
+    if not steps:
+        return {"loss": None, "mfu": None}
+    last = steps[-tail:]
+    avg = lambda k: sum(s[k] for s in last if s.get(k) is not None) / len(last)
+    return {"loss": round(avg("loss"), 4), "mfu": round(avg("mfu"), 4),
+            "tokens_per_s": round(avg("tokens_per_s"))}
+
+
+def run_variant(v, tokens, data_dir, dry_run, results_root):
+    model = {**BASE_MODEL, **v.get("model", {})}
+    train = {**BASE_TRAIN, **v.get("train", {})}
+    if dry_run:                                   # shrink for CPU/CI
+        model.update(d_model=64, n_layers=2, n_heads=4,
+                     n_kv_heads=min(4, model.get("n_kv_heads", 4)), max_seq_len=64)
+        if v["name"] == "mha":
+            model["n_kv_heads"] = 4
+        train.update(seq_len=64, batch_size=4, grad_accum=1, warmup_steps=2,
+                     device="cpu", dtype="float32", compile=False)
+    ckpt_dir = os.path.join(results_root, v["name"])
+    train["ckpt_dir"] = ckpt_dir
+
+    mcfg = ModelConfig(**model)
+    tcfg = TrainConfig(total_steps=steps_for_tokens(tokens, TrainConfig(**train)),
+                       **train)
+    active = Transformer(mcfg).num_params(non_embedding=True)
+
+    if dry_run or not data_dir:
+        stream = SyntheticStream(mcfg.vocab_size, tcfg.batch_size, tcfg.seq_len,
+                                 seed=0, device=tcfg.device)
+    else:
+        stream = BinStream(shard_paths(data_dir), tcfg.batch_size, tcfg.seq_len,
+                           seed=0, device=tcfg.device)
+
+    print(f"\n=== variant: {v['name']} | steps={tcfg.total_steps} "
+          f"active={active/1e6:.1f}M kv_heads={mcfg.n_kv_heads} "
+          f"qk_norm={mcfg.qk_norm} ===")
+    Trainer(mcfg, tcfg, stream).train()
+    m = final_metrics(ckpt_dir)
+    return {"name": v["name"], "active_params_m": round(active / 1e6, 1),
+            "n_kv_heads": mcfg.n_kv_heads, "qk_norm": mcfg.qk_norm,
+            "optimizer": tcfg.optimizer,
+            "final_loss": m["loss"], "mfu": m["mfu"],
+            "tokens_per_s": m.get("tokens_per_s")}
+
+
+def write_table(rows, tokens, out_md, out_json):
+    os.makedirs(os.path.dirname(out_md), exist_ok=True)
+    os.makedirs(os.path.dirname(out_json), exist_ok=True)
+    json.dump({"tokens_per_variant": tokens, "rows": rows},
+              open(out_json, "w"), indent=2)
+    lines = [
+        "# Architecture Ablations",
+        "",
+        f"Each variant trained on **{tokens/1e6:.0f}M tokens** of the same data, "
+        "identical except for the column under test.",
+        "",
+        "| Variant | Active params | n_kv_heads | QK-norm | Optimizer | Final loss | MFU | tok/s |",
+        "|---------|--------------|-----------|---------|-----------|-----------|-----|-------|",
+    ]
+    for r in rows:
+        mfu = f"{r['mfu']*100:.1f}%" if r["mfu"] is not None else "—"
+        tps = f"{r['tokens_per_s']:,}" if r.get("tokens_per_s") else "—"
+        loss = r["final_loss"] if r["final_loss"] is not None else "—"
+        lines.append(f"| {r['name']} | {r['active_params_m']}M | {r['n_kv_heads']} "
+                     f"| {r['qk_norm']} | {r.get('optimizer','adamw')} | {loss} "
+                     f"| {mfu} | {tps} |")
+    open(out_md, "w").write("\n".join(lines) + "\n")
+    print(f"\nwrote {out_md} and {out_json}")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data-dir", default=None)
+    ap.add_argument("--tokens", type=int, default=300_000_000)
+    ap.add_argument("--dry-run", action="store_true")
+    ap.add_argument("--only", nargs="*", default=None,
+                    help="subset of variant names to run")
+    ap.add_argument("--results", default=str(ROOT / "results" / "ablations"))
+    args = ap.parse_args()
+
+    tokens = 50_000 if args.dry_run else args.tokens
+    variants = [v for v in VARIANTS
+                if args.only is None or v["name"] in args.only]
+    rows = [run_variant(v, tokens, args.data_dir, args.dry_run, args.results)
+            for v in variants]
+    write_table(rows, tokens,
+                out_md=str(ROOT / "docs" / "ABLATIONS.md"),
+                out_json=str(Path(args.results) / "ablations.json"))
+    print("\n".join(f"  {r['name']:12} loss={r['final_loss']} mfu={r['mfu']}"
+                    for r in rows))
+
+
+if __name__ == "__main__":
+    main()

scripts/launch_vast.sh

@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+# Spot-safe launch wrapper for a Vast.ai A100 instance.
+#
+# The trainer itself traps SIGTERM and checkpoints before exit (train.py), so
+# this wrapper's job is environment setup + syncing checkpoints to durable
+# storage so an instance death doesn't lose them.
+#
+#   bash scripts/launch_vast.sh configs/calibration.json   # MFU smoke test
+#   bash scripts/launch_vast.sh configs/base_124m.json      # v1: 124M / FineWeb-Edu
+#   bash scripts/launch_vast.sh configs/base_350m.json      # v2: 350M / SmolLM mix
+#
+# Set REMOTE to an rclone/s3 target to enable checkpoint upload on exit.
+# Set DATA_DIR to override the default per-config data path.
+
+set -euo pipefail
+
+CONFIG="${1:-configs/base_124m.json}"
+REMOTE="${REMOTE:-}"                       # e.g. s3://my-bucket/matilda or gdrive:matilda
+CKPT_DIR="$(python -c "import json;print(json.load(open('$CONFIG'))['train']['ckpt_dir'])")"
+USE_LIGER="$(python -c "import json;print(json.load(open('$CONFIG'))['model'].get('use_liger', False))")"
+
+# Default DATA_DIR depends on which config is being run.
+if [ -z "${DATA_DIR:-}" ]; then
+  case "$CONFIG" in
+    *base_350m*) DATA_DIR="data/smollm_mix" ;;
+    *)           DATA_DIR="data/fwedu"      ;;
+  esac
+fi
+
+sync_checkpoints() {
+  [ -z "$REMOTE" ] && { echo "[sync] REMOTE unset, skipping upload"; return; }
+  echo "[sync] uploading $CKPT_DIR -> $REMOTE"
+  if command -v aws >/dev/null;     then aws s3 sync "$CKPT_DIR" "$REMOTE" || true
+  elif command -v rclone >/dev/null; then rclone copy "$CKPT_DIR" "$REMOTE" || true
+  fi
+}
+trap sync_checkpoints EXIT               # runs on normal exit AND on spot kill
+
+echo "[setup] installing deps"
+pip install -q -r requirements.txt
+if [ "$USE_LIGER" = "True" ]; then
+  echo "[setup] config has use_liger=True; ensuring liger-kernel is importable"
+  python -c "import liger_kernel" 2>/dev/null || pip install -q liger-kernel
+fi
+
+# Pull checkpoints back first so a relaunched instance resumes (no-op if absent).
+if [ -n "$REMOTE" ]; then
+  mkdir -p "$CKPT_DIR"
+  (command -v aws >/dev/null && aws s3 sync "$REMOTE" "$CKPT_DIR") || \
+  (command -v rclone >/dev/null && rclone copy "$REMOTE" "$CKPT_DIR") || true
+fi
+
+if [ ! -f "$DATA_DIR/manifest.json" ]; then
+  case "$CONFIG" in
+    *base_350m*)
+      echo "[data] no manifest in $DATA_DIR; building SmolLM 75/15/10 mix (~15B tokens)"
+      python scripts/prepare_smollm_data.py --out-dir "$DATA_DIR" \
+             --target-tokens 15000000000
+      ;;
+    *)
+      echo "[data] no manifest in $DATA_DIR; tokenizing FineWeb-Edu (~3B tokens)"
+      python scripts/prepare_data.py --out-dir "$DATA_DIR" \
+             --target-tokens 3000000000
+      ;;
+  esac
+fi
+
+echo "[train] launching $CONFIG"
+python run.py --config "$CONFIG" --data-dir "$DATA_DIR"

scripts/prepare_data.py

@@ -0,0 +1,73 @@
+"""Tokenize a HuggingFace text dataset into verifiable uint16 shards.
+
+Default: FineWeb-Edu sample-10BT (best small-scale benchmark movement, max
+comparability to Pythia/SmolLM). Dataset-agnostic by design so DCLM /
+Nemotron-CC become drop-in data-ablation rows:
+
+    python scripts/prepare_data.py --target-tokens 3_000_000_000 \
+        --dataset HuggingFaceFW/fineweb-edu --name sample-10BT --out-dir data/fwedu
+
+    # ablation rows (same flags, different source):
+    --dataset mlfoundations/dclm-baseline-1.0 --out-dir data/dclm
+    --dataset nvidia/Nemotron-CC            --out-dir data/nemotron
+
+Streams the source (no full download), so disk holds only the tokenized output.
+"""
+
+from __future__ import annotations
+
+import sys
+import argparse
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+from matilda.data import ShardWriter, verify_manifest  # noqa: E402
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu")
+    ap.add_argument("--name", default="sample-10BT")
+    ap.add_argument("--split", default="train")
+    ap.add_argument("--text-key", default="text")
+    ap.add_argument("--tokenizer", default="gpt2")
+    ap.add_argument("--target-tokens", type=int, default=3_000_000_000)
+    ap.add_argument("--shard-tokens", type=int, default=100_000_000)
+    ap.add_argument("--out-dir", default="data/fwedu")
+    args = ap.parse_args()
+
+    import tiktoken
+    from datasets import load_dataset
+
+    enc = tiktoken.get_encoding(args.tokenizer)
+    eot = enc.eot_token
+    assert enc.n_vocab <= 65535, "vocab > uint16; use a smaller tokenizer"
+
+    ds = load_dataset(args.dataset, name=args.name, split=args.split,
+                      streaming=True)
+    writer = ShardWriter(args.out_dir, shard_tokens=args.shard_tokens)
+
+    n_docs = 0
+    for doc in ds:
+        ids = enc.encode_ordinary(doc[args.text_key])
+        ids.append(eot)                       # document boundary
+        writer.add(ids)
+        n_docs += 1
+        if n_docs % 1000 == 0:
+            print(f"\rdocs={n_docs:,} tokens={writer.total_tokens:,}", end="")
+        if writer.total_tokens >= args.target_tokens:
+            break
+
+    manifest = writer.close(meta={
+        "dataset": args.dataset, "name": args.name, "split": args.split,
+        "tokenizer": args.tokenizer, "eot_token": eot, "n_docs": n_docs,
+    })
+    print(f"\nwrote {manifest['total_tokens']:,} tokens in "
+          f"{len(manifest['shards'])} shards -> {args.out_dir}")
+    verify_manifest(args.out_dir)
+    print("manifest verified (checksums + sizes OK)")
+
+
+if __name__ == "__main__":
+    main()

scripts/prepare_smollm_data.py

@@ -0,0 +1,122 @@
+"""Build the SmolLM-corpus 75/15/10 mix into verifiable uint16 shards.
+
+Adopts the SmolLM data recipe (proven sub-1B winner over Pythia-410M /
+MobileLLM-350M / Qwen2-500M) but re-tokenized through our locked GPT-2 BPE so
+that loader, shard format, and tokenizer-comparability with v1 are unchanged.
+
+Mix targets are token-based, not document-based: documents have wildly
+different lengths, so picking-by-doc-count would drift far from the intended
+mix. We pick from whichever source is most under-represented in *tokens*
+emitted so far, which converges to exactly the configured weights.
+
+Default 15B-token target matches configs/base_350m.json. Streams the three
+HuggingFaceTB/smollm-corpus subsets so disk only holds the tokenized output.
+
+Usage (on the GPU instance):
+
+    python scripts/prepare_smollm_data.py \
+        --target-tokens 15_000_000_000 \
+        --out-dir data/smollm_mix
+"""
+
+from __future__ import annotations
+
+import sys
+import argparse
+import logging
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+from matilda.data import ShardWriter, verify_manifest  # noqa: E402
+
+log = logging.getLogger("prepare_smollm_data")
+
+
+SOURCES: tuple[tuple[str, float], ...] = (
+    ("fineweb-edu-dedup", 0.75),   # bulk: web reading-comprehension signal
+    ("cosmopedia-v2",     0.15),   # the Cosmopedia premium (synthetic textbooks)
+    ("python-edu",        0.10),   # code: HumanEval signal, modest cost
+)
+DATASET = "HuggingFaceTB/smollm-corpus"
+TEXT_KEY = "text"
+
+
+def _open_stream(name: str):
+    from datasets import load_dataset
+    return iter(load_dataset(DATASET, name=name, split="train", streaming=True))
+
+
+def _pick_source(accumulated: dict[str, int]) -> str:
+    """Return the source whose token deficit relative to its weight is largest.
+    Equivalent to argmin(accumulated[s] / weight[s]) over s in SOURCES."""
+    return min(
+        (name for name, _ in SOURCES),
+        key=lambda n: accumulated[n] / next(w for s, w in SOURCES if s == n),
+    )
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--target-tokens", type=int, default=15_000_000_000)
+    ap.add_argument("--shard-tokens", type=int, default=100_000_000)
+    ap.add_argument("--out-dir", default="data/smollm_mix")
+    ap.add_argument("--tokenizer", default="gpt2")
+    ap.add_argument("--log-every", type=int, default=1000)
+    args = ap.parse_args()
+
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+    import tiktoken
+
+    enc = tiktoken.get_encoding(args.tokenizer)
+    eot = enc.eot_token
+    assert enc.n_vocab <= 65535, "vocab > uint16; loader assumes uint16 shards"
+
+    streams = {name: _open_stream(name) for name, _ in SOURCES}
+    accumulated = {name: 0 for name, _ in SOURCES}
+    n_docs = {name: 0 for name, _ in SOURCES}
+    writer = ShardWriter(args.out_dir, shard_tokens=args.shard_tokens)
+
+    while writer.total_tokens < args.target_tokens:
+        pick = _pick_source(accumulated)
+        try:
+            doc = next(streams[pick])
+        except StopIteration:
+            # Defensive: at 15B target with smollm-corpus sizes, this branch
+            # shouldn't trigger. If it does, restart the source so the mix
+            # ratio is preserved rather than silently rebalancing.
+            log.warning("source %s exhausted at %d tokens; restarting stream",
+                        pick, writer.total_tokens)
+            streams[pick] = _open_stream(pick)
+            doc = next(streams[pick])
+
+        ids = enc.encode_ordinary(doc[TEXT_KEY])
+        ids.append(eot)                                    # document boundary
+        writer.add(ids)
+        accumulated[pick] += len(ids)
+        n_docs[pick] += 1
+
+        total_docs = sum(n_docs.values())
+        if total_docs % args.log_every == 0:
+            mix = {n: f"{accumulated[n] / max(1, writer.total_tokens) * 100:5.2f}%"
+                   for n in accumulated}
+            log.info("docs=%d tokens=%d mix=%s",
+                     total_docs, writer.total_tokens, mix)
+
+    manifest = writer.close(meta={
+        "dataset": DATASET,
+        "splits": [name for name, _ in SOURCES],
+        "weights": {name: weight for name, weight in SOURCES},
+        "tokenizer": args.tokenizer,
+        "eot_token": eot,
+        "n_docs_per_source": n_docs,
+        "tokens_per_source": accumulated,
+    })
+    log.info("wrote %d tokens in %d shards -> %s",
+             manifest["total_tokens"], len(manifest["shards"]), args.out_dir)
+    verify_manifest(args.out_dir)
+    log.info("manifest verified (checksums + sizes OK)")
+
+
+if __name__ == "__main__":
+    main()

src/matilda/__init__.py

@@ -0,0 +1,4 @@
+from .config import ModelConfig, DEV_TINY, BASE_124M, BASE_350M
+from .model import Transformer
+
+__all__ = ["ModelConfig", "DEV_TINY", "BASE_124M", "BASE_350M", "Transformer"]

src/matilda/checkpoint.py

@@ -0,0 +1,106 @@
+"""Crash-safe checkpointing.
+
+A resumed run must continue *identically*, not just "without an obvious spike".
+That requires saving everything that influences the next step:
+  model, optimizer, scheduler, step, AND all RNG states AND the dataloader
+  position. Dropping the RNG or data position is the classic cause of a loss
+  blip on resume (you re-see data / re-sample dropout differently).
+
+Writes are atomic (write tmp -> os.replace) so an instance dying mid-save can
+never leave a half-written checkpoint that fails to load.
+"""
+
+from __future__ import annotations
+
+import os
+import glob
+import random
+from dataclasses import asdict
+
+import numpy as np
+import torch
+
+
+def _rng_state() -> dict:
+    state = {
+        "python": random.getstate(),
+        "numpy": np.random.get_state(),
+        "torch": torch.get_rng_state(),
+    }
+    if torch.cuda.is_available():
+        state["cuda"] = torch.cuda.get_rng_state_all()
+    return state
+
+
+def _set_rng_state(state: dict) -> None:
+    random.setstate(state["python"])
+    np.random.set_state(state["numpy"])
+    # RNG states must be CPU ByteTensors. A checkpoint loaded with
+    # map_location="cuda" moves them to GPU, which set_rng_state rejects.
+    torch.set_rng_state(state["torch"].cpu())
+    if "cuda" in state and torch.cuda.is_available():
+        torch.cuda.set_rng_state_all([s.cpu() for s in state["cuda"]])
+
+
+def save_checkpoint(path, *, model, optimizer, scheduler, step,
+                    config=None, data_state=None, extra=None) -> None:
+    """Atomically write a complete checkpoint to `path`.
+
+    `extra` carries run provenance (TrainConfig, git SHA) so a resume can detect
+    a changed schedule rather than silently corrupting the LR curve.
+    """
+    payload = {
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "scheduler": scheduler.state_dict() if scheduler is not None else None,
+        "step": step,
+        "rng": _rng_state(),
+        "data_state": data_state,
+        "config": asdict(config) if hasattr(config, "__dataclass_fields__") else config,
+        "extra": extra or {},
+    }
+    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+    tmp = f"{path}.tmp.{os.getpid()}"
+    torch.save(payload, tmp)
+    os.replace(tmp, path)  # atomic on POSIX and Windows
+
+
+def load_checkpoint(path, *, model, optimizer=None, scheduler=None,
+                    map_location="cpu", restore_rng=True) -> dict:
+    """Restore in-place. Returns the raw payload (for step, data_state, config).
+
+    weights_only=False is required to unpickle optimizer state. Safe here (we only
+    load our own checkpoints); never point this at an untrusted checkpoint.
+    """
+    ck = torch.load(path, map_location=map_location, weights_only=False)
+    model.load_state_dict(ck["model"])
+    if optimizer is not None and ck.get("optimizer") is not None:
+        optimizer.load_state_dict(ck["optimizer"])
+    if scheduler is not None and ck.get("scheduler") is not None:
+        scheduler.load_state_dict(ck["scheduler"])
+    if restore_rng and ck.get("rng") is not None:
+        _set_rng_state(ck["rng"])
+    return ck
+
+
+def latest_checkpoint(directory) -> str | None:
+    """Highest-step checkpoint matching ckpt_*.pt, or None."""
+    paths = glob.glob(os.path.join(directory, "ckpt_*.pt"))
+    if not paths:
+        return None
+    return max(paths, key=lambda p: int(os.path.basename(p)[5:-3]))
+
+
+def rotate_checkpoints(directory, keep_last: int, protect: set[str] | None = None) -> None:
+    """Delete oldest ckpt_*.pt beyond keep_last. `protect` = basenames to keep."""
+    protect = protect or set()
+    paths = sorted(
+        glob.glob(os.path.join(directory, "ckpt_*.pt")),
+        key=lambda p: int(os.path.basename(p)[5:-3]),
+    )
+    deletable = [p for p in paths if os.path.basename(p) not in protect]
+    for p in deletable[:-keep_last] if keep_last > 0 else deletable:
+        try:
+            os.remove(p)
+        except OSError as e:
+            print(f"[warn] checkpoint rotation could not delete {p}: {e}")

src/matilda/config.py

@@ -0,0 +1,94 @@
+"""Model + run configuration.
+
+One frozen dataclass is the single source of truth for a run. It gets logged
+verbatim (alongside the git SHA) so any run is exactly reproducible. Nothing in
+the training stack reads a hyperparameter from anywhere else.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field, asdict
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    # --- vocab / sequence ---
+    vocab_size: int = 50257          # GPT-2 (tiktoken "gpt2"); fits uint16
+    max_seq_len: int = 1024
+
+    # --- transformer shape ---
+    d_model: int = 768
+    n_layers: int = 12
+    n_heads: int = 12
+    n_kv_heads: int = 4              # GQA: n_heads must be divisible by this
+    mlp_ratio: float = 8 / 3         # SwiGLU keeps params ~= 4x dense MLP
+    mlp_multiple_of: int = 256       # round hidden dim up to this for kernels
+
+    # --- numerics / regularization ---
+    norm_eps: float = 1e-6
+    rope_theta: float = 10000.0
+    init_std: float = 0.02
+    dropout: float = 0.0             # pretraining: keep at 0
+
+    # --- structural switches ---
+    tie_weights: bool = True         # share embedding <-> lm_head
+    qk_norm: bool = True             # RMSNorm on Q,K before attention
+    attn_logit_softcap: float = 0.0  # >0 caps scores via tanh; needed for the
+                                     # qk_norm=False ablation so it can't NaN.
+                                     # 0 keeps the fast fused SDPA path.
+
+    # --- loss / kernel knobs (no-op at v1 defaults) ---
+    z_loss_coef: float = 0.0         # penalize log-partition (PaLM/Pythia); 0 = off
+    use_liger: bool = False          # swap RMSNorm / RoPE / linear+CE for Liger
+                                     # fused kernels on GPU. Requires liger-kernel.
+
+    @property
+    def head_dim(self) -> int:
+        assert self.d_model % self.n_heads == 0, "d_model must divide n_heads"
+        return self.d_model // self.n_heads
+
+    def __post_init__(self) -> None:
+        assert self.n_heads % self.n_kv_heads == 0, \
+            "n_heads must be divisible by n_kv_heads (GQA grouping)"
+        assert self.head_dim % 2 == 0, "head_dim must be even for RoPE"
+
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        d["head_dim"] = self.head_dim
+        return d
+
+
+# Tiny config for free correctness validation on Colab T4 / RTX 3060.
+# Real vocab kept so tokenizer wiring is exercised; dims shrunk so it runs fast.
+DEV_TINY = ModelConfig(
+    vocab_size=50257,
+    max_seq_len=256,
+    d_model=128,
+    n_layers=2,
+    n_heads=4,
+    n_kv_heads=2,
+)
+
+# ~124M params (GPT-2 base footprint). The original A100 long-run config.
+BASE_124M = ModelConfig()
+
+
+# ~363M params total / ~315M non-embed. The 350M-family hero config.
+# Shape rationale: locked GQA-4 + locked n_layer=32 + clean head_dim (80) +
+# embedding tying. Depth ratio 32/960 = 0.033 (vs Pythia-410M's 24/1024 = 0.023)
+# = ~50% deeper than the closest published anchor at comparable width.
+# Liger Kernel + z-loss on; both no-ops if the flags are flipped off.
+BASE_350M = ModelConfig(
+    vocab_size=50257,
+    max_seq_len=2048,
+    d_model=960,
+    n_layers=32,
+    n_heads=12,             # head_dim = 80; 12 % 4 (n_kv_heads) == 0 ✓
+    n_kv_heads=4,           # GQA: 3 query heads per KV head
+    mlp_ratio=8 / 3,        # SwiGLU; 8/3 * 960 = 2560 lands on the 256 grid
+    mlp_multiple_of=256,
+    tie_weights=True,
+    qk_norm=True,
+    z_loss_coef=1e-4,
+    use_liger=True,
+)

src/matilda/data.py

@@ -0,0 +1,205 @@
+"""Data streams + tokenized-shard tooling.
+
+All streams expose the same interface so the training loop never changes when
+we swap synthetic data for real tokenized FineWeb shards:
+
+    x, y = stream.next()            # (B, T) input, (B, T) targets
+    state = stream.state_dict()     # resumable position / RNG
+    stream.load_state_dict(state)
+
+On-disk format: token ids as a flat `uint16` `.bin` (valid for vocab <= 65535,
+which GPT-2's 50257 satisfies). A `manifest.json` records per-shard token counts
+and SHA-256 so the prepared data is verifiable and reproducible.
+"""
+
+from __future__ import annotations
+
+import os
+import json
+import glob
+import hashlib
+
+import numpy as np
+import torch
+
+DTYPE = np.uint16
+
+
+# --------------------------------------------------------------------------
+# shard writing / verification (used by scripts/prepare_data.py, unit-tested)
+# --------------------------------------------------------------------------
+def sha256_file(path, chunk=1 << 20) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for block in iter(lambda: f.read(chunk), b""):
+            h.update(block)
+    return h.hexdigest()
+
+
+class ShardWriter:
+    """Accumulate token ids and flush fixed-size `.bin` shards + a manifest."""
+
+    def __init__(self, out_dir, shard_tokens=100_000_000, prefix="shard"):
+        self.out_dir = out_dir
+        self.shard_tokens = shard_tokens
+        self.prefix = prefix
+        os.makedirs(out_dir, exist_ok=True)
+        self._buf: list[np.ndarray] = []
+        self._buf_len = 0
+        self._shard_idx = 0
+        self.total_tokens = 0
+        self.manifest_entries: list[dict] = []
+
+    def add(self, tokens) -> None:
+        arr = np.asarray(tokens, dtype=DTYPE)
+        self._buf.append(arr)
+        self._buf_len += arr.size
+        self.total_tokens += arr.size
+        while self._buf_len >= self.shard_tokens:
+            self._flush(self.shard_tokens)
+
+    def _flush(self, n) -> None:
+        flat = np.concatenate(self._buf) if len(self._buf) > 1 else self._buf[0]
+        out, rest = flat[:n], flat[n:]
+        path = os.path.join(self.out_dir, f"{self.prefix}_{self._shard_idx:05d}.bin")
+        out.tofile(path)
+        self.manifest_entries.append({
+            "file": os.path.basename(path),
+            "tokens": int(out.size),
+            "sha256": sha256_file(path),
+        })
+        self._shard_idx += 1
+        self._buf = [rest] if rest.size else []
+        self._buf_len = rest.size
+
+    def close(self, meta: dict | None = None) -> dict:
+        if self._buf_len > 0:
+            self._flush(self._buf_len)        # final partial shard
+        manifest = {
+            "total_tokens": int(self.total_tokens),
+            "shards": self.manifest_entries,
+            **(meta or {}),
+        }
+        with open(os.path.join(self.out_dir, "manifest.json"), "w") as f:
+            json.dump(manifest, f, indent=2)
+        return manifest
+
+
+def verify_manifest(out_dir) -> bool:
+    """Re-checksum every shard against the manifest. Raises on mismatch."""
+    with open(os.path.join(out_dir, "manifest.json")) as f:
+        manifest = json.load(f)
+    for entry in manifest["shards"]:
+        path = os.path.join(out_dir, entry["file"])
+        actual = sha256_file(path)
+        if actual != entry["sha256"]:
+            raise ValueError(f"checksum mismatch for {entry['file']}")
+        if os.path.getsize(path) != entry["tokens"] * 2:  # uint16 = 2 bytes
+            raise ValueError(f"size mismatch for {entry['file']}")
+    return True
+
+
+def shard_paths(out_dir) -> list[str]:
+    with open(os.path.join(out_dir, "manifest.json")) as f:
+        manifest = json.load(f)
+    return [os.path.join(out_dir, e["file"]) for e in manifest["shards"]]
+
+
+# --------------------------------------------------------------------------
+# streams
+# --------------------------------------------------------------------------
+class SyntheticStream:
+    """Deterministic random tokens. For dev/CI and loop testing on CPU."""
+
+    def __init__(self, vocab_size, batch_size, seq_len, seed=0, device="cpu"):
+        self.vocab_size = vocab_size
+        self.B = batch_size
+        self.T = seq_len
+        self.device = device
+        self.gen = torch.Generator().manual_seed(seed)
+        self.pos = 0
+
+    def next(self):
+        seq = torch.randint(0, self.vocab_size, (self.B, self.T + 1),
+                            generator=self.gen)
+        x = seq[:, :-1].contiguous().to(self.device)
+        y = seq[:, 1:].contiguous().to(self.device)
+        self.pos += 1
+        return x, y
+
+    def state_dict(self):
+        return {"pos": self.pos, "gen": self.gen.get_state()}
+
+    def load_state_dict(self, s):
+        self.pos = s["pos"]
+        # set_state needs a CPU ByteTensor; map_location="cuda" can move it.
+        self.gen.set_state(s["gen"].cpu())
+
+
+class BinStream:
+    """Packed windows sampled from memory-mapped uint16 token shards.
+
+    Each batch element is a random T+1 window from a shard chosen with
+    probability proportional to its size; x/y are the next-token shift. Random
+    sampling (rather than a sequential cursor) means epochs are implicit and
+    resume only needs the RNG state.
+    """
+
+    def __init__(self, bin_paths, batch_size, seq_len, seed=0, device="cpu"):
+        assert bin_paths, "no shards provided"
+        self.paths = list(bin_paths)
+        self.B = batch_size
+        self.T = seq_len
+        self.device = device
+        # Drop shards too small to sample a T+1 window (e.g. the small final
+        # overflow shard prepare_data emits). Filter by on-disk size FIRST so we
+        # never open mmaps we'll discard (those hold a file lock on Windows).
+        min_bytes = (seq_len + 2) * 2          # uint16 = 2 bytes/token
+        keep = [p for p in self.paths if os.path.getsize(p) >= min_bytes]
+        dropped = len(self.paths) - len(keep)
+        if dropped:
+            print(f"[data] skipping {dropped} shard(s) smaller than seq_len+1")
+        assert keep, "no shard is large enough for seq_len+1"
+        self.paths = keep
+        self.arrays = [np.memmap(p, dtype=DTYPE, mode="r") for p in keep]
+        self.sizes = torch.tensor([float(a.size) for a in self.arrays])
+        self.weights = self.sizes / self.sizes.sum()
+        self.gen = torch.Generator().manual_seed(seed)
+        self.pos = 0
+        # int32 halves H2D bytes vs int64; cast to long happens on the GPU.
+        # uint16 max (50257 vocab) exceeds int16 range, so int32 is the floor.
+        pin = (str(device) != "cpu")
+        self._xbuf = torch.empty((batch_size, seq_len), dtype=torch.int32,
+                                 pin_memory=pin)
+        self._ybuf = torch.empty((batch_size, seq_len), dtype=torch.int32,
+                                 pin_memory=pin)
+
+    def next(self):
+        # vectorized sampling: one multinomial + one rand call, no per-element sync
+        shard_ids = torch.multinomial(self.weights, self.B, replacement=True,
+                                      generator=self.gen)
+        u = torch.rand(self.B, generator=self.gen)
+        sizes = self.sizes[shard_ids]
+        starts = (u * (sizes - (self.T + 1)).clamp(min=0)).long()
+        xb = np.empty((self.B, self.T), dtype=np.int32)
+        yb = np.empty((self.B, self.T), dtype=np.int32)
+        for i in range(self.B):
+            arr = self.arrays[int(shard_ids[i])]
+            s = int(starts[i])
+            window = arr[s:s + self.T + 1].astype(np.int32)
+            xb[i] = window[:-1]
+            yb[i] = window[1:]
+        self.pos += 1
+        self._xbuf.copy_(torch.from_numpy(xb))
+        self._ybuf.copy_(torch.from_numpy(yb))
+        x = self._xbuf.to(self.device, non_blocking=True).long()
+        y = self._ybuf.to(self.device, non_blocking=True).long()
+        return x, y
+
+    def state_dict(self):
+        return {"pos": self.pos, "gen": self.gen.get_state()}
+
+    def load_state_dict(self, s):
+        self.pos = s["pos"]
+        # set_state needs a CPU ByteTensor; map_location="cuda" can move it.
+        self.gen.set_state(s["gen"].cpu())

src/matilda/model.py

@@ -0,0 +1,278 @@
+"""Modernized dense decoder-only transformer.
+
+Block recipe (pre-norm): x = x + attn(rmsnorm(x)); x = x + swiglu(rmsnorm(x)).
+Modernizations baked in from the start (validated free on Colab before any
+paid run): RoPE, RMSNorm, SwiGLU, GQA, QK-Norm, weight tying, residual-scaled
+init. This is the architecture the long A100 run uses.
+
+Optional knobs for the 350M hero run, all default-off so the CPU/v1 path is
+unchanged:
+  - z_loss_coef > 0: PaLM-style log-Z penalty, applied alongside CE
+  - use_liger: swap RMSNorm/RoPE/(linear+CE) for Liger Triton kernels. The
+    fused linear+CE skips materializing the (B*T, vocab) logit tensor, which
+    is the dominant memory cost at 350M with vocab=50257.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import ModelConfig
+
+
+def _import_liger():
+    """Lazy import so CPU tests don't pay for liger-kernel install.
+    Raises ImportError with a helpful message if use_liger=True without it."""
+    try:
+        from liger_kernel.transformers.rms_norm import LigerRMSNorm
+        from liger_kernel.transformers.rope import liger_rotary_pos_emb
+        from liger_kernel.transformers.fused_linear_cross_entropy import (
+            LigerFusedLinearCrossEntropyLoss,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "use_liger=True requires liger-kernel. "
+            "Install: `pip install liger-kernel` (GPU only)."
+        ) from e
+    return LigerRMSNorm, liger_rotary_pos_emb, LigerFusedLinearCrossEntropyLoss
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm with the reduction done in fp32 for mixed-precision safety."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
+        # keep the scale multiply in fp32, then cast back once
+        return ((x * rms) * self.weight.float()).to(dtype)
+
+
+def build_rope_cache(head_dim: int, max_seq_len: int, theta: float,
+                     device=None, dtype=torch.float32):
+    """Precompute (cos, sin) of shape (max_seq_len, head_dim).
+
+    Half-rotation (Llama) convention: freqs are computed for head_dim/2 pairs
+    and concatenated with themselves so they line up with rotate_half.
+    """
+    i = torch.arange(0, head_dim, 2, device=device, dtype=torch.float32)
+    inv_freq = 1.0 / (theta ** (i / head_dim))             # (head_dim/2,)
+    t = torch.arange(max_seq_len, device=device, dtype=torch.float32)
+    freqs = torch.outer(t, inv_freq)                        # (T, head_dim/2)
+    emb = torch.cat([freqs, freqs], dim=-1)                 # (T, head_dim)
+    return emb.cos().to(dtype), emb.sin().to(dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    # x: (B, n_heads, T, head_dim); cos/sin: (T, head_dim) -> broadcast
+    cos = cos[None, None, :, :]
+    sin = sin[None, None, :, :]
+    return (x * cos) + (_rotate_half(x) * sin)
+
+
+class Attention(nn.Module):
+    """Causal grouped-query attention with optional QK-Norm and RoPE."""
+
+    def __init__(self, cfg: ModelConfig, liger=None):
+        super().__init__()
+        self.n_heads = cfg.n_heads
+        self.n_kv_heads = cfg.n_kv_heads
+        self.head_dim = cfg.head_dim
+        self.n_rep = cfg.n_heads // cfg.n_kv_heads
+
+        self.wq = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(cfg.n_heads * self.head_dim, cfg.d_model, bias=False)
+
+        self.qk_norm = cfg.qk_norm
+        if cfg.qk_norm:
+            norm_cls = liger[0] if liger else RMSNorm
+            self.q_norm = norm_cls(self.head_dim, cfg.norm_eps)
+            self.k_norm = norm_cls(self.head_dim, cfg.norm_eps)
+        self.dropout = cfg.dropout
+        self.softcap = cfg.attn_logit_softcap
+        self._liger_rope = liger[1] if liger else None
+
+    def _apply_rope(self, q, k, cos, sin):
+        if self._liger_rope is not None:
+            return self._liger_rope(q, k, cos, sin)
+        return apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+
+    def forward(self, x, cos, sin):
+        B, T, _ = x.shape
+        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
+        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
+        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        # (B, n_heads, T, head_dim)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        q, k = self._apply_rope(q, k, cos, sin)
+
+        # expand KV heads to match Q heads (GQA). repeat_interleave on head dim.
+        if self.n_rep > 1:
+            k = k.repeat_interleave(self.n_rep, dim=1)
+            v = v.repeat_interleave(self.n_rep, dim=1)
+
+        if self.softcap > 0:
+            out = self._attn_softcap(q, k, v, T)        # eager path, no Flash
+        else:
+            out = F.scaled_dot_product_attention(
+                q, k, v, is_causal=True,
+                dropout_p=self.dropout if self.training else 0.0,
+            )
+        out = out.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.wo(out)
+
+    def _attn_softcap(self, q, k, v, T):
+        # tanh logit soft-cap (Gemma-2 style). Disables the fused SDPA kernel,
+        # so only used for the qk_norm=False stability ablation.
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
+        scores = self.softcap * torch.tanh(scores / self.softcap)
+        mask = torch.ones(T, T, dtype=torch.bool, device=q.device).tril()
+        scores = scores.masked_fill(~mask, float("-inf"))
+        attn = F.softmax(scores, dim=-1)
+        return torch.matmul(attn, v)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        hidden = int(cfg.mlp_ratio * cfg.d_model)
+        m = cfg.mlp_multiple_of
+        hidden = ((hidden + m - 1) // m) * m
+        self.gate = nn.Linear(cfg.d_model, hidden, bias=False)
+        self.up = nn.Linear(cfg.d_model, hidden, bias=False)
+        self.down = nn.Linear(hidden, cfg.d_model, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: ModelConfig, liger=None):
+        super().__init__()
+        norm_cls = liger[0] if liger else RMSNorm
+        self.norm1 = norm_cls(cfg.d_model, cfg.norm_eps)
+        self.attn = Attention(cfg, liger=liger)
+        self.norm2 = norm_cls(cfg.d_model, cfg.norm_eps)
+        self.mlp = SwiGLU(cfg)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+
+        # Liger bundle: (RMSNorm class, rope fn, FLCE class) or None.
+        # Built once and threaded through Block/Attention so every layer uses
+        # the same kernel family — avoids per-layer import cost.
+        liger = _import_liger() if cfg.use_liger else None
+        self._liger = liger
+        self._flce = None
+        if liger is not None:
+            self._flce = liger[2](
+                ignore_index=-1,
+                lse_square_scale=cfg.z_loss_coef,    # z-loss handled in-kernel
+                label_smoothing=0.0,
+                reduction="mean",
+            )
+
+        norm_cls = liger[0] if liger else RMSNorm
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.blocks = nn.ModuleList(
+            [Block(cfg, liger=liger) for _ in range(cfg.n_layers)])
+        self.norm_f = norm_cls(cfg.d_model, cfg.norm_eps)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        if cfg.tie_weights:
+            self.lm_head.weight = self.embed.weight
+
+        cos, sin = build_rope_cache(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+        self.apply(self._init_weights)
+        # Residual-projection scaling (GPT-2 trick): keep residual-stream growth
+        # bounded with depth by shrinking the layers that write back into it.
+        scale = 1.0 / math.sqrt(2 * cfg.n_layers)
+        for name, p in self.named_parameters():
+            if name.endswith("wo.weight") or name.endswith("down.weight"):
+                with torch.no_grad():
+                    p.mul_(scale)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.cfg.init_std)
+            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=self.cfg.init_std)
+
+    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
+        B, T = idx.shape
+        assert T <= self.cfg.max_seq_len, "sequence longer than rope cache"
+        x = self.embed(idx)
+        cos = self.rope_cos[:T]
+        sin = self.rope_sin[:T]
+        for block in self.blocks:
+            x = block(x, cos, sin)
+        x = self.norm_f(x)
+
+        # Fused linear+CE path: skip logit materialization (the 350M+vocab=50k
+        # memory hot-spot). Returns logits=None — training never uses them.
+        if self._flce is not None and targets is not None:
+            loss = self._flce(
+                self.lm_head.weight,
+                x.view(-1, x.size(-1)),
+                targets.view(-1),
+            )
+            return None, loss
+
+        logits = self.lm_head(x)
+        if targets is None:
+            return logits, None
+
+        loss = F.cross_entropy(
+            logits.view(-1, logits.size(-1)),
+            targets.view(-1),
+            ignore_index=-1,
+        )
+        if self.cfg.z_loss_coef > 0:
+            # PaLM/Pythia-style: penalize log(Z) magnitude to keep logits bounded.
+            log_z = logits.logsumexp(dim=-1)
+            loss = loss + self.cfg.z_loss_coef * (log_z ** 2).mean()
+        return logits, loss
+
+    def num_params(self, non_embedding: bool = True) -> int:
+        n = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            # tied: lm_head shares embed, so subtract the one embedding table
+            n -= self.embed.weight.numel()
+        return n

src/matilda/monitor.py

@@ -0,0 +1,89 @@
+"""Throughput / utilization observability.
+
+MFU (Model FLOPs Utilization) is the headline number a training engineer lives
+in: what fraction of the GPU's theoretical bf16 throughput you actually extract.
+We also track a rolling step-time window so a *degradation* mid-run (thermal
+throttling, a noisy Vast neighbour) is visible, not just the instantaneous rate.
+"""
+
+from __future__ import annotations
+
+import time
+from collections import deque
+
+# Peak bf16 *dense* TFLOPs (no 2:4 sparsity). Substring-matched on device name.
+PEAK_TFLOPS_BF16 = {
+    "A100": 312.0,
+    "H100": 494.0,
+    "H800": 494.0,
+    "4090": 165.0,
+    "3090": 71.0,
+    "3060": 51.0,
+    "T4": 65.0,    # fp16; T4 has no native bf16 (Turing) -> emulated, ignore MFU
+    "A10": 125.0,
+    "L4": 121.0,
+}
+
+
+def peak_tflops(device_name: str, default: float = 312.0) -> float:
+    for key, val in PEAK_TFLOPS_BF16.items():
+        if key in device_name:
+            return val
+    return default
+
+
+def mfu(n_params_active: int, tokens_per_step: int, dt_seconds: float,
+        peak_flops_per_sec: float) -> float:
+    """Raw 6N-flops/token MFU (no attention term). Kept for quick estimates."""
+    if dt_seconds <= 0:
+        return 0.0
+    achieved = 6 * n_params_active * tokens_per_step / dt_seconds
+    return achieved / peak_flops_per_sec
+
+
+def flops_per_token(n_params, n_layers, d_model, seq_len) -> float:
+    """Karpathy/PaLM estimate: 6N (all matmuls incl. QKVO projections) plus the
+    attention score+value matmuls that scale with sequence length and are NOT
+    captured by the parameter count: 12 * L * d_model * T per token.
+    """
+    return 6 * n_params + 12 * n_layers * d_model * seq_len
+
+
+class Throughput:
+    """Rolling step-time tracker; reports tokens/sec and MFU.
+
+    `warmup` ticks are excluded from the running average so torch.compile /
+    cuDNN autotuning on the first steps doesn't depress the reported MFU.
+    """
+
+    def __init__(self, flops_per_step, tokens_per_step, peak_flops_per_sec,
+                 window=50, warmup=10):
+        self.fps = flops_per_step
+        self.tps = tokens_per_step
+        self.peak = peak_flops_per_sec
+        self.times = deque(maxlen=window)
+        self.warmup = warmup
+        self._n = 0
+        self._t0 = None
+
+    def _mfu(self, dt):
+        return (self.fps / dt) / self.peak if dt > 0 else 0.0
+
+    def tick(self) -> dict | None:
+        now = time.perf_counter()
+        if self._t0 is None:
+            self._t0 = now
+            return None
+        dt = now - self._t0
+        self._t0 = now
+        self._n += 1
+        if self._n > self.warmup:
+            self.times.append(dt)
+        avg = (sum(self.times) / len(self.times)) if self.times else dt
+        return {
+            "dt_s": dt,
+            "dt_avg_s": avg,
+            "tokens_per_s": self.tps / dt,
+            "mfu": self._mfu(dt),
+            "mfu_avg": self._mfu(avg),
+        }

src/matilda/optim.py

@@ -0,0 +1,277 @@
+"""Optimizer + LR schedule construction.
+
+Two details that materially affect quality and that most tutorials get wrong:
+  1. No weight decay on 1-D params (norms, biases) or the embedding table.
+     Decaying norm/embedding weights quietly hurts.
+  2. The LR schedule includes linear warmup; starting at peak LR on random
+     weights diverges.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+def _adamw(groups, lr, betas, eps):
+    """AdamW with fused kernel when available, falling back if unsupported."""
+    fused = torch.cuda.is_available()
+    try:
+        return torch.optim.AdamW(groups, lr=lr, betas=betas, eps=eps, fused=fused)
+    except (RuntimeError, TypeError):
+        return torch.optim.AdamW(groups, lr=lr, betas=betas, eps=eps)
+
+
+def build_adamw(model, lr=3e-4, weight_decay=0.1,
+                betas=(0.9, 0.95), eps=1e-8) -> torch.optim.AdamW:
+    decay, no_decay = [], []
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        # 2-D matmul weights get decay; norms/biases (1-D) and embeddings don't.
+        if p.ndim < 2 or "embed" in name:
+            no_decay.append(p)
+        else:
+            decay.append(p)
+    groups = [
+        {"params": decay, "weight_decay": weight_decay},
+        {"params": no_decay, "weight_decay": 0.0},
+    ]
+    return _adamw(groups, lr, betas, eps)
+
+
+@torch.no_grad()
+def zeropower_via_newtonschulz5(G, steps=5, eps=1e-7):
+    """Orthogonalize a 2-D gradient via the quintic Newton-Schulz iteration
+    (Keller Jordan). Pushes all singular values toward 1 in ~5 matmuls, so the
+    update weights every direction equally instead of being dominated by large
+    singular directions. Runs in bf16, as in the reference implementation.
+    """
+    assert G.ndim == 2
+    a, b, c = 3.4445, -4.7750, 2.0315
+    X = G.bfloat16()
+    X = X / (X.norm() + eps)
+    transposed = G.size(0) > G.size(1)
+    if transposed:
+        X = X.T
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * (A @ A)
+        X = a * X + B @ X
+    if transposed:
+        X = X.T
+    return X.to(G.dtype)
+
+
+class Muon(torch.optim.Optimizer):
+    """Muon: momentum + orthogonalized update for 2-D parameters only.
+
+    Use ONLY for interior matrices (attn/MLP weights). Embeddings, norms, biases
+    must stay on AdamW (see build_optimizer). Muon's natural LR is ~0.02, much
+    higher than Adam's, because the orthogonalized update has ~unit scale.
+    """
+
+    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5):
+        super().__init__(params, dict(lr=lr, momentum=momentum,
+                                      nesterov=nesterov, ns_steps=ns_steps))
+
+    @torch.no_grad()
+    def step(self):
+        for group in self.param_groups:
+            lr, mom, nesterov = group["lr"], group["momentum"], group["nesterov"]
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                g = p.grad
+                state = self.state[p]
+                if "m" not in state:
+                    state["m"] = torch.zeros_like(g)
+                buf = state["m"]
+                buf.mul_(mom).add_(g)
+                g = g.add(buf, alpha=mom) if nesterov else buf
+                g = zeropower_via_newtonschulz5(g, steps=group["ns_steps"])
+                # scale so the update RMS matches across differently-shaped matrices
+                scale = max(1.0, p.size(0) / p.size(1)) ** 0.5
+                p.add_(g, alpha=-lr * scale)
+
+
+class HybridOptimizer:
+    """Presents several optimizers as one (unified step/zero_grad/state_dict and
+    a concatenated param_groups so a single LR scheduler drives all of them)."""
+
+    def __init__(self, optimizers):
+        self.optimizers = optimizers
+
+    @property
+    def param_groups(self):
+        return [g for o in self.optimizers for g in o.param_groups]
+
+    def zero_grad(self, set_to_none=True):
+        for o in self.optimizers:
+            o.zero_grad(set_to_none=set_to_none)
+
+    def step(self):
+        for o in self.optimizers:
+            o.step()
+
+    def state_dict(self):
+        return {"opts": [o.state_dict() for o in self.optimizers]}
+
+    def load_state_dict(self, sd):
+        for o, s in zip(self.optimizers, sd["opts"]):
+            o.load_state_dict(s)
+
+
+def build_optimizer(model, name="adamw", lr=3e-4, weight_decay=0.1,
+                    betas=(0.9, 0.95), eps=1e-8, muon_lr=0.02):
+    """Dispatch: 'adamw' (default) or 'muon' (Muon on 2-D interior weights +
+    AdamW on embeddings/norms/biases)."""
+    if name == "adamw":
+        return build_adamw(model, lr, weight_decay, betas, eps)
+    if name != "muon":
+        raise ValueError(f"unknown optimizer: {name}")
+
+    muon_p, adamw_decay, adamw_nodecay = [], [], []
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if p.ndim == 2 and "embed" not in n:     # interior matrices -> Muon
+            muon_p.append(p)
+        elif p.ndim < 2 or "embed" in n:          # norms/biases/embeddings -> AdamW
+            adamw_nodecay.append(p)
+        else:
+            adamw_decay.append(p)
+    adamw = _adamw(
+        [{"params": adamw_decay, "weight_decay": weight_decay},
+         {"params": adamw_nodecay, "weight_decay": 0.0}],
+        lr, betas, eps)
+    return HybridOptimizer([Muon(muon_p, lr=muon_lr), adamw])
+
+
+class WarmupCosine:
+    """Linear warmup then cosine decay to min_lr_ratio * peak.
+
+    Works on any object exposing `param_groups` (torch optimizers AND our
+    HybridOptimizer, which torch's LambdaLR rejects). Applies one multiplier to
+    every group, scaling each group's own base lr (so Muon's 0.02 and AdamW's
+    3e-4 warm up/decay together). state_dict captures the step for exact resume.
+    """
+
+    def __init__(self, optimizer, warmup_steps, total_steps, min_lr_ratio=0.1):
+        self.opt = optimizer
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.min_lr_ratio = min_lr_ratio
+        self.base_lrs = [g["lr"] for g in optimizer.param_groups]
+        self.last_step = -1
+        self.step()                       # apply step 0
+
+    def _scale(self, step):
+        if step < self.warmup_steps:
+            return (step + 1) / max(1, self.warmup_steps)
+        if step >= self.total_steps:
+            return self.min_lr_ratio
+        progress = (step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps)
+        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+        return self.min_lr_ratio + (1 - self.min_lr_ratio) * cosine
+
+    def step(self):
+        self.last_step += 1
+        s = self._scale(self.last_step)
+        for group, base in zip(self.opt.param_groups, self.base_lrs):
+            group["lr"] = base * s
+
+    def get_last_lr(self):
+        return [g["lr"] for g in self.opt.param_groups]
+
+    def state_dict(self):
+        return {"last_step": self.last_step, "base_lrs": self.base_lrs}
+
+    def load_state_dict(self, sd):
+        self.last_step = sd["last_step"]
+        self.base_lrs = sd["base_lrs"]
+        # re-apply so lrs match the restored step
+        s = self._scale(self.last_step)
+        for group, base in zip(self.opt.param_groups, self.base_lrs):
+            group["lr"] = base * s
+
+
+def cosine_warmup_scheduler(optimizer, warmup_steps, total_steps, min_lr_ratio=0.1):
+    return WarmupCosine(optimizer, warmup_steps, total_steps, min_lr_ratio)
+
+
+class WarmupStableDecay:
+    """Linear warmup -> hold at peak -> linear decay to min_lr_ratio * peak.
+
+    MiniCPM-style WSD: empirically beats cosine at sub-1B scale, and intermediate
+    checkpoints stay useful because they're sampled at peak LR (no slow-roll decay).
+    `stable_share` is the fraction of the post-warmup steps spent at peak LR
+    (rest is the decay phase). 0.8 is MiniCPM's reported sweet spot.
+
+    Same interface contract as WarmupCosine so train.py just dispatches by name
+    and the rest of the stack (resume, multi-group LRs, HybridOptimizer) is
+    unchanged.
+    """
+
+    def __init__(self, optimizer, warmup_steps, total_steps,
+                 stable_share=0.8, min_lr_ratio=0.1):
+        assert 0.0 < stable_share < 1.0, "stable_share must be in (0, 1)"
+        self.opt = optimizer
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.stable_share = stable_share
+        self.min_lr_ratio = min_lr_ratio
+        post_warm = max(1, total_steps - warmup_steps)
+        self.decay_start = warmup_steps + int(stable_share * post_warm)
+        self.base_lrs = [g["lr"] for g in optimizer.param_groups]
+        self.last_step = -1
+        self.step()                       # apply step 0
+
+    def _scale(self, step):
+        if step < self.warmup_steps:
+            return (step + 1) / max(1, self.warmup_steps)
+        if step < self.decay_start:
+            return 1.0
+        if step >= self.total_steps:
+            return self.min_lr_ratio
+        progress = (step - self.decay_start) / max(
+            1, self.total_steps - self.decay_start)
+        return 1.0 - progress * (1.0 - self.min_lr_ratio)
+
+    def step(self):
+        self.last_step += 1
+        s = self._scale(self.last_step)
+        for group, base in zip(self.opt.param_groups, self.base_lrs):
+            group["lr"] = base * s
+
+    def get_last_lr(self):
+        return [g["lr"] for g in self.opt.param_groups]
+
+    def state_dict(self):
+        return {"last_step": self.last_step, "base_lrs": self.base_lrs}
+
+    def load_state_dict(self, sd):
+        self.last_step = sd["last_step"]
+        self.base_lrs = sd["base_lrs"]
+        s = self._scale(self.last_step)
+        for group, base in zip(self.opt.param_groups, self.base_lrs):
+            group["lr"] = base * s
+
+
+def wsd_scheduler(optimizer, warmup_steps, total_steps,
+                  stable_share=0.8, min_lr_ratio=0.1):
+    return WarmupStableDecay(optimizer, warmup_steps, total_steps,
+                             stable_share, min_lr_ratio)
+
+
+def build_scheduler(name, optimizer, warmup_steps, total_steps,
+                    stable_share=0.8, min_lr_ratio=0.1):
+    """Dispatch by name: 'cosine' (v1 default) or 'wsd' (MiniCPM)."""
+    if name == "cosine":
+        return cosine_warmup_scheduler(optimizer, warmup_steps, total_steps,
+                                       min_lr_ratio)
+    if name == "wsd":
+        return wsd_scheduler(optimizer, warmup_steps, total_steps,
+                             stable_share, min_lr_ratio)
+    raise ValueError(f"unknown lr_schedule: {name}")

src/matilda/train.py

@@ -0,0 +1,308 @@
+"""Training loop with the reliability guards a long run needs.
+
+Designed so an overnight Vast run survives the things that actually kill runs:
+  - one bad batch (NaN/Inf loss or grad) -> skip it, log loudly, keep going
+  - spot-instance SIGTERM -> checkpoint before exiting
+  - process death -> resume bit-for-bit from latest checkpoint on restart
+
+Throughput (MFU, tokens/s, step-time) is logged every `log_every` steps. The
+data source is any object with .next()/.state_dict()/.load_state_dict() (see
+data.py), so swapping synthetic data for real FineWeb shards changes nothing.
+"""
+
+from __future__ import annotations
+
+import os
+import json
+import signal
+import subprocess
+from dataclasses import dataclass, asdict
+
+import torch
+
+from .model import Transformer
+from .config import ModelConfig
+from .optim import build_optimizer, build_scheduler
+from .monitor import Throughput, peak_tflops, flops_per_token
+from .checkpoint import (
+    save_checkpoint, load_checkpoint, latest_checkpoint, rotate_checkpoints,
+)
+
+
+def get_git_sha() -> str:
+    try:
+        return subprocess.check_output(
+            ["git", "rev-parse", "HEAD"], text=True,
+            stderr=subprocess.DEVNULL).strip()
+    except Exception:
+        return "unknown"
+
+
+def setup_backends():
+    """Free A100 throughput: TF32 matmuls + cuDNN autotuner."""
+    if torch.cuda.is_available():
+        torch.set_float32_matmul_precision("high")
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+        torch.backends.cudnn.benchmark = True
+
+
+@dataclass
+class TrainConfig:
+    total_steps: int = 1000
+    warmup_steps: int = 100
+    grad_accum: int = 1
+    lr: float = 3e-4
+    weight_decay: float = 0.1
+    grad_clip: float = 1.0
+    min_lr_ratio: float = 0.1
+    optimizer: str = "adamw"          # "adamw" or "muon" (Muon+AdamW hybrid)
+    muon_lr: float = 0.02
+    lr_schedule: str = "cosine"       # "cosine" (v1) or "wsd" (MiniCPM, hero)
+    wsd_stable_share: float = 0.8     # only used when lr_schedule="wsd"
+
+    batch_size: int = 8
+    seq_len: int = 1024
+
+    log_every: int = 10
+    ckpt_every: int = 500
+    keep_last: int = 3
+    ckpt_dir: str = "checkpoints"
+
+    device: str = "cuda"
+    dtype: str = "bfloat16"            # bf16 on Ampere+, fp32 fallback on CPU
+    compile: bool = False
+    seed: int = 1234
+    max_skips: int = 20               # abort if too many bad batches in a row
+    wandb_project: str | None = None
+
+
+class Trainer:
+    def __init__(self, model_cfg: ModelConfig, train_cfg: TrainConfig, stream):
+        self.mcfg = model_cfg
+        self.cfg = train_cfg
+        self.stream = stream
+        self.step = 0
+        self.consecutive_skips = 0
+        self._interrupted = False
+
+        setup_backends()
+        self.git_sha = get_git_sha()
+        torch.manual_seed(train_cfg.seed)
+        self.device = torch.device(
+            train_cfg.device if torch.cuda.is_available()
+            or train_cfg.device == "cpu" else "cpu")
+
+        self.model = Transformer(model_cfg).to(self.device)
+        if train_cfg.compile:
+            self.model = torch.compile(self.model)
+        self.opt = build_optimizer(self.model, name=train_cfg.optimizer,
+                                   lr=train_cfg.lr,
+                                   weight_decay=train_cfg.weight_decay,
+                                   muon_lr=train_cfg.muon_lr)
+        self.sched = build_scheduler(
+            train_cfg.lr_schedule, self.opt,
+            train_cfg.warmup_steps, train_cfg.total_steps,
+            stable_share=train_cfg.wsd_stable_share,
+            min_lr_ratio=train_cfg.min_lr_ratio)
+
+        # bf16 only where supported; CPU/T4 fall back to fp32 for correctness.
+        want_bf16 = (train_cfg.dtype == "bfloat16"
+                     and self.device.type == "cuda"
+                     and torch.cuda.is_bf16_supported())
+        self.amp_dtype = torch.bfloat16 if want_bf16 else None
+
+        tokens_per_step = (train_cfg.batch_size * train_cfg.seq_len
+                           * train_cfg.grad_accum)
+        dev_name = (torch.cuda.get_device_name(self.device)
+                    if self.device.type == "cuda" else "cpu")
+        fps = flops_per_token(self._active_params(), model_cfg.n_layers,
+                              model_cfg.d_model, train_cfg.seq_len) * tokens_per_step
+        self.monitor = Throughput(
+            flops_per_step=fps,
+            tokens_per_step=tokens_per_step,
+            peak_flops_per_sec=peak_tflops(dev_name) * 1e12,
+        )
+        self.metrics_path = os.path.join(train_cfg.ckpt_dir, "metrics.jsonl")
+        self.wandb = self._init_wandb()
+        self._install_signal_handlers()
+
+    # --- helpers -----------------------------------------------------------
+    def _unwrap(self):
+        m = self.model
+        if hasattr(m, "_orig_mod"):          # torch.compile
+            m = m._orig_mod
+        if hasattr(m, "module"):             # DDP
+            m = m.module
+        return m
+
+    def _active_params(self):
+        return self._unwrap().num_params(non_embedding=True)
+
+    def _init_wandb(self):
+        if not self.cfg.wandb_project:
+            return None
+        try:
+            import wandb
+        except ImportError:
+            print("[warn] wandb not installed; using metrics.jsonl only")
+            return None
+        try:
+            wandb.init(project=self.cfg.wandb_project,
+                       config={**self.mcfg.to_dict(), **vars(self.cfg)})
+            return wandb
+        except Exception as e:                 # network/auth/quota — non-fatal
+            print(f"[warn] wandb init failed: {e}")
+            return None
+
+    def _install_signal_handlers(self):
+        def handler(signum, frame):
+            print(f"[signal] {signum} received -> checkpoint and exit")
+            self._interrupted = True
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            try:
+                signal.signal(sig, handler)
+            except (ValueError, OSError) as e:
+                # not in main thread (e.g. under pytest) -> can't install
+                print(f"[warn] could not install handler for signal {sig}: {e}")
+
+    def _autocast(self):
+        if self.amp_dtype is not None:
+            return torch.autocast(device_type="cuda", dtype=self.amp_dtype)
+        return torch.autocast(device_type="cpu", enabled=False)
+
+    # --- checkpoint --------------------------------------------------------
+    def _ckpt_path(self):
+        return os.path.join(self.cfg.ckpt_dir, f"ckpt_{self.step}.pt")
+
+    def save(self):
+        save_checkpoint(self._ckpt_path(), model=self.model, optimizer=self.opt,
+                        scheduler=self.sched, step=self.step, config=self.mcfg,
+                        data_state=self.stream.state_dict(),
+                        extra={"train_config": asdict(self.cfg),
+                               "git_sha": self.git_sha})
+        rotate_checkpoints(self.cfg.ckpt_dir, self.cfg.keep_last)
+
+    def maybe_resume(self):
+        path = latest_checkpoint(self.cfg.ckpt_dir)
+        if path is None:
+            return False
+        ck = load_checkpoint(path, model=self.model, optimizer=self.opt,
+                             scheduler=self.sched,
+                             map_location=self.device)
+        self.step = ck["step"]
+        if ck.get("data_state") is not None:
+            self.stream.load_state_dict(ck["data_state"])
+        # warn if the schedule changed since the checkpoint (silent LR corruption)
+        prev = (ck.get("extra") or {}).get("train_config", {})
+        for k in ("total_steps", "warmup_steps", "lr"):
+            if k in prev and prev[k] != getattr(self.cfg, k):
+                print(f"[warn] {k} changed on resume: {prev[k]} -> "
+                      f"{getattr(self.cfg, k)}; LR schedule will differ")
+        print(f"[resume] from {path} at step {self.step}")
+        return True
+
+    # --- one optimizer step (grad accum + guards) --------------------------
+    def _step(self):
+        self.opt.zero_grad(set_to_none=True)
+        total_loss = 0.0
+        for micro in range(self.cfg.grad_accum):
+            x, y = self.stream.next()
+            x = x.to(self.device, non_blocking=True)
+            y = y.to(self.device, non_blocking=True)
+            sync = (micro == self.cfg.grad_accum - 1)
+            ctx = (self.model.no_sync()
+                   if (not sync and hasattr(self.model, "no_sync"))
+                   else _nullcontext())
+            with ctx, self._autocast():
+                _, loss = self.model(x, y)
+                loss = loss / self.cfg.grad_accum
+            if not torch.isfinite(loss):
+                return None  # bad micro-batch -> abort this step
+            loss.backward()
+            total_loss += loss.item()
+
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            self.model.parameters(), self.cfg.grad_clip)
+        if not torch.isfinite(grad_norm):
+            return None  # non-finite grad -> skip update
+
+        self.opt.step()
+        self.sched.step()
+        return total_loss, grad_norm.item()
+
+    # --- main loop ---------------------------------------------------------
+    def train(self):
+        os.makedirs(self.cfg.ckpt_dir, exist_ok=True)
+        self.maybe_resume()
+        self._record({"event": "start", "git_sha": self.git_sha,
+                      "model": self.mcfg.to_dict(), "train": asdict(self.cfg),
+                      "step": self.step})
+        self.model.train()
+        while self.step < self.cfg.total_steps:
+            result = self._step()
+            if result is None:
+                self.consecutive_skips += 1
+                data_pos = self.stream.state_dict().get("pos")
+                print(f"[nan-guard] step {self.step} skipped "
+                      f"({self.consecutive_skips}/{self.cfg.max_skips}) "
+                      f"near data_pos={data_pos}")
+                self._record({"event": "nan_skip", "step": self.step,
+                              "data_pos": data_pos})
+                self.opt.zero_grad(set_to_none=True)
+                if self.consecutive_skips >= self.cfg.max_skips:
+                    raise RuntimeError("too many non-finite steps; aborting run")
+                continue
+            self.consecutive_skips = 0
+            loss, grad_norm = result
+            self.step += 1
+
+            stats = self.monitor.tick()
+            if self.step % self.cfg.log_every == 0:
+                self._log(loss, grad_norm, stats)
+            if self.step % self.cfg.ckpt_every == 0:
+                self.save()
+            if self._interrupted:
+                self.save()
+                print(f"[exit] checkpointed at step {self.step}")
+                break
+        else:
+            self.save()  # final checkpoint on clean completion
+        return self.step
+
+    def _record(self, row: dict):
+        """Always-on JSONL metric log (survives wandb outage)."""
+        try:
+            with open(self.metrics_path, "a") as f:
+                f.write(json.dumps(row) + "\n")
+        except OSError as e:
+            print(f"[warn] could not write metrics: {e}")
+
+    def _log(self, loss, grad_norm, stats):
+        lr = self.sched.get_last_lr()[0]
+        row = {"event": "step", "step": self.step, "loss": loss, "lr": lr,
+               "grad_norm": grad_norm}
+        if stats:
+            row.update({"tokens_per_s": stats["tokens_per_s"],
+                        "mfu": stats["mfu_avg"], "step_time_s": stats["dt_avg_s"]})
+        if self.device.type == "cuda":
+            row["gpu_mem_peak_gb"] = torch.cuda.max_memory_allocated() / 1e9
+        self._record(row)
+
+        msg = (f"step {self.step:>6} | loss {loss:7.4f} | lr {lr:.2e} "
+               f"| gnorm {grad_norm:5.2f}")
+        if stats:
+            msg += (f" | {stats['tokens_per_s']:8.0f} tok/s "
+                    f"| mfu {stats['mfu_avg']*100:4.1f}%")
+        print(msg)
+        if self.wandb:
+            self.wandb.log({k: v for k, v in row.items()
+                            if k not in ("event",)}, step=self.step)
+
+
+class _nullcontext:
+    def __enter__(self):
+        return None
+
+    def __exit__(self, *exc):
+        return False

tests/test_ablate.py

@@ -0,0 +1,40 @@
+"""Ablation harness: variants run, metrics captured, table emitted (dry CPU)."""
+
+import os
+import sys
+import json
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "scripts"))
+
+import ablate  # noqa: E402
+
+
+def test_steps_for_tokens():
+    from matilda.train import TrainConfig
+    tc = TrainConfig(batch_size=24, grad_accum=8, seq_len=1024)
+    # 300M / (24*8*1024) ~= 1526
+    assert ablate.steps_for_tokens(300_000_000, tc) == 1526
+
+
+def test_run_variant_dry_and_table(tmp_path):
+    results = str(tmp_path / "abl")
+    rows = []
+    for vname in ("baseline", "mqa"):
+        v = next(x for x in ablate.VARIANTS if x["name"] == vname)
+        row = ablate.run_variant(v, tokens=50_000, data_dir=None,
+                                 dry_run=True, results_root=results)
+        rows.append(row)
+        assert row["final_loss"] is not None      # metrics captured
+        assert row["mfu"] is not None
+    assert rows[1]["n_kv_heads"] == 1              # mqa override applied
+
+    md = tmp_path / "ABLATIONS.md"
+    js = tmp_path / "ablations.json"
+    ablate.write_table(rows, 50_000, str(md), str(js))
+    assert md.exists() and js.exists()
+    text = md.read_text()
+    assert "baseline" in text and "mqa" in text and "MFU" in text
+    data = json.loads(js.read_text())
+    assert len(data["rows"]) == 2

tests/test_checkpoint.py

@@ -0,0 +1,99 @@
+"""Resume must continue *identically* to an uninterrupted run.
+
+Strategy: run a short reference training. Run it again but checkpoint halfway,
+throw the objects away, build *fresh* model/optimizer/scheduler/data-stream,
+restore from the checkpoint, and continue. The post-resume per-step losses must
+match the reference run's second half exactly. If any piece of state is missing
+(opt moments, scheduler, RNG, data position) the curves diverge and this fails.
+"""
+
+import os
+import random
+
+import numpy as np
+import torch
+
+from matilda import Transformer, ModelConfig
+from matilda.data import SyntheticStream
+from matilda.optim import build_adamw, cosine_warmup_scheduler
+from matilda.checkpoint import (
+    save_checkpoint, load_checkpoint, latest_checkpoint, rotate_checkpoints,
+)
+
+CFG = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64,
+                  n_layers=2, n_heads=4, n_kv_heads=2)
+TOTAL, WARMUP, HALF = 12, 2, 6
+
+
+def _stream(seed):
+    return SyntheticStream(CFG.vocab_size, batch_size=4, seq_len=16, seed=seed)
+
+
+def _seed_all(seed=1234):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def _build():
+    model = Transformer(CFG).train()
+    opt = build_adamw(model, lr=1e-3)
+    sched = cosine_warmup_scheduler(opt, WARMUP, TOTAL)
+    return model, opt, sched
+
+
+def _train_steps(model, opt, sched, stream, n):
+    losses = []
+    for _ in range(n):
+        x, y = stream.next()
+        _, loss = model(x, y)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step()
+        sched.step()
+        losses.append(loss.item())
+    return losses
+
+
+def test_resume_is_bit_for_bit(tmp_path):
+    # --- reference: uninterrupted run ---
+    _seed_all()
+    m, o, s = _build()
+    ref_stream = _stream(42)
+    ref_losses = _train_steps(m, o, s, ref_stream, TOTAL)
+
+    # --- interrupted run: train HALF, checkpoint, drop everything ---
+    _seed_all()
+    m, o, s = _build()
+    stream = _stream(42)
+    first_half = _train_steps(m, o, s, stream, HALF)
+    assert first_half == ref_losses[:HALF], "training is not deterministic"
+
+    ckpt = os.path.join(tmp_path, f"ckpt_{HALF}.pt")
+    save_checkpoint(ckpt, model=m, optimizer=o, scheduler=s, step=HALF,
+                    config=CFG, data_state=stream.state_dict())
+
+    # --- fresh objects, restore, continue ---
+    m2, o2, s2 = _build()                    # fresh random init + zeroed moments
+    stream2 = _stream(999)      # deliberately wrong seed...
+    ck = load_checkpoint(ckpt, model=m2, optimizer=o2, scheduler=s2)
+    stream2.load_state_dict(ck["data_state"])  # ...corrected by restore
+    assert ck["step"] == HALF
+
+    resumed = _train_steps(m2, o2, s2, stream2, TOTAL - HALF)
+
+    for i, (a, b) in enumerate(zip(resumed, ref_losses[HALF:])):
+        assert abs(a - b) < 1e-6, f"resume diverged at step {HALF+i}: {a} vs {b}"
+
+
+def test_latest_and_rotation(tmp_path):
+    for step in (10, 20, 30, 40):
+        p = os.path.join(tmp_path, f"ckpt_{step}.pt")
+        torch.save({"step": step}, p)
+    assert latest_checkpoint(tmp_path).endswith("ckpt_40.pt")
+
+    rotate_checkpoints(tmp_path, keep_last=2, protect={"ckpt_10.pt"})
+    remaining = sorted(os.path.basename(p) for p in
+                       __import__("glob").glob(os.path.join(tmp_path, "ckpt_*.pt")))
+    # keep last 2 (30, 40) + protected 10; drop 20
+    assert remaining == ["ckpt_10.pt", "ckpt_30.pt", "ckpt_40.pt"]

tests/test_data.py

@@ -0,0 +1,90 @@
+"""Data pipeline: shard writing/verification + BinStream correctness."""
+
+import os
+import json
+
+import numpy as np
+import pytest
+import torch
+
+from matilda.data import (
+    ShardWriter, verify_manifest, shard_paths, BinStream, DTYPE,
+)
+
+
+def test_shardwriter_roundtrip_and_manifest(tmp_path):
+    w = ShardWriter(str(tmp_path), shard_tokens=100)
+    # 250 tokens -> two full shards (100) + one partial (50)
+    w.add(list(range(0, 120)))
+    w.add(list(range(120, 250)))
+    manifest = w.close(meta={"tokenizer": "gpt2"})
+
+    assert manifest["total_tokens"] == 250
+    assert [s["tokens"] for s in manifest["shards"]] == [100, 100, 50]
+    assert manifest["tokenizer"] == "gpt2"
+    assert verify_manifest(str(tmp_path)) is True
+
+    # reconstructed token stream matches the input exactly
+    rebuilt = np.concatenate(
+        [np.fromfile(p, dtype=DTYPE) for p in shard_paths(str(tmp_path))])
+    assert rebuilt.tolist() == list(range(250))
+
+
+def test_verify_detects_corruption(tmp_path):
+    w = ShardWriter(str(tmp_path), shard_tokens=100)
+    w.add(list(range(150)))
+    w.close()
+    # corrupt the first shard's bytes without changing its size
+    p = shard_paths(str(tmp_path))[0]
+    data = bytearray(open(p, "rb").read())
+    data[0] ^= 0xFF
+    open(p, "wb").write(data)
+    with pytest.raises(ValueError, match="checksum mismatch"):
+        verify_manifest(str(tmp_path))
+
+
+def _write_ramp_bin(path, n):
+    # values 0..n-1 so a correct next-token window always has y == x + 1
+    np.arange(n, dtype=DTYPE).tofile(path)
+
+
+def test_binstream_shift_is_next_token(tmp_path):
+    p = os.path.join(tmp_path, "ramp.bin")
+    _write_ramp_bin(p, 5000)
+    s = BinStream([p], batch_size=8, seq_len=32, seed=0)
+    x, y = s.next()
+    assert x.shape == (8, 32) and y.shape == (8, 32)
+    assert torch.equal(y, x + 1)              # packed next-token target
+
+
+def test_binstream_resume_is_deterministic(tmp_path):
+    p = os.path.join(tmp_path, "ramp.bin")
+    _write_ramp_bin(p, 5000)
+    s = BinStream([p], batch_size=4, seq_len=16, seed=7)
+    snap = s.state_dict()
+    x1, y1 = s.next()
+    s.load_state_dict(snap)
+    x2, y2 = s.next()
+    assert torch.equal(x1, x2) and torch.equal(y1, y2)
+
+
+def test_binstream_skips_too_small_shards(tmp_path):
+    # a tiny final shard (< seq_len+1) must be skipped, not crash (regression)
+    big = os.path.join(tmp_path, "big.bin")
+    tiny = os.path.join(tmp_path, "tiny.bin")
+    _write_ramp_bin(big, 5000)
+    _write_ramp_bin(tiny, 10)                  # smaller than seq_len+1
+    s = BinStream([big, tiny], batch_size=4, seq_len=32, seed=0)
+    assert len(s.arrays) == 1                  # tiny dropped
+    x, y = s.next()
+    assert torch.equal(y, x + 1)
+
+
+def test_binstream_multishard_weighting(tmp_path):
+    p1 = os.path.join(tmp_path, "a.bin")
+    p2 = os.path.join(tmp_path, "b.bin")
+    _write_ramp_bin(p1, 2000)
+    _write_ramp_bin(p2, 2000)
+    s = BinStream([p1, p2], batch_size=4, seq_len=8, seed=1)
+    x, y = s.next()
+    assert torch.equal(y, x + 1)              # still correct across shards

tests/test_model.py

@@ -0,0 +1,158 @@
+"""Sanity tests that gate the paid run. All must be green on Colab first.
+
+  - test_forward_shapes / test_weight_tying / test_param_count: wiring is correct
+  - test_causal_mask: no information leaks from future tokens (the bug that
+    silently inflates eval and is invisible in the loss curve)
+  - test_overfit_single_batch: the model can actually learn (loss -> ~0 on one
+    fixed batch). The cheapest, highest-signal correctness check in ML.
+"""
+
+import torch
+import pytest
+
+from matilda import Transformer, ModelConfig, DEV_TINY, BASE_350M
+
+
+def _tiny():
+    return Transformer(DEV_TINY).eval()
+
+
+def test_forward_shapes():
+    model = _tiny()
+    B, T = 2, 16
+    idx = torch.randint(0, DEV_TINY.vocab_size, (B, T))
+    logits, loss = model(idx)
+    assert logits.shape == (B, T, DEV_TINY.vocab_size)
+    assert loss is None
+    targets = torch.randint(0, DEV_TINY.vocab_size, (B, T))
+    _, loss = model(idx, targets)
+    assert loss is not None and loss.ndim == 0
+
+
+def test_weight_tying():
+    model = _tiny()
+    assert model.lm_head.weight.data_ptr() == model.embed.weight.data_ptr()
+
+
+def test_loss_at_init_is_near_uniform():
+    # untrained model should be ~ -log(1/V) = log(V)
+    model = _tiny()
+    idx = torch.randint(0, DEV_TINY.vocab_size, (4, 32))
+    tgt = torch.randint(0, DEV_TINY.vocab_size, (4, 32))
+    _, loss = model(idx, tgt)
+    expected = torch.log(torch.tensor(float(DEV_TINY.vocab_size)))
+    assert abs(loss.item() - expected.item()) < 1.0
+
+
+def test_causal_mask_no_future_leak():
+    """Changing token at position t must not alter logits at positions < t."""
+    model = _tiny()
+    torch.manual_seed(0)
+    idx = torch.randint(0, DEV_TINY.vocab_size, (1, 24))
+    with torch.no_grad():
+        base, _ = model(idx)
+        idx2 = idx.clone()
+        idx2[0, -1] = (idx2[0, -1] + 1) % DEV_TINY.vocab_size  # perturb last token
+        perturbed, _ = model(idx2)
+    # all positions except the last must be identical
+    assert torch.allclose(base[:, :-1], perturbed[:, :-1], atol=1e-5)
+    assert not torch.allclose(base[:, -1], perturbed[:, -1], atol=1e-5)
+
+
+def test_gqa_kv_head_counts():
+    model = _tiny()
+    attn = model.blocks[0].attn
+    assert attn.wk.out_features == DEV_TINY.n_kv_heads * DEV_TINY.head_dim
+    assert attn.wq.out_features == DEV_TINY.n_heads * DEV_TINY.head_dim
+
+
+def test_softcap_path_is_finite_and_causal():
+    # qk_norm OFF + soft-cap ON: the ablation config must stay finite and causal
+    cfg = ModelConfig(vocab_size=200, max_seq_len=64, d_model=64, n_layers=2,
+                      n_heads=4, n_kv_heads=2, qk_norm=False,
+                      attn_logit_softcap=20.0)
+    model = Transformer(cfg).eval()
+    idx = torch.randint(0, cfg.vocab_size, (2, 24))
+    with torch.no_grad():
+        logits, _ = model(idx)
+        assert torch.isfinite(logits).all()
+        idx2 = idx.clone()
+        idx2[0, -1] = (idx2[0, -1] + 1) % cfg.vocab_size
+        perturbed, _ = model(idx2)
+    assert torch.allclose(logits[:, :-1], perturbed[:, :-1], atol=1e-5)
+
+
+def test_zloss_off_matches_plain_ce():
+    """z_loss_coef=0 must be a perfect no-op vs the v1 loss path."""
+    torch.manual_seed(0)
+    cfg_off = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64,
+                          n_layers=2, n_heads=4, n_kv_heads=2, z_loss_coef=0.0)
+    model = Transformer(cfg_off).eval()
+    idx = torch.randint(0, cfg_off.vocab_size, (2, 16))
+    tgt = torch.randint(0, cfg_off.vocab_size, (2, 16))
+    with torch.no_grad():
+        logits, loss = model(idx, tgt)
+        expected = torch.nn.functional.cross_entropy(
+            logits.view(-1, logits.size(-1)), tgt.view(-1), ignore_index=-1)
+    assert torch.allclose(loss, expected, atol=1e-6)
+
+
+def test_zloss_adds_lse_square_term():
+    """With z_loss_coef>0, loss = CE + coef * mean(logsumexp(logits)^2)."""
+    torch.manual_seed(0)
+    coef = 1e-3
+    cfg = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64, n_layers=2,
+                      n_heads=4, n_kv_heads=2, z_loss_coef=coef)
+    model = Transformer(cfg).eval()
+    idx = torch.randint(0, cfg.vocab_size, (2, 16))
+    tgt = torch.randint(0, cfg.vocab_size, (2, 16))
+    with torch.no_grad():
+        logits, loss = model(idx, tgt)
+        ce = torch.nn.functional.cross_entropy(
+            logits.view(-1, logits.size(-1)), tgt.view(-1), ignore_index=-1)
+        log_z = logits.logsumexp(dim=-1)
+        z = coef * (log_z ** 2).mean()
+    assert torch.allclose(loss, ce + z, atol=1e-6)
+    assert loss.item() > ce.item()                # z-loss strictly adds to CE
+
+
+def test_base_350m_shape_validates_and_constructs():
+    """BASE_350M is the hero config; must build cleanly with the locked shape.
+    Skip if liger isn't installed (use_liger=True requires it on GPU)."""
+    cfg = BASE_350M
+    assert cfg.head_dim == 80                     # 960 / 12
+    assert cfg.n_heads % cfg.n_kv_heads == 0      # GQA divides
+    assert cfg.tie_weights and cfg.qk_norm
+    assert cfg.z_loss_coef > 0
+    # Hero config has use_liger=True; on a CPU dev box without liger installed,
+    # this asserts the import-error message rather than building the giant model.
+    if cfg.use_liger:
+        try:
+            import liger_kernel  # noqa: F401
+        except ImportError:
+            pytest.skip("liger-kernel not installed (expected on CPU dev boxes)")
+    # If we reach here, liger is present; param count check is the cheap part.
+    model = Transformer(cfg)
+    n = model.num_params(non_embedding=True)
+    assert 280_000_000 < n < 360_000_000, \
+        f"non-embed params {n:,} outside the 280M-360M sub-1B target band"
+
+
+@pytest.mark.slow
+def test_overfit_single_batch():
+    """The model must drive loss toward zero on one fixed batch."""
+    cfg = ModelConfig(vocab_size=256, max_seq_len=64, d_model=128,
+                      n_layers=2, n_heads=4, n_kv_heads=2)
+    model = Transformer(cfg).train()
+    torch.manual_seed(0)
+    idx = torch.randint(0, cfg.vocab_size, (4, 32))
+    tgt = torch.randint(0, cfg.vocab_size, (4, 32))
+    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
+    losses = []
+    for _ in range(300):
+        _, loss = model(idx, tgt)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step()
+        losses.append(loss.item())
+    assert losses[-1] < 0.1, f"failed to overfit; final loss={losses[-1]:.3f}"

tests/test_optim.py

@@ -0,0 +1,122 @@
+"""Muon optimizer + Muon/AdamW hybrid."""
+
+import pytest
+import torch
+
+from matilda import Transformer, ModelConfig
+from matilda.optim import (
+    build_optimizer, cosine_warmup_scheduler, build_scheduler,
+    wsd_scheduler, WarmupStableDecay, Muon, HybridOptimizer,
+    zeropower_via_newtonschulz5,
+)
+
+
+def test_newtonschulz_orthogonalizes():
+    torch.manual_seed(0)
+    G = torch.randn(32, 16)
+    X = zeropower_via_newtonschulz5(G, steps=5).float()
+    # singular values should be pushed toward 1
+    s = torch.linalg.svdvals(X)
+    assert (s > 0.5).all() and (s < 1.5).all()
+
+
+def test_hybrid_splits_params():
+    cfg = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64,
+                      n_layers=2, n_heads=4, n_kv_heads=2)
+    opt = build_optimizer(Transformer(cfg), name="muon")
+    assert isinstance(opt, HybridOptimizer)
+    muon, adamw = opt.optimizers
+    assert isinstance(muon, Muon)
+    # every Muon param is a 2-D matrix
+    assert all(p.ndim == 2 for g in muon.param_groups for p in g["params"])
+
+
+@pytest.mark.slow
+def test_muon_overfits_single_batch():
+    cfg = ModelConfig(vocab_size=256, max_seq_len=64, d_model=128,
+                      n_layers=2, n_heads=4, n_kv_heads=2)
+    model = Transformer(cfg).train()
+    torch.manual_seed(0)
+    idx = torch.randint(0, cfg.vocab_size, (4, 32))
+    tgt = torch.randint(0, cfg.vocab_size, (4, 32))
+    opt = build_optimizer(model, name="muon", lr=3e-3, muon_lr=0.02)
+    sched = cosine_warmup_scheduler(opt, warmup_steps=10, total_steps=300)
+    last = None
+    for _ in range(300):
+        _, loss = model(idx, tgt)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step()
+        sched.step()
+        last = loss.item()
+    assert last < 0.5, f"Muon failed to overfit; final loss={last:.3f}"
+
+
+def _dummy_opt(base_lr=1.0):
+    """Minimal optimizer just to exercise scheduler interface."""
+    p = torch.nn.Parameter(torch.zeros(2, 2))
+    return torch.optim.SGD([p], lr=base_lr)
+
+
+def test_wsd_schedule_three_phases():
+    """Warmup ramps linearly to peak, stable holds peak, decay drops to min."""
+    base = 1.0
+    warmup, total, stable = 10, 100, 0.8
+    sched = wsd_scheduler(_dummy_opt(base), warmup_steps=warmup,
+                          total_steps=total, stable_share=stable,
+                          min_lr_ratio=0.1)
+    lrs = []
+    for _ in range(total):
+        lrs.append(sched.get_last_lr()[0])
+        sched.step()
+    # warmup: lr at step warmup-1 should equal base (peak)
+    assert lrs[0] == pytest.approx(base / warmup, rel=1e-6)
+    assert lrs[warmup - 1] == pytest.approx(base, rel=1e-6)
+    # stable phase: held at peak for stable_share of post-warmup
+    decay_start = warmup + int(stable * (total - warmup))
+    assert lrs[warmup] == pytest.approx(base, rel=1e-6)
+    assert lrs[decay_start - 1] == pytest.approx(base, rel=1e-6)
+    # decay: monotone non-increasing
+    decay_lrs = lrs[decay_start:]
+    assert all(decay_lrs[i] >= decay_lrs[i + 1] for i in range(len(decay_lrs) - 1))
+    # ends at min_lr_ratio * base (within one-step linear quantum)
+    assert lrs[-1] < base                          # has decayed
+    assert lrs[-1] >= 0.1 * base * 0.5             # not below floor
+
+
+def test_wsd_state_dict_roundtrip():
+    """Resume must restore step + base_lrs and re-apply the LR exactly."""
+    opt = _dummy_opt(1.0)
+    s1 = wsd_scheduler(opt, warmup_steps=10, total_steps=100, stable_share=0.8)
+    for _ in range(50):
+        s1.step()
+    sd = s1.state_dict()
+    s2 = wsd_scheduler(_dummy_opt(1.0), warmup_steps=10, total_steps=100,
+                       stable_share=0.8)
+    s2.load_state_dict(sd)
+    assert s2.last_step == s1.last_step
+    assert s2.get_last_lr()[0] == pytest.approx(s1.get_last_lr()[0], rel=1e-9)
+
+
+def test_build_scheduler_dispatch():
+    assert isinstance(
+        build_scheduler("wsd", _dummy_opt(), 10, 100), WarmupStableDecay)
+    cos = build_scheduler("cosine", _dummy_opt(), 10, 100)
+    assert hasattr(cos, "step") and hasattr(cos, "get_last_lr")
+    with pytest.raises(ValueError):
+        build_scheduler("bogus", _dummy_opt(), 10, 100)
+
+
+def test_hybrid_state_dict_roundtrip():
+    cfg = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64,
+                      n_layers=2, n_heads=4, n_kv_heads=2)
+    model = Transformer(cfg)
+    opt = build_optimizer(model, name="muon")
+    # take a step so state is populated
+    _, loss = model(torch.randint(0, 128, (2, 16)), torch.randint(0, 128, (2, 16)))
+    loss.backward()
+    opt.step()
+    sd = opt.state_dict()
+    opt2 = build_optimizer(Transformer(cfg), name="muon")
+    opt2.load_state_dict(sd)        # must not raise
+    assert len(sd["opts"]) == 2

tests/test_run.py

@@ -0,0 +1,52 @@
+"""The launch entrypoint: configs parse, build, override, and run."""
+
+import sys
+import json
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT))
+
+import run  # noqa: E402
+
+
+def _load(name):
+    return json.loads((ROOT / "configs" / name).read_text())
+
+
+def test_shipped_configs_build():
+    for name in ("calibration.json", "base_124m.json"):
+        mcfg, tcfg = run.build(_load(name))
+        assert mcfg.d_model == 768 and mcfg.n_layers == 12
+        assert tcfg.seq_len == 1024
+
+
+def test_base_config_targets_about_3B_tokens():
+    _, tcfg = run.build(_load("base_124m.json"))
+    tokens = tcfg.total_steps * tcfg.batch_size * tcfg.grad_accum * tcfg.seq_len
+    assert 2.5e9 < tokens < 3.5e9
+
+
+def test_overrides_coerce_types():
+    cfg = {"train": {"batch_size": 24}}
+    out = run.apply_overrides(cfg, ["train.batch_size=48", "train.compile=true",
+                                    "train.lr=0.0006"])
+    assert out["train"]["batch_size"] == 48 and isinstance(
+        out["train"]["batch_size"], int)
+    assert out["train"]["compile"] is True
+    assert out["train"]["lr"] == 0.0006
+
+
+def test_dry_run_builds_synthetic_and_steps(tmp_path):
+    cfg = _load("calibration.json")
+    cfg = run.apply_overrides(cfg, [
+        "model.d_model=64", "model.n_layers=2", "model.n_heads=4",
+        "model.n_kv_heads=2", "model.vocab_size=256", "model.max_seq_len=64",
+        "train.total_steps=4", "train.warmup_steps=1", "train.batch_size=4",
+        "train.seq_len=64", "train.device=cpu", "train.dtype=float32",
+        "train.compile=false", f"train.ckpt_dir={tmp_path.as_posix()}",
+    ])
+    mcfg, tcfg = run.build(cfg)
+    stream = run.build_stream(mcfg, tcfg, data_dir=None, dry_run=True)
+    from matilda.train import Trainer
+    assert Trainer(mcfg, tcfg, stream).train() == 4

tests/test_train.py

@@ -0,0 +1,79 @@
+"""Training-loop reliability guards (all CPU-testable)."""
+
+import os
+import glob
+
+import torch
+import pytest
+
+from matilda import ModelConfig
+from matilda.data import SyntheticStream
+from matilda.train import Trainer, TrainConfig
+from matilda.monitor import mfu, peak_tflops
+
+MCFG = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64,
+                  n_layers=2, n_heads=4, n_kv_heads=2)
+
+
+def _make(tmp_path, **over):
+    kw = dict(total_steps=20, warmup_steps=2, batch_size=4, seq_len=16,
+              log_every=5, ckpt_every=10, keep_last=2, device="cpu",
+              dtype="float32", ckpt_dir=str(tmp_path))
+    kw.update(over)
+    tc = TrainConfig(**kw)
+    stream = SyntheticStream(MCFG.vocab_size, tc.batch_size, tc.seq_len, seed=0)
+    return Trainer(MCFG, tc, stream)
+
+
+def test_loop_runs_and_checkpoints(tmp_path):
+    t = _make(tmp_path)
+    final = t.train()
+    assert final == 20
+    # final checkpoint on clean completion + rotation kept <= keep_last
+    ckpts = glob.glob(os.path.join(tmp_path, "ckpt_*.pt"))
+    assert len(ckpts) <= 2
+    assert os.path.exists(os.path.join(tmp_path, "ckpt_20.pt"))
+
+
+def test_loop_resume_continues(tmp_path):
+    # run only 10 steps, leaving a checkpoint at step 10
+    t1 = _make(tmp_path, total_steps=10, ckpt_every=10)
+    assert t1.train() == 10
+    # a fresh trainer in the same dir must resume from step 10, not restart
+    t2 = _make(tmp_path, total_steps=15, ckpt_every=10)
+    assert t2.maybe_resume() is True
+    assert t2.step == 10
+    assert t2.train() == 15
+
+
+def test_nan_guard_aborts_after_max_skips(tmp_path):
+    t = _make(tmp_path, total_steps=50, max_skips=3)
+    nan = torch.tensor(float("nan"))
+    t.model.forward = lambda x, targets=None: (None, nan)  # shadow bound method
+    with pytest.raises(RuntimeError, match="non-finite"):
+        t.train()
+    assert t.step == 0  # never advanced past a bad batch
+
+
+def test_nan_guard_skips_then_recovers(tmp_path):
+    t = _make(tmp_path, total_steps=5, max_skips=10)
+    real_forward = t.model.forward
+    calls = {"n": 0}
+
+    def flaky(x, targets=None):
+        calls["n"] += 1
+        if calls["n"] <= 2:                     # first two micro-batches are bad
+            return None, torch.tensor(float("inf"))
+        return real_forward(x, targets)
+
+    t.model.forward = flaky
+    assert t.train() == 5                        # recovered and finished
+    assert t.consecutive_skips == 0
+
+
+def test_mfu_sanity():
+    # 100M params, 500k tokens/step, 2.0s, A100 -> ~0.48 MFU
+    val = mfu(100_000_000, 500_000, 2.0, peak_tflops("A100") * 1e12)
+    assert 0.0 < val < 1.0
+    assert peak_tflops("A100") == 312.0
+    assert peak_tflops("totally-unknown-gpu") == 312.0  # default