matilda-mini-v2

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.

Three configs, one stack

The v1.5 hero run (152M) is designed around two claims: (1) direct token- matched comparison vs Pythia-160M (same width, 50% more depth, better data, better optimizer); (2) approaching Pythia-410M at the same token count despite being 2.7× smaller — the "beats above weight" framing. Eight things differ from v1: deeper architecture (18L vs 12L), seq_len 2048 (vs 1024), Muon hybrid (vs AdamW), WSD schedule (vs cosine), Liger fused linear+CE, z-loss, ReLU² FFN (vs SwiGLU), final logit soft-cap.

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:

Training infrastructure (the actual deliverable)

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:

Target (124M, ~3B tokens, vs Pythia-160M): HellaSwag ~30-35%, ARC-easy ~40-45%, PIQA ~60%.

Quickstart

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

# v1.5 hero run (152M / SmolLM 83/17 / ~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_152m_v2.json --data-dir data/smollm_mix

Full GPU procedure (validate → calibrate → ablate → train → eval) is in 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.

pytest tests/ -q