Merge emerging-optimizers change from dev to main by skyw · Pull Request #4060 · NVIDIA/Megatron-LM

skyw · 2026-03-30T16:45:05Z

What does this PR do ?

Emerging optimizer registry: New emerging_optimizers.py with registry pattern for Muon, AdaptiveMuon, SOAP, Lion, and future
optimizers. Muon code consolidated from standalone muon.py (now a backward-compat shim) into the registry framework.
Layer-wise optimizer improvements: Simplified LayerWiseDistributedOptimizer init (removed async_allgather param, uses config
directly), improved error messages for pre-wrapped optimizers.
Config updates: New fields (use_layer_wise_distributed_optimizer, SOAP/AdaptiveMuon params), muon_use_nesterov renamed to
muon_nesterov, AdamOptimizerConfig/SGDOptimizerConfig simplified to aliases.
CLI args: Expanded --optimizer choices (soap, adaptive_muon), --muon-nesterov rename, dist_muon deprecation with auto-redirect to
muon + layer-wise.
ChainedOptimizer fix: _split_state_dict now uses model chunk identity mapping instead of sequential offset, fixing state dict
splitting for mixed optimizer types.

Add --muon-coefficient-type argument for Muon optimizer #3927 — Add --muon-coefficient-type argument for Muon optimizer
Add more emerging optimizers #3907 — Add more emerging optimizers
Fix emerging optimizer init_group for ckpt loading #3897 — Fix emerging optimizer init_group for ckpt loading
Add Lion optimizer support #3813 — Add Lion optimizer support
Continue emerging optimizer refactoring #3737 — Continue emerging optimizer refactoring
Scaling for MuP over Muon optimizer. #3715 — Scaling for MuP over Muon optimizer
[dev] refactor to support emerging optimizers beyond muon #3618 — Refactor to support emerging optimizers beyond muon
add --overlap-param-gather support for layer-wise optimizer. lots of unit tests. #3524 — Add --overlap-param-gather support for layer-wise optimizer

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Bring over optimizer improvements from the dev branch including: - Emerging optimizers registry and refactoring (Lion, AdaptiveMuon, etc.) - Optimizer state and master weight CPU offloading - Layer-wise optimizer improvements with --overlap-param-gather support - Muon optimizer cleanup and consolidation into emerging optimizers framework - Updated unit tests (new: emerging_optimizers, lion, state_offloading; removed: muon_optimizer) - Updated functional test golden values Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…rad_buffer.py - arguments.py: Add CLI args for new optimizers (soap, adaptive_muon), --offload-optimizer-states flag, --muon-nesterov rename, dist_muon deprecation handling, use_layer_wise_distributed_optimizer logic, and emerging optimizer validation checks - param_and_grad_buffer.py: Add grad mode tracking for cached param views to avoid stale views when grad mode changes Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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

Reverts: OptimizerStateOffloader, distrib_optimizer quantized tensor changes, offload_optimizer_states config/CLI arg, and offloading test. These will be merged separately. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

copy-pr-bot · 2026-03-30T16:45:09Z

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.

Lion is part of the emerging_optimizers package, so a single HAVE_EMERGING_OPTIMIZERS check is sufficient. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Signed-off-by: Hao Wu <skyw@nvidia.com>

skyw · 2026-03-30T17:01:02Z

/ok to test 04e0d35

skyw · 2026-03-30T17:02:05Z

@FDecaYed plz take a look. Pay attention to the muon-scalar-optimizer logic which could be breaking.

And some uv settings may be out of sync. CI returned strange missing package errors.

Signed-off-by: Hao Wu <skyw@nvidia.com>

skyw · 2026-03-30T17:20:07Z

/ok to test 0765b35

Signed-off-by: Hao Wu <skyw@nvidia.com>

skyw · 2026-03-31T01:41:02Z

/ok to test 2a21bba

Signed-off-by: Hao Wu <skyw@nvidia.com>

skyw · 2026-03-31T03:41:38Z

/ok to test ce7174e

FDecaYed

some minor issues. otherwise looks good to me

FDecaYed · 2026-03-31T11:38:09Z

        if state_dict is not None:
            if len(self.model_chunks) == 1:
-                state_dicts[0] = state_dict
+                # When there is only one global model chunk, all sub-optimizers


changes from #3676
@kunlunl also agree to merge here in one go

FDecaYed · 2026-03-31T11:51:19Z

-    # Muon optimizer check
-    if 'muon' in args.optimizer:
+    # emerging optimizer check
+    if not hasattr(args, 'use_layer_wise_distributed_optimizer'):


maybe just add this as a usual parameter below, and set to default false there? it might help allow downstream properly set it(bridge for example). because sometime they just create args without running code logics here.
not a hard request though

…o-main # Conflicts: # pyproject.toml # uv.lock

skyw · 2026-03-31T16:35:46Z

/ok to test f4d9301

skyw · 2026-03-31T19:12:12Z

Closing.
In favor of small PRs instead of combined ones.

skyw and others added 4 commits March 30, 2026 09:22

Add muon_coefficient_type config field back to OptimizerConfig

b1be5d1

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

