Add --muon-coefficient-type argument for Muon optimizer by mchrzanowski · Pull Request #3927 · NVIDIA/Megatron-LM

mchrzanowski · 2026-03-18T16:26:52Z

Allow users to select the Newton-Schulz polynomial coefficient set (e.g. simple, quintic, polar_express, aol) via the new --muon-coefficient-type CLI flag. Supported types are discovered dynamically from the installed emerging_optimizers package so that upstream additions are picked up automatically without code changes.

Allow users to select the Newton-Schulz polynomial coefficient set (e.g. simple, quintic, polar_express, aol) via the new --muon-coefficient-type CLI flag. Supported types are discovered dynamically from the installed emerging_optimizers package so that upstream additions are picked up automatically without code changes. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

copy-pr-bot · 2026-03-18T16:26:56Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-03-18T16:27:05Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

Derive supported coefficient types from the public NSCoeffT Literal type via typing.get_args() rather than reading keys from the private _COEFFICIENT_SETS dict. Tests likewise avoid importing _COEFFICIENT_SETS. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

emerging_optimizers must be installed to use Muon, so there is no need for a hardcoded fallback list. get_supported_coefficient_types() now asserts the package is present and reads NSCoeffT directly. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Replace hardcoded _NS_STEPS_FOR_COEFF_TYPE mapping with dynamic discovery via get_supported_coefficient_types() (backed by NSCoeffT). Since get_coefficient_iterator cycles/repeats coefficients, a single default step count works for all types. Remove redundant duplicate test. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

deepakn94 · 2026-03-19T04:42:32Z

@@ -288,6 +311,7 @@ def lion_init_state_fn(opt, config=None):
        "use_nesterov": config.muon_use_nesterov,


BTW, should we check for certain versions of emerging_optimizers? I believe this got re-named to nesterov in the latest release.

I feel it is better to keep the code only support one version. And I think having that version be the tagged in pyproject is reasonable. Optionally we can add a global check.
We will bump main to support v0.2.0 soon(after dev refactor+bump and main2dev sync, both are finalizing).

The Lion class moved in emerging_optimizers 0.2. Gate the import behind an explicit version check so users get a clear error instead of a silent ImportError on older versions. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Tests were failing with AssertionError because they unconditionally called get_supported_coefficient_types() which requires emerging_optimizers >= 0.2. Added pytestmark skip conditions and guarded module-level collection code. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

validate_coefficient_type was unconditionally calling get_supported_coefficient_types() which asserts HAVE_EO_V02. This caused dist_checkpointing tests to fail when emerging_optimizers < 0.2 was installed, since TensorParallelMuon.__init__ always calls validate_coefficient_type. Skip validation when HAVE_EO_V02 is False, as the coefficient_type kwarg is not passed to newton_schulz_tp in that case anyway. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…>= 0.2 "quintic" is the default coefficient type supported before emerging_optimizers 0.2, so validate against it rather than skipping validation entirely. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…ient_types Now that validate_coefficient_type falls back to ("quintic",) without emerging_optimizers >= 0.2, most tests only need HAVE_EMERGING_OPTIMIZERS. Only the three tests that directly call get_supported_coefficient_types() still require HAVE_EO_V02. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…sion The coefficient_type parameter is supported in emerging_optimizers >= 0.1, so there's no need to gate it behind HAVE_EO_V02 (>= 0.2). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

mchrzanowski · 2026-03-26T03:24:15Z

/ok to test 10bf5eb

svcnvidia-nemo-ci · 2026-03-26T04:24:12Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/23577425224

Co-authored-by: root <root@pool0-0084.cm.cluster> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: root <root@pool0-0393.cm.cluster> Co-authored-by: root <root@pool0-0154.cm.cluster> Co-authored-by: root <root@pool0-0021.cm.cluster> Co-authored-by: root <root@pool0-0331.cm.cluster> Co-authored-by: root <root@pool0-0079.cm.cluster> Co-authored-by: root <root@pool0-0137.cm.cluster>

Co-authored-by: root <root@pool0-0084.cm.cluster> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: root <root@pool0-0393.cm.cluster> Co-authored-by: root <root@pool0-0154.cm.cluster> Co-authored-by: root <root@pool0-0021.cm.cluster> Co-authored-by: root <root@pool0-0331.cm.cluster> Co-authored-by: root <root@pool0-0079.cm.cluster> Co-authored-by: root <root@pool0-0137.cm.cluster> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

Co-authored-by: root <root@pool0-0084.cm.cluster> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: root <root@pool0-0393.cm.cluster> Co-authored-by: root <root@pool0-0154.cm.cluster> Co-authored-by: root <root@pool0-0021.cm.cluster> Co-authored-by: root <root@pool0-0331.cm.cluster> Co-authored-by: root <root@pool0-0079.cm.cluster> Co-authored-by: root <root@pool0-0137.cm.cluster> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

