fix(muon): avoid EP/FSDP AGRS OOM - #2069

Merged

jayhenry merged 15 commits into

InternLM:mainfrom

jayhenry:fix_muon_oom

Sep 21, 2026

Commits on Sep 21, 2026

fix(muon): expose padded all-to-all remainder strategy
jayhenry
committed
7520645
fix(muon): use FSDP-local shard dimension for EP
jayhenry
committed
bbf9f57
Fix Muon EP FSDP batch dimensions
jayhenry
committed
3f38933
Stabilize Muon FSDP parity test
jayhenry
committed
2a126aa
Linearize Muon task creation
jayhenry
committed
29c6720
[Refactor] Split Muon._create_muon_tasks into a mesh-plan/task hierarchy
jayhenry
committed
8243228
[Refactor] Drop redundant None-fallback and clarify Muon mesh-plan steps
jayhenry
committed
08c587a
[Refactor] Return Muon mesh comm plans per strategy
jayhenry
committed
1aceaa3
[Fix] Keep MoE experts on the Muon sub-group all-gather path
jayhenry
committed
5b2870d
[Fix] Default Muon remainder batches to pad_all2all
jayhenry
committed
6703074
[Feature] Exchange evenly sharded Muon remainders with a ragged all-to-all
jayhenry
committed
004bc25
[Refactor] Let the ragged Muon all-to-all carry uneven shards
jayhenry
committed
3bbb742
[Feature] Expose ragged_all_to_all as an explicit Muon remainder strategy
jayhenry
committed
6c6147c
[Refactor] Remove internal Muon OOM investigation notes and repro script

jayhenry
and
claude
committed
e47c33a
[Fix] Wrap Muon task-creation docstring for docformatter
jayhenry
committed
d9e48fb