# Matilda-Mini A sub-200M-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**. ## 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, at ~124M params: | Component | Choice | |-----------|--------| | Positions | **RoPE** (rotary, half-rotation convention) | | Normalization | **RMSNorm**, pre-norm, fp32 reduction | | MLP | **SwiGLU** (2/3 sizing) | | Attention | **GQA** (12 query / 4 KV heads) + **QK-Norm** | | Stability | residual-projection init scaled 1/√(2·n_layers); optional attn logit soft-cap | | Tying | embedding ↔ LM head | | Size | 114M total / 75.5M non-embedding · d=768 · 12 layers · seq 1024 · GPT-2 50k vocab | ## 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 # 30 tests: correctness, resume, NaN guard, data integrity # 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 # real run (after tokenizing data — see docs/INSTANCE_RUNBOOK.md) python scripts/prepare_data.py --out-dir data/fwedu --target-tokens 3000000000 python run.py --config configs/base_124m.json --data-dir data/fwedu ``` 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 (tokenize), ablate.py (experiments), launch_vast.sh configs/ calibration.json (MFU tuning), base_124m.json (the run) tests/ 30 tests — model, checkpoint, train loop, data, optim, ablation, run docs/ INSTANCE_RUNBOOK.md (operating manual) run.py training entrypoint (--config + --set overrides) ``` ## Testing 30 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. ```bash pytest tests/ -q ```