Co-authored-by: root <root@pool0-0084.cm.cluster> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: root <root@pool0-0393.cm.cluster> Co-authored-by: root <root@pool0-0154.cm.cluster> Co-authored-by: root <root@pool0-0021.cm.cluster> Co-authored-by: root <root@pool0-0331.cm.cluster> Co-authored-by: root <root@pool0-0079.cm.cluster> Co-authored-by: root <root@pool0-0137.cm.cluster>

mchrzanowski requested review from a team as code owners March 18, 2026 16:26

svcnvidia-nemo-ci marked this pull request as draft March 18, 2026 16:27

root and others added 4 commits March 18, 2026 09:34

Run autoformatter on test file

17597ac

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

mchrzanowski marked this pull request as ready for review March 18, 2026 23:02

Merge branch 'main' into add-muon-coefficient-type

8a025d3

mchrzanowski requested a review from mkhona-nvidia March 18, 2026 23:02

svcnvidia-nemo-ci requested a review from a team March 18, 2026 23:02

svcnvidia-nemo-ci added the complexity: low label Mar 18, 2026

mkhona-nvidia approved these changes Mar 18, 2026

View reviewed changes

Comment thread megatron/core/optimizer/optimizer_config.py

chtruong814 approved these changes Mar 19, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Mar 19, 2026

deepakn94 reviewed Mar 19, 2026

View reviewed changes

cuichenx approved these changes Mar 19, 2026

View reviewed changes

skyw reviewed Mar 19, 2026

View reviewed changes

Comment thread megatron/core/optimizer/muon.py

Merge branch 'main' into add-muon-coefficient-type

678c832

ericharper reviewed Mar 19, 2026

View reviewed changes

Comment thread megatron/core/optimizer/muon.py Outdated

root and others added 2 commits March 19, 2026 15:29

Narrow broad Exception catch to PackageNotFoundError

02bad44

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

mchrzanowski enabled auto-merge March 19, 2026 23:52

mchrzanowski requested review from deepakn94 and ericharper March 21, 2026 09:27

ericharper approved these changes Mar 25, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Mar 25, 2026

svcnvidia-nemo-ci added this to the Core 0.16 milestone Mar 25, 2026

root and others added 2 commits March 25, 2026 10:43

Fix pylint E0606 for conditional imports of emerging_optimizers

99e8ad4

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Fix black formatting for flash_attn imports in attention.py

8d9f7c5

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

copy-pr-bot Bot temporarily deployed to test March 25, 2026 17:52 Inactive

root and others added 2 commits March 25, 2026 13:03

Run black autoformatting on test_muon_optimizer.py

a1dd63d

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

copy-pr-bot Bot temporarily deployed to test March 25, 2026 20:45 Inactive

skyw requested review from skyw and removed request for skyw March 25, 2026 20:47

mchrzanowski and others added 4 commits March 25, 2026 17:19

Merge branch 'main' into add-muon-coefficient-type

b2936b2

copy-pr-bot Bot temporarily deployed to test March 25, 2026 21:36 Inactive

Always pass coefficient_type to newton_schulz_tp regardless of EO ver…

10bf5eb

…sion The coefficient_type parameter is supported in emerging_optimizers >= 0.1, so there's no need to gate it behind HAVE_EO_V02 (>= 0.2). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

copy-pr-bot Bot temporarily deployed to test March 26, 2026 03:25 Inactive

mchrzanowski added this pull request to the merge queue Mar 26, 2026

Merged via the queue into NVIDIA:main with commit 0842ca2 Mar 26, 2026
63 checks passed

mchrzanowski deleted the add-muon-coefficient-type branch March 26, 2026 04:57

skyw mentioned this pull request Mar 31, 2026

Merge emerging-optimizers change from dev to main #4060

Closed

5 tasks

Navigation Menu

Uh oh!

Add --muon-coefficient-type argument for Muon optimizer - #3927

Add --muon-coefficient-type argument for Muon optimizer#3927
mchrzanowski merged 25 commits into
NVIDIA:mainfrom
mchrzanowski:add-muon-coefficient-type

mchrzanowski commented Mar 18, 2026

Uh oh!

copy-pr-bot Bot commented Mar 18, 2026

Uh oh!

github-actions Bot commented Mar 18, 2026

Uh oh!

Uh oh!

deepakn94 Mar 19, 2026

Uh oh!

FDecaYed Mar 19, 2026

Uh oh!

Uh oh!

Uh oh!

mchrzanowski commented Mar 26, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

10 participants

		@@ -288,6 +311,7 @@ def lion_init_state_fn(opt, config=None):
		"use_nesterov": config.muon_use_nesterov,

Uh oh!

Conversation

mchrzanowski commented Mar 18, 2026

Uh oh!

copy-pr-bot Bot commented Mar 18, 2026

Uh oh!

github-actions Bot commented Mar 18, 2026

Uh oh!

Uh oh!

deepakn94 Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

FDecaYed Mar 19, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

mchrzanowski commented Mar 26, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

10 participants