Add muon and layerwise distributed optimizer - #2241

Merged

Phlip79 merged 29 commits into

NVIDIA:mainfrom

FDecaYed:deyuf/muon

Jan 14, 2026

Commits on Nov 13, 2025

ADLR/megatron-lm!4106 - [DEV] Add muon and layer-wise distributed optimizer

committed
3446b5a
ADLR/megatron-lm!4232 - [DEV] improve muon and layer-wise dist opt unit tests
FDecaYed
committed
d9df6ab
[DEV] support split qkv in muon (#1915)
FDecaYed
committed
000c4ac
[Dev] Support LayerWiseDistributedOptimizer with torch_dist checkpoint format (#1928)

BoxiangW
and
FDecaYed
committed
338062a
[DEV] torch_dist fixes, speed improvements and memory reduction for layerwise distopt (#1965)
FDecaYed
committed
6a2ecf4
[Dev]Revert torch ckpt format change for LayerwiseDistOpt (#2228)

BoxiangW
authored and
FDecaYed
committed
24901a1
[Dev] Add more tests for LayerwiseDistOpt with dist_ckpt (#2132)

BoxiangW
authored and
FDecaYed
committed
ec472e8

Commits on Nov 18, 2025

update emerging-optimizer to v0.1.0

skyw
authored and
FDecaYed
committed
28dfadc
update uv lock

skyw
authored and
FDecaYed
committed
baf883c
add setuptools back for dev

skyw
authored and
FDecaYed
committed
6524168
fix layerwise torch_dist fails due to empty rank
FDecaYed
committed
f4b3c2e
fix empty group on pp chunk
FDecaYed
committed
bfb736f

Commits on Nov 19, 2025

address comments
FDecaYed
committed
c993fdd

Commits on Jan 9, 2026

Merge branch 'main' into deyuf/muon
FDecaYed
committed
98283e6
post merge update due to optimizer refactor
FDecaYed
committed
e96207e
attempt to fix build dep
FDecaYed
committed
cbe01aa
fix conflict due to new paramgroupoverrides change
FDecaYed
committed
7d7ac52
lint
FDecaYed
committed
0637b15
fix black/isort lint conflicts
FDecaYed
committed
f5fda5e
fix merge typo that accidentally removed a line
FDecaYed
committed
ae7dbb4

Commits on Jan 10, 2026

Fix CI issue. Disable PP=4 fore now
BoxiangW
committed
fd297c3
Lint
BoxiangW
committed
9b0f109
Merge branch 'main' into deyuf/muon
ko3n1g
authored
cd3fe36

Commits on Jan 13, 2026

Add assertion for overlapping
BoxiangW
committed
053b21a
Merge branch 'main' into deyuf/muon
BoxiangW
committed
a35730b
Revert uv.lock changes
BoxiangW
committed
e1da8dd

Commits on Jan 14, 2026

Update optimizer_config.py
Phlip79
authored
3dd2f5c
Merge branch 'main' into deyuf/muon
Phlip79
authored
98189ec
Merge branch 'main' into deyuf/muon
Phlip79
authored
4dafbfa