Cormac Shannon
84b04d12e9
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
74 lines
3.7 KiB
Markdown
74 lines
3.7 KiB
Markdown
# 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)
|
||
|
||
```sh
|
||
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
|
||
|
||
```sh
|
||
.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.
|