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).