[faresobeid]

Stacked on #4117.

Two opt-in Muon knobs, both from the DeepSeek-V4 report (section 2.4, Algorithm 1). Defaults are unchanged and keep dion's own Newton-Schulz function.

• optim.newton_schulz_coefficients = "deepseek_v4": 8 Newton-Schulz steps with (3.4445, -4.7750, 2.0315), then 2 steps with (2, -1.5, 0.5). Default "muon" is dion's 5-step schedule.
• optim.update_rms (default 0.2; DeepSeek-V4 uses 0.18): the RMS of a Muon update per unit of learning rate.
  • dion's rms_norm adjustment fixes this at 0.2 (Moonlight).
  • The new Newton-Schulz function scales its output by update_rms / 0.2, so weight decay (1 - lr * wd) is unchanged.
• Both use trainer/optim/newton_schulz.py, which is dion's bfloat16 iteration with the schedule as an argument. With "muon" and 0.2 it is bitwise equal to dion's function (checked on CPU). It is only installed when one of the knobs differs from its default.

Measurements

Newton-Schulz alone on 1 H200, inputs with singular values spread over 1..1e-3:

shape	dion 5-step	deepseek_v4 10-step	singular values of output (dion / deepseek_v4)
16x1408x4096 (GLM-4.5-Air experts)	6.1 ms	11.4 ms	0.07-1.15 / 0.998-1.003
12288x4096 (GLM-4.5-Air q_proj)	7.6 ms	14.9 ms	0.03-1.15 / 0.998-1.003
16x768x2048 (Qwen3-30B experts)	1.1 ms	2.1 ms	0.07-1.15 / 0.998-1.003

The deepseek_v4 schedule roughly doubles Newton-Schulz time. On GLM-4.5-Air (2 nodes) all of Muon's step after the dion bump takes 0.63 s, so the cost is at most a few tenths of a second per step.

Loss A/B: Qwen3-0.6B SFT on PrimeIntellect/Reverse-Text-SFT, 1 GPU, Muon lr 5e-5, batch 32, seed 0. Mean train loss per window:

run	steps 1-10	steps 11-31	steps 32-62
default	4.787	2.255	0.286
default, rerun	4.787	2.254	0.286
deepseek_v4	4.647	2.047	0.233
update_rms = 0.18	4.843	2.437	0.430
deepseek_v4 + update_rms = 0.18	4.726	2.224	0.292

• deepseek_v4 pushes small singular values up to 1, which gives a larger effective step. That is consistent with the lower early loss.
• 0.18 is a 10% smaller step.
• Together they land near the default. This run measures the step size, not final quality.

Checks

• CPU: with the defaults the function equals dion's exactly. With deepseek_v4 and 0.18 the output singular values are 0.896-0.903 (expected 0.9).
• tests/unit/orchestrator + tests/unit/test_configs.py -m "not gpu": 350 passed. ruff check, ruff format --check.

---

Note

Low Risk
Training-path change is opt-in with unchanged defaults; only Muon’s orthogonalization step is affected when non-default config is set.

Overview
Adds opt-in Muon optimizer knobs aligned with DeepSeek-V4 while keeping defaults on dion’s built-in Newton-Schulz path.

MuonConfig gains newton_schulz_coefficients ("muon" vs "deepseek_v4") and update_rms (default 0.2). A new trainer/optim/newton_schulz.py implements bfloat16 Newton-Schulz with selectable coefficient schedules and output scaling so update_rms can differ from Moonlight’s fixed 0.2 without changing weight decay behavior.

Muon construction only passes a custom newton_schulz_func into dion when either knob differs from default; otherwise behavior stays on dion’s native function. Choosing "deepseek_v4" uses a longer 10-step schedule (roughly ~2× Newton-Schulz cost) aimed at singular values closer to 1.

^{Reviewed by Cursor Bugbot for commit 4aaa0a8. Bugbot is set up for automated code reviews on this repo. Configure here.}