skyw and others added 2 commits March 30, 2026 09:52

Remove HAVE_LION guard, use HAVE_EMERGING_OPTIMIZERS instead

605e615

Lion is part of the emerging_optimizers package, so a single HAVE_EMERGING_OPTIMIZERS check is sufficient. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

update lion logic

04e0d35

Signed-off-by: Hao Wu <skyw@nvidia.com>

skyw marked this pull request as ready for review March 30, 2026 16:54

skyw requested review from a team as code owners March 30, 2026 16:54

svcnvidia-nemo-ci requested a review from a team March 30, 2026 16:54

svcnvidia-nemo-ci added Final Review PR is in the "final review" stage complexity: high labels Mar 30, 2026

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

skyw requested a review from FDecaYed March 30, 2026 17:02

skyw marked this pull request as draft March 30, 2026 17:06

skyw and others added 2 commits March 30, 2026 10:19

fix lion import

dedf3fb

Signed-off-by: Hao Wu <skyw@nvidia.com>

Merge branch 'main' into merge-optimizer-dev-to-main

0765b35

skyw marked this pull request as ready for review March 30, 2026 17:19

svcnvidia-nemo-ci added complexity: high and removed complexity: high labels Mar 30, 2026

skyw marked this pull request as draft March 30, 2026 17:24

skyw removed the Final Review PR is in the "final review" stage label Mar 30, 2026

Merge branch 'main' into merge-optimizer-dev-to-main

51785e0

skyw marked this pull request as ready for review March 31, 2026 01:23

svcnvidia-nemo-ci added Final Review PR is in the "final review" stage and removed complexity: high labels Mar 31, 2026

svcnvidia-nemo-ci added the complexity: high label Mar 31, 2026

skyw added 2 commits March 30, 2026 18:36

add safe global for loading torch ckpt

3c9dd05

Signed-off-by: Hao Wu <skyw@nvidia.com>

remove overlap assert for muon

2a21bba

Signed-off-by: Hao Wu <skyw@nvidia.com>

skyw marked this pull request as draft March 31, 2026 01:41

update uv lock

ce7174e

Signed-off-by: Hao Wu <skyw@nvidia.com>

skyw marked this pull request as ready for review March 31, 2026 03:41

skyw marked this pull request as draft March 31, 2026 03:42

copy-pr-bot Bot temporarily deployed to test March 31, 2026 03:42 Inactive

svcnvidia-nemo-ci added complexity: high and removed complexity: high labels Mar 31, 2026

skyw removed the Final Review PR is in the "final review" stage label Mar 31, 2026

FDecaYed approved these changes Mar 31, 2026

View reviewed changes

skyw marked this pull request as ready for review March 31, 2026 15:27

svcnvidia-nemo-ci added Final Review PR is in the "final review" stage complexity: high and removed complexity: high labels Mar 31, 2026

Merge remote-tracking branch 'origin/main' into merge-optimizer-dev-t…

f4d9301

…o-main # Conflicts: # pyproject.toml # uv.lock

skyw added core_r0.17.0 Auto-cherrypick to release branch. Apply before merge; cherrypick happens after merge. and removed core_r0.17.0 Auto-cherrypick to release branch. Apply before merge; cherrypick happens after merge. labels Mar 31, 2026

copy-pr-bot Bot temporarily deployed to test March 31, 2026 16:36 Inactive

skyw requested a review from deepakn94 March 31, 2026 16:41

skyw closed this Mar 31, 2026

skyw deleted the merge-optimizer-dev-to-main branch May 6, 2026 20:48

Navigation Menu

Uh oh!

Merge emerging-optimizers change from dev to main - #4060

Merge emerging-optimizers change from dev to main#4060
skyw wants to merge 13 commits into
NVIDIA:mainfrom
skyw:merge-optimizer-dev-to-main

skyw commented Mar 30, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Mar 30, 2026

Uh oh!

skyw commented Mar 30, 2026

Uh oh!

skyw commented Mar 30, 2026 •

edited

Loading

Uh oh!

skyw commented Mar 30, 2026

Uh oh!

skyw commented Mar 31, 2026

Uh oh!

skyw commented Mar 31, 2026

Uh oh!

FDecaYed left a comment

Uh oh!

Uh oh!

FDecaYed Mar 31, 2026

Uh oh!

Uh oh!

FDecaYed Mar 31, 2026

Uh oh!

skyw commented Mar 31, 2026

Uh oh!

skyw commented Mar 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

skyw commented Mar 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Mar 30, 2026

Uh oh!

skyw commented Mar 30, 2026

Uh oh!

skyw commented Mar 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

skyw commented Mar 30, 2026

Uh oh!

skyw commented Mar 31, 2026

Uh oh!

skyw commented Mar 31, 2026

Uh oh!

FDecaYed left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

FDecaYed Mar 31, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

FDecaYed Mar 31, 2026

Choose a reason for hiding this comment

Uh oh!

skyw commented Mar 31, 2026

Uh oh!

skyw commented Mar 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

skyw commented Mar 30, 2026 •

edited

Loading

skyw commented Mar 30, 2026 •

edited

Loading