optimizer/COMPARISON.md

# Optimizer Comparison — Llama 77M pretraining, 60/ 120 min each

**Hardware:** RX 9070 XT 16GB, ROCm 7.2, torch 2.13, transformers 5.14.1, bitsandbytes 0.50.0
**Common recipe:** effective batch 256×1024 tok (micro 8 × accum 32), cosine→5% floor, warmup 200, bf16 autocast + fp32 master weights, seed 42, same data order. Per-folder details in each `*/REPORT.md`.

## Headline table

| optimizer | steps/60min | final eval ppl | held-out bpb | bake /100 | peak GPU | verdict |
|---|---|---|---|---|---|---|
| **adamw_torch** (baseline) | **848** | **62.4** | **1.4044** | 39 | 7.6 GB | reference — solid |
| adamw_torch_fused | 785 | 65.4 | 1.4215 | 39 | 7.6 GB | identical math, no win on ROCm |
| adamw_bnb_8bit | 783 | 65.6 | 1.4224 | 39 | 7.2 GB | quality-free 8-bit states |
| **ademamix_8bit** | 785 | 65.2 | 1.4208 | 39 | 7.4 GB | **only one to beat AdamW per step** |
| lion_8bit | 795 | 149.7 | 1.6900 | 22 | 7.1 GB | far behind at safe LR — skip |
| gefen_x | 783 | 77.8 | 1.4790 | 33 | **7.0 GB** | works on ROCm; quality gap — skip at this scale |
| **muonq** | 739 | **55.3** | **1.3657** | **42** | 7.2 GB | **series winner — best everything despite fewest steps** |

Final-ppl/bpb columns partly reflect tokens seen (baseline ran +8% more steps — see caveat below). The fair signal is quality at *matched* step:

## Quality at matched steps (eval loss; lower better)

| step ≈ | baseline | fused | bnb-adam | ademamix | lion* | gefen_x* | **muonq*** |
|---|---|---|---|---|---|---|---|
| 330 | 4.7845 | 4.7755 | 4.7786 | 4.7728 | 5.71 | 4.9730 | **4.49** (@315) |
| 395 | 4.6056 | 4.6127 | 4.6102 | 4.5977 | 5.56 | 4.7985 | **4.35** (@377) |
| 525 | 4.3906 | 4.4056 | 4.3862 | 4.3809 | 5.31 | 4.5690 | **4.17** (@499) |
| 720 | 4.1832 | 4.1808 | 4.1836 | 4.1776 | 5.01 | 4.3547 | **4.01** (@683) |

\* Lion, Gefen, and MuonQ use their own documented lr/wd recipes; not equal-settings comparisons. MuonQ's eval steps don't align with the others (fewest steps/hour), so its column shows the *earlier* step in parentheses — it leads by −0.17…−0.29 even with that handicap.

- fused and bnb-adam are statistically identical to baseline (|Δ|≤0.02, no consistent sign) → 8-bit Adam states cost **nothing**.
- **AdEMAMix is ahead at every single eval point** (−0.006…−0.012). Small in one hour, but the AdEMAMix paper predicts this gap *grows* with training length — its slow EMA (β₃≈0.9999) compounds. On a multi-day full run this is the one worth betting on.

## Recommendation

1. **For the real 19.6B-token run: `muonq`** — won every quality metric by a wide margin (final ppl 55.3 vs 62.4; −0.17…−0.29 eval loss at matched steps; series-best bake 42/100 and bpb 1.3657) despite the fewest steps/hour. Caveats: research-grade code, single 1-hour seed, and an interrupted-resume should be verified before a multi-day run — details in `muonq/REPORT.md`. **Suggested gate: a 3-hour muonq vs ademamix_8bit head-to-head before committing.** → *Update 2026-08-13: gate run as a 2-hour head-to-head — muonq confirmed, and ademamix picked up a stability black mark. See "V2 follow-up" at the end.*
2. **Conservative / battle-tested choice: `ademamix_8bit`.** Only *Adam-family* optimizer that beat AdamW at matched steps, systematically; stable with standard AdamW hyperparameters; 8-bit states. Its paper predicts the small edge grows with run length.
3. `adamw_bnb_8bit` — use when VRAM is the binding constraint (it isn't at 77M).
4. Skip `adamw_torch_fused` (no benefit on this ROCm stack) and `lion_8bit` (needs long runs and tuning to compete; badly behind in an hour).
5. `gefen_x` (quantized ~1 B/param AdamW drop-in, run 6) proved the ROCm question — its pure-PyTorch fallback trains stably at full speed and the lowest peak memory of the series (7.0 GB) — but trails AdamW by +0.17–0.19 eval loss at matched steps with its documented 0.6× LR recipe. Its memory edge only matters at model sizes far beyond 77M. Skip for this project; details in `gefen_x/REPORT.md`.

## Caveats

- **Throughput anomaly:** runs 2–5 all landed at 783–795 steps (±1%), the baseline got 848 (+8%). GPU clocks were verified at max (3058 MHz, 74 °C junction) mid-series, ruling out thermal throttling; the baseline's edge is an unexplained first-run effect (allocator/page-cache/background state). Treat "steps/hour vs baseline" with skepticism; comparisons among runs 2–5 are clean.
- All runs are 1-hour snapshots (~1% of the corpus); bake scores of 39/100 ("DOUGH") are on-trend, not a quality ceiling.
- Batch geometry is 8×32 (not the original 16×16) because the 13.5 GB footprint OOMed at eval boundaries on 16 GB; identical across all runs. Side discovery: micro-batch 8 was also *faster* per token than 16 on this GPU (~3.0 vs ~4.5 s/step average in early testing).
- One eval every 5 min steals ~35–40 s/hour of training each run, identically.

## Ops notes (for future runs)

- `torch.cuda.empty_cache()` mid-training is **harmful** on this ROCm build: released memory did not come back to the process (~6 GB stranded), turning a fragmentation stumble into a hard OOM. Callback was removed from `base_train.py`.
- zsh: `>` won't overwrite (noclobber) — use `>|`; unmatched globs abort the whole command.

## V2 follow-up: 2-hour head-to-head, muonq vs ademamix_8bit

The recommended gate before the real run: the top two optimizers, re-run for **120 min** with **warmup 400** on the **fixed dataset** (train slightly bigger, validation slightly smaller — v2 numbers are not directly comparable to the 1-hour tables above). Full details in `muonq-v2/REPORT.md` and `ademamix_8bit-v2/REPORT.md`.

| | muonq-v2 | ademamix_8bit-v2 |
|---|---|---|
| steps in 120 min | 1 497 (748/h) | 1 556 (778/h) |
| final eval ppl | **43.6** | 48.7 (best, @1438) |
| eval loss at matched steps | **−0.11…−0.20** | — |
| held-out bpb | **1.2852** | 1.3271 |
| bake /100 | **49.5** | 46.5 |
| chat bpb | **1.1716** | 1.2297 |
| peak GPU | **7.2 GB** | 7.4 GB |
| stability | spotless (grad max 5.7) | **grad spike @step ~1490** (norm 37, clipped; eval ppl 48.7→91; recovered in ~50 steps) |

**muonq confirmed — it wins every metric, and its matched-step lead over ademamix (−0.11…−0.20) is the same magnitude as its v1 lead over AdamW.** The margin did not shrink with doubled training length.

**New negative signal for ademamix:** its late-run gradient spike happened on a data region muonq crossed with grad_norm 1.03 (same seed ⇒ identical batch order), so the batch is innocent — the slow-EMA state is the suspect. For a multi-day run that risk now weighs against it.

**Decision: use muonq for the real 19.6B-token run**, after verifying trainer checkpoint-resume on MuonQ state (one interrupted muonq-v2 attempt showed HF auto-resume *engages* silently from a leftover checkpoint — it worked, but verify correctness deliberately, and mind that silent auto-resume when re-running experiments in a folder with old checkpoints).