feat(megatron): expose muon_coefficient_type for the Muon optimizer - #9591
Conversation
ms-swift exposes 8 Muon hyperparameters but not muon_coefficient_type, which Megatron-Core supports (--muon-coefficient-type / OptimizerConfig.muon_coefficient_type, default 'quintic'). It selects the Newton-Schulz coefficient set (quintic, polar_express, ...) and affects convergence; recipes using non-default coefficients cannot currently be expressed via Megatron-SWIFT. Add the field (str, matching Megatron's CLI type/default); it forwards through the existing asdict()-based mechanism like the other muon_* args. Add zh/en docs. Signed-off-by: yuchenwang3 <eang333cms@gmail.com>
There was a problem hiding this comment.
Code Review
This pull request introduces a new command-line parameter muon_coefficient_type for the Muon optimizer, updating both the Python arguments configuration and the Chinese and English documentation. The feedback points out minor markdown formatting issues in both documentation files, specifically unbalanced backticks around the --muon-coefficient-type parameter name.
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.
| - muon_use_nesterov: 是否在内部 SGD 中使用 Nesterov 风格的动量，默认为False。 | ||
| - muon_scale_mode: Muon 优化器的缩放模式。可选为'spectral', 'unit_rms_norm', 'shape_scaling'。默认为'spectral'。 | ||
| - muon_fp32_matmul_prec: Newton-Schulz 迭代的 FP32 矩阵乘法精度，可选为'low', 'medium', 'high'。默认为'medium'。 | ||
| - muon_coefficient_type: Muon 优化器 Newton-Schulz 迭代的系数类型，传递给 Megatron 的 `--muon-coefficient-type`。可选值取决于所安装的 emerging_optimizers 版本（如'quintic', 'polar_express', 'simple', 'cans', 'aol', 'deepseekv4', 'cubic5', 'custom'）。默认为'quintic'。 |
There was a problem hiding this comment.
There is a markdown formatting issue on this line. The backticks around --muon-coefficient-type are unbalanced (two backticks at the start and one at the end: `--muon-coefficient-type`). Please update it to use single backticks for proper rendering:
- muon_coefficient_type: Muon 优化器 Newton-Schulz 迭代的系数类型，传递给 Megatron 的 `--muon-coefficient-type`。可选值取决于所安装的 emerging_optimizers 版本（如'quintic', 'polar_express', 'simple', 'cans', 'aol', 'deepseekv4', 'cubic5', 'custom'）。默认为'quintic'。| - muon_use_nesterov: Whether to use Nesterov-style momentum in the internal SGD. Default is False. | ||
| - muon_scale_mode: Scale mode for the Muon optimizer. Options include 'spectral', 'unit_rms_norm', and 'shape_scaling'. Default is 'spectral'. | ||
| - muon_fp32_matmul_prec: FP32 matrix multiplication precision for Newton-Schulz iteration. Options include 'low', 'medium', and 'high'. Default is 'medium'. | ||
| - muon_coefficient_type: Newton-Schulz coefficient type for the Muon optimizer, forwarded to Megatron's `--muon-coefficient-type`. Available options depend on the installed emerging_optimizers version (e.g. 'quintic', 'polar_express', 'simple', 'cans', 'aol', 'deepseekv4', 'cubic5', 'custom'). Default is 'quintic'. |
There was a problem hiding this comment.
There is a markdown formatting issue on this line. The backticks around --muon-coefficient-type are unbalanced (two backticks at the start and one at the end: `--muon-coefficient-type`). Please update it to use single backticks for proper rendering:
- muon_coefficient_type: Newton-Schulz coefficient type for the Muon optimizer, forwarded to Megatron's `--muon-coefficient-type`. Available options depend on the installed emerging_optimizers version (e.g. 'quintic', 'polar_express', 'simple', 'cans', 'aol', 'deepseekv4', 'cubic5', 'custom'). Default is 'quintic'.
What
Megatron-SWIFT's
MegatronArgumentsexposes 8 Muon hyperparameters (muon_momentum,muon_split_qkv,muon_use_nesterov,muon_scale_mode,muon_fp32_matmul_prec,muon_num_ns_steps,muon_tp_mode,muon_extra_scale_factor) but is missingmuon_coefficient_type, which Megatron-Core supports:--muon-coefficient-type(megatron/training/arguments.py), default'quintic'OptimizerConfig.muon_coefficient_type(megatron/core/optimizer/optimizer_config.py)muon_coefficient_typeselects the Newton–Schulz coefficient set (quintic,polar_express, …), which affects convergence. Without it exposed, recipes that rely on a non-default coefficient set (e.g.polar_express) cannot be expressed through Megatron-SWIFT.Change
muon_coefficient_type: str = 'quintic'toMegatronArguments(matching Megatron's CLI default andtype=str).No forwarding code is needed:
MegatronArgumentsis forwarded to Megatron via the existingasdict()mechanism (swift/megatron/pipelines/train/sft.py), exactly like the othermuon_*args, and Megatron already has the matching--muon-coefficient-typeargument.Notes
Typed as
str(rather than aLiteral) to mirror Megatron's own CLI, since the supported set is version-dependent and validated by Megatron at runtime (emerging_optimizers.validate_coefficient_type). Happy to switch to aLiteralif preferred. Verified against currentmain; the field was the onlymuon_*argument absent relative to Megatron's CLI.