Axolotl hardens SinkGD for width-stable fine-tuning

What changed

axolotl-ai-cloud/axolotl landed a substantial SinkGD optimizer update in commit 6ccb24e.

The commit expands SinkGD beyond a plain scalar learning-rate scale. The important pieces:

• Width-aware LR transfer via sinkgd_base_width and sinkgd_lr_width_exponent. Tune sinkgd_lr_scale at one hidden size, then scale 2D-linear updates as roughly 1/d_in when you move to a wider model. Existing configs stay unchanged if you leave the new knobs unset.
• Spectral normalization via sinkgd_spectral_norm and sinkgd_spectral_target. The muon target pins operator norm to sqrt(d_out/d_in) and is intended to own width transfer; the default unit target is a conditioning stabilizer that can sit on top of sinkgd_base_width. Using both sinkgd_base_width and spectral_target: muon is rejected, because that double-counts width correction.
• Fused Triton kernels (sinkgd_fused_kernel: true), reported at about 1.4–1.75× on the optimizer step single-GPU on B200, and up to about 1.8× on 2-rank FSDP2, with automatic fallback when Triton, sharding, or problem size make the compiled path better.
• Experimental MD sphere (sinkgd_md_sphere) that holds SinkGD-routed 2D weights on a fixed Frobenius sphere. Tight activation growth control, but width-dependent LR and off by default.

Docs and tests moved with the code: optimizer docs, Triton kernels, FSDP2 multigpu coverage, and unit tests for the new behaviour.

Why it matters

Fine-tuning still dies of boring problems: an LR that worked on a 1–2B probe collapses or undertrains when you jump to a wider model, or the optimizer step becomes the free lunch you never get because it is too slow under FSDP.

SinkGD already sat in the interesting middle of modern matrix-aware optimizers. This change makes that path less of a “remember the folklore” exercise:

1. You can encode the width you tuned against.
2. You can condition updates with spectral norm instead of hand-waving about scale.
3. You can spend less wall-clock on the optimizer step when fused kernels apply.

For agent and product work, that matters when the loop is “ship a specialist model, not a science project.” Fewer wasted full runs means faster iteration on domain data, tool-use traces, and client-specific behaviour.

My read

Worth a spike if Axolotl + SinkGD is already on the training path. Watch only if you are still on plain AdamW/Muon defaults and not touching SinkGD.

Sensible test plan:

1. Pin current Axolotl; train a known small SinkGD config as baseline.
2. Enable sinkgd_base_width at the width you originally tuned, or sinkgd_spectral_norm with spectral_target: muon — not both width-owning modes.
3. Repeat at a second model width with the same data recipe.
4. Compare loss curves, final eval, and step time with/without sinkgd_fused_kernel.
5. Only then promote the knobs into shared training configs.

Ignore the MD-sphere path unless you are deliberately chasing activation stability and have budget for LR retuning.

Bottom line

Axolotl just made SinkGD more usable as a production fine-tuning optimizer, not just a clever default. If you move model width often, spike the new transfer and spectral-norm options; if you do not use SinkGD, leave the pin alone and keep watching.