Refactor emerging optimizer integration by skyw · Pull Request #4113 · NVIDIA/Megatron-LM

skyw · 2026-04-02T20:03:00Z

What does this PR do ?

Cherry-picks three related PRs from dev that refactor the optimizer infrastructure to generalize beyond Muon and support additional emerging optimizers (SOAP, Lion, etc.):

[dev] refactor to support emerging optimizers beyond muon #3618 — Initial refactor: extracts emerging_optimizers.py module, generalizes get_megatron_optimizer to handle all optimizer types uniformly, deprecates get_megatron_muon_optimizer
Continue emerging optimizer refactoring #3737 — Continues refactoring: adds SOAP optimizer support, renames test_muon_optimizer.py → test_emerging_optimizers.py, expands test coverage including coefficient type validation
Fix emerging optimizer init_group for ckpt loading #3897 — Fixes init_group in emerging optimizer to support checkpoint loading/resuming

combined those instead cherry picking 1-by-1 because there are immediate changes among them that is not needed anymore.

⚠️ 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.

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

copy-pr-bot · 2026-04-02T20:03:04Z

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-04-02T20:03:10Z

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.

skyw · 2026-04-02T20:39:11Z

/ok to test 5275694

deepakn94 · 2026-04-02T20:59:33Z

+    Handles both standard optimizers (Adam, SGD) and emerging optimizers (e.g. Muon).
    We use separate optimizers for expert parameters and non-expert parameters.
+    For emerging optimizers with ``config.use_layer_wise_distributed_optimizer=True``,
+    the optimizer is automatically wrapped with :class:`LayerWiseDistributedOptimizer`.


I think we should eventually re-name DistributedOptimizer to something else (maybe ParameterWiseDistributedOptimizer)?

(can be done in a separate PR)

I've been using element wise. can debate when we actually do it.

…LayerWise wrapping The test was manually extracting base optimizers from Float16OptimizerWithFloat16Params wrappers and re-wrapping them in LayerWiseDistributedOptimizer, but the fp32 master param updates never propagated back to the bf16 model params. Use use_layer_wise_distributed_optimizer=True in the config instead, which lets _get_megatron_emerging_optimizer handle wrapping correctly.

skyw · 2026-04-02T21:31:23Z

/ok to test c74dd1a

…er=True

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

skyw · 2026-04-02T22:21:23Z

/ok to test 90c5747

…async_allgather

skyw · 2026-04-02T22:58:25Z

/ok to test 16f2e53

svcnvidia-nemo-ci · 2026-04-02T23:51:16Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-04-03T00:18:52Z

🔄 Merge queue validation started!

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

This reverts commit 5b512b4. Signed-off-by: oliver könig <okoenig@nvidia.com>

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

skyw added 6 commits April 2, 2026 12:45

[dev] refactor to support emerging optimizers beyond muon (NVIDIA#3618)

5990608

Continue emerging optimizer refactoring (NVIDIA#3737)

9fefb32

Fix emerging optimizer init_group for ckpt loading (NVIDIA#3897)

710faf3

Fix missing imports in test_emerging_optimizers.py

1f8ab3d

Fix TensorParallelMuon kwarg: mode -> tp_mode in test

c65de58

remove test shouldn't exit yet

5275694

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

skyw requested review from a team as code owners April 2, 2026 20:03

svcnvidia-nemo-ci marked this pull request as draft April 2, 2026 20:03

skyw marked this pull request as ready for review April 2, 2026 20:09

svcnvidia-nemo-ci requested a review from a team April 2, 2026 20:09

svcnvidia-nemo-ci added the complexity: high label Apr 2, 2026

skyw requested a review from deepakn94 April 2, 2026 20:13

svcnvidia-nemo-ci added this to the Core 0.16 milestone Apr 2, 2026

copy-pr-bot Bot temporarily deployed to test April 2, 2026 20:40 Inactive

Fix duplicate _EMERGING_OPTIMIZERS definition for autodoc2

a4c8749

deepakn94 reviewed Apr 2, 2026

View reviewed changes

deepakn94 approved these changes Apr 2, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to test April 2, 2026 21:32 Inactive

yaoyu-33 approved these changes Apr 2, 2026

View reviewed changes

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Apr 2, 2026

skyw added 2 commits April 2, 2026 15:17

Fix test overlap_param_gather: add use_layer_wise_distributed_optimiz…

762ca99

…er=True

remove out of date assert

90c5747

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

copy-pr-bot Bot temporarily deployed to test April 2, 2026 22:22 Inactive

Pass overlap_param_gather config to LayerWiseDistributedOptimizer as …

16f2e53

…async_allgather

copy-pr-bot Bot temporarily deployed to test April 2, 2026 22:59 Inactive

skyw added this pull request to the merge queue Apr 2, 2026

skyw changed the title ~~Refactor emerging optimizer integrate~~ Refactor emerging optimizer integration Apr 3, 2026

Merged via the queue into NVIDIA:main with commit 5b512b4 Apr 3, 2026
63 checks passed

skyw deleted the emerging-optimizer-refactor branch April 3, 2026 01:46

ko3n1g added a commit to ko3n1g/Megatron-LM that referenced this pull request Apr 6, 2026

Revert "Refactor emerging optimizer integration (NVIDIA#4113)"

bf4a1d6

This reverts commit 5b512b4. Signed-off-by: oliver könig <okoenig@nvidia.com>

ko3n1g mentioned this pull request Apr 6, 2026

revert: Revert emerging optimizer refactor (#4113, #4119) #4156

Closed

3 tasks

ko3n1g added a commit to ko3n1g/Megatron-LM that referenced this pull request Apr 6, 2026

Revert "Refactor emerging optimizer integration (NVIDIA#4113)"

5f9f3ff

This reverts commit 5b512b4. Signed-off-by: oliver könig <okoenig@nvidia.com>

yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026

Refactor emerging optimizer integration (NVIDIA#4113)

901f072

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

ko3n1g mentioned this pull request May 26, 2026

[dev-branch] gpt3_moe_mcore_te_ep8_resume_torch_dist_muon: KeyError: 'state' in optimizer #3885

Closed

sbhavani mentioned this pull request May 26, 2026

[ROADMAP][2026 Q2] Megatron Core Roadmap #4997

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor emerging optimizer integration#4113

Refactor emerging optimizer integration#4113
skyw merged 11 commits into
NVIDIA:mainfrom
skyw:emerging-optimizer-refactor

skyw commented Apr 2, 2026

Uh oh!

copy-pr-bot Bot commented Apr 2, 2026

Uh oh!

github-actions Bot commented Apr 2, 2026

Uh oh!

skyw commented Apr 2, 2026

Uh oh!

deepakn94 Apr 2, 2026

Uh oh!

deepakn94 Apr 2, 2026

Uh oh!

skyw Apr 2, 2026

Uh oh!

skyw commented Apr 2, 2026

Uh oh!

skyw commented Apr 2, 2026

Uh oh!

skyw commented Apr 2, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 2, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

skyw commented Apr 2, 2026

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 Apr 2, 2026

Uh oh!

github-actions Bot commented Apr 2, 2026

Uh oh!

skyw commented Apr 2, 2026

Uh oh!

deepakn94 Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

deepakn94 Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

skyw Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

skyw commented Apr 2, 2026

Uh oh!

skyw commented Apr 2, 2026

Uh oh!

skyw commented Apr 2, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 2, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants