Skip to content

feat(trainer): Nesterov momentum for Muon by default - #4113

Merged
samsja merged 1 commit into
mainfrom
sebastian/muon-nesterov-default
Oct 17, 2026
Merged

samsja merged 1 commit into
mainfrom
sebastian/muon-nesterov-default

Conversation

@snimu

@snimu snimu commented Oct 16, 2026 •

Copy link
Copy Markdown
Collaborator

Turns on Nesterov momentum for Muon by default.

What changes

  • MuonConfig gets nesterov: bool = True. It is passed to dion's Muon, whose own default is False.
  • This changes the default for every SFT and RL run that uses Muon. Set optim.nesterov = false to get the old behavior.
  • Resuming works: the momentum buffers are the same with and without Nesterov; only the update read from them differs.

Why

  • In our pretraining runs (MoE, d=384, Muon), Nesterov lowered the final loss by 0.0128 ± 0.0017 nats, paired over 4 seeds and several learning rates. These are internal, unpublished experiments.
  • Keller Jordan's Muon and Moonlight use Nesterov momentum.
  • dion's own reference implementation (muon_reference) defaults to nesterov=True.

The effect on SFT and RL fine-tuning has not been tested.

Validation

  • Ruff passes. The config and SFT unit tests pass (366).
  • SFT smoke run with Muon on 1 GPU (configs/debug/fake/sft.toml, Qwen3-0.6B, lr 1e-3, 20 steps, H200): trains without errors, nesterov = true in the resolved config, loss 8.46 at step 1 and 6.61 at step 20 (noisy on the fake data: 4.80 at step 12, 8.37 at step 14).

MuonConfig gets nesterov (default True), passed to dion's Muon, whose own default is False.
@snimu
snimu marked this pull request as ready for review October 16, 2026 08:39
@samsja
samsja merged commit 884400a into main Oct 17, 2026
25 of 26 checks passed
@samsja
samsja deleted the sebastian/muon-nesterov-default branch October 17, 2026 06:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants