fix(muon): muon_use_nesterov silently dropped (arg/config name mismatch)

[yuchenwang3]

What

ms-swift exposes a user-facing Muon argument muon_use_nesterov (swift/megatron/arguments/megatron_args.py), but Megatron's OptimizerConfig field is named muon_nesterov. The optimizer config is built in swift/megatron/trainers/base.py by copying kwargs by the config field name:

kwargs = {
    f.name: getattr(args, f.name)
    for f in dataclasses.fields(config_cls) if hasattr(args, f.name) and f.name != 'loss_scale'
}
config = config_cls(**kwargs)

Because the loop keys off the Megatron field name (muon_nesterov) and checks hasattr(args, 'muon_nesterov') — which is False (the arg is muon_use_nesterov) — the user's setting is silently dropped. --muon_use_nesterov true is a no-op: muon_nesterov stays at its default False.

All of the other muon_* args (muon_momentum, muon_split_qkv, muon_scale_mode, muon_coefficient_type, ...) match the Megatron OptimizerConfig field names 1:1 and transfer correctly through this same loop — muon_use_nesterov is the only one whose name diverges.

Change

Add a one-line alias in _check_muon (inside the if 'muon' in self.optimizer: block) so the user's value reaches Megatron:

self.muon_nesterov = self.muon_use_nesterov

This is additive and non-breaking (the public CLI arg --muon_use_nesterov is unchanged).

Testing

py_compile clean. Traced the consumption path (megatron_args arg → _check_muon → trainers/base.py field-name copy → Megatron OptimizerConfig.muon_nesterov). Confirmed muon_use_nesterov has no other consumer in the tree, and muon_nesterov is the actual Megatron field (megatron/core/optimizer/optimizer_config.py).

Note

Alternative: rename the arg muon_use_nesterov → muon_nesterov to match the 1:1 convention of the other muon_* args. That is cleaner but a breaking CLI change, so I went with the additive alias. Happy to switch to the rename if you prefer.

[gemini-code-assist]

Code Review

This pull request updates swift/megatron/arguments/megatron_args.py to alias the user-facing argument muon_use_nesterov to self.muon_nesterov. This ensures that the setting is correctly passed to Megatron's OptimizerConfig and not silently dropped. There are no review comments, and I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Jintao-Huang]

thanks

It looks like a breaking change from mcore 0.16 to 0.17.