Add initial Muon+M-FSDPv2 implementation by janEbert · Pull Request #6425 · NVIDIA/Megatron-LM

janEbert · 2026-08-11T00:09:15Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Add Muon support for M-FSDPv2 based on an experimental optimized M-FSDPv1 prototype.
No integration into MCore yet.

This version uses all-Flat placements and uses an owner-based design:

Non-matrix parameters are simply computed like usual.
Unsharded, fully locally available matrix parameters are computed fully locally.
For each sharded matrix parameter, one owner (that owns at least one shard of that parameter) is picked, on which
1. the full orthogonalization input is reconstructed,
2. orthogonalization performed, and
3. the updates are scattered back to the original destinations.

Pseudocode:

**** Fwd
for each FSDP unit:
    1. COMM: all-gather weight shards, each rank has full weights
    2. execute fwd, each rank has activations over local data

**** Bwd
for each FSDP unit:
    1. COMM: all-gather weight shards, each rank has full weights
    2. execute bwd, each rank has full gradients over local data
    3. COMM: reduce-scatter full gradients, each rank has gradient
       shards over global data for local weight shards

**** Opt step
for each FSDP unit:
    1. update local optimizer state shards using local gradient shards
       and local optimizer state shards
    2. produce local orthogonalization input shards using local
       optimizer state shards and (optionally, with Nesterov) local
       gradient shards
    3. COMM: on owner, gather orthogonalization input shards, each
       owner has full local orthogonalization inputs

       owner is heuristically picked depending on orthogonalization
       compute cost
    4. on owner, compute orthogonalization outputs (the full updates)
    5. COMM: owner scatters full updates, each rank has local update
       shards
    6. update local weight shards using local update shards

Internal design doc with more details: https://docs.google.com/document/d/11FGBYk-hdvGBGh_s-1tGBv2JYLicRPXIHApILtD6Mzw/edit?tab=t.vd156zcgkxyj

Minor information

The change to the get_containing_parameter_group signature is due to it being too strict; I needed to pass raw torch.Tensors to this function.

Issue tracking

Linked issue: Ref #5533.

github-actions · 2026-08-11T00:09:39Z

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.

copy-pr-bot · 2026-08-11T00:09:40Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

wujingyue

Thanks for the PR!

https://google.github.io/eng-practices/review/developer/small-cls.html

This PR is too large for me to review right now. Could we reduce the first iteration to under 500 LOC (including tests), even if that means cutting features or simplifying the algorithm?

wujingyue · 2026-08-11T03:58:33Z

+
+
+@dataclasses.dataclass(frozen=True)
+class ShardPlan:


Apparently, most of the PR is spent on deciding the parameter-to-rank mapping. Can we separate this into another PR and make this PR focused on the most basic allgather-each-parameter-group solution? We may even keep this as a baseline for numerics testing.