Files
Cormac Shannon 84b04d12e9 Tech spike: four smoke-test runs and a divergence-stop test
Pipeline works on MPS (~105 s per 300-step run), results are reproducible,
and the divergence stop works. Not experimental evidence.

Claude-Session: https://claude.ai/code/session_01S8KN7NuUCm8nXLmy6eDuVv
2026-10-28 18:14:26 +00:00

3.7 KiB
Raw Permalink Blame History

06-experiments

Plain nanoGPT + TinyStories on an M4 Pro (PyTorch MPS), with QK-norm as a single on/off switch and every run logged automatically.

Path What
preregistration.md Hypotheses and how each will be tested, written before the main runs
run.py Launch a run. Writes runs/<id>/run.json + metrics and adds a row to runs/index.csv
configs/ Run configs (spike.py = tech-spike smoke test only)
nanogpt/ Vendored nanoGPT; README-thesis.md lists every change from upstream
data/prepare_tinystories.py Downloads + tokenises TinyStories, records SHA-256 checksums
runs/ One folder per run (run.json, config.json, metrics.jsonl tracked; checkpoints and raw stdout ignored)
runs/index.csv One row per run: status, QK-norm, LR, depth, seed, best val loss, diverged-at, commit

Setup (once)

cd 06-experiments
uv venv --python 3.12 .venv && uv pip install --python .venv/bin/python -r requirements.txt
.venv/bin/python data/prepare_tinystories.py --spike   # 22 MB, ~1 min
.venv/bin/python data/prepare_tinystories.py           # full dataset, ~2.2 GB download

Running

.venv/bin/python run.py my-run-name configs/spike.py qk_norm=True learning_rate=3e-3 seed=1 \
    --note "why I'm running this"

Habit: commit before running. run.py records the commit SHA. If there are uncommitted changes, it saves them as diff.patch in the run folder, so a result can always be traced back to the exact code, but a clean commit is easier to cite.

What gets logged (run.json)

name, note, status (completed / diverged / failed / interrupted), start/end time, exact command, requested + resolved config (LR, size, iters, batch, seed, qk_norm…), git commit + dirty flag, upstream nanoGPT commit, dataset name + SHA-256 of raw files + token counts, machine / macOS / Python / torch versions, and a summary (param count, first/last/best loss, iteration where it diverged).

Tech-spike results (28/10/2026)

See runs/index.csv. Pipeline check only: 4 layers × 256 wide, 300 steps, 4.9M-token subset, one seed. Not evidence for or against any hypothesis.

Run QK-norm LR Best val loss Time
spike-qkoff-lr1e-3 off 1e-3 3.459 103 s
spike-qkon-lr1e-3 on 1e-3 3.430 105 s
spike-qkoff-lr1e-1 off 1e-1 5.295 103 s
spike-qkon-lr1e-1 on 1e-1 3.903 106 s
test-divergence-stop off 1e-1 (stopped) 6 s

What the spike showed:

  • Training runs on MPS at ~0.29 s/step for this size; the whole pipeline works.
  • Reproducible: re-running QK-off at 1e-3 gave the identical val loss (3.4590) as a pre-commit test run.
  • At LR 0.1, QK-off was clearly worse but did not diverge: grad clipping (1.0) kept loss finite. So "diverged" needs a definition beyond NaN for the pre-registration (e.g. final loss vs a threshold).
  • The divergence stop works (tested with an artificially low threshold; the run was killed and logged).
  • The QK-on vs off direction at high LR matches H2, but this is one seed, 300 steps: a hint for the pilot, nothing more.

Known issues / decisions still open

  • Vocabulary size dominates the parameter count. With GPT-2's 50k BPE vocab, a 4-layer 256-wide model has ~16M params, ~13M of them in the embedding. For a "10–15M param" model the non-embedding count is far smaller than it sounds. [Your call: report non-embedding params (as Kaplan et al. and Wortsman et al. do), and/or train a smaller tokenizer on TinyStories.]
  • MPS runs in float32 (nanoGPT only autocasts on CUDA). Fine for this study; note it in methods.
  • nanoGPT's estimate_mfu assumes an A100, so the mfu number in the logs is meaningless on a Mac.