fsdp: avoid double sharding of MoE experts when EP is enabled by CodersAcademy006 · Pull Request #2833 · NVIDIA/Megatron-LM

CodersAcademy006 · 2026-01-06T16:56:35Z

When Expert Parallelism (EP) is enabled, MoE expert parameters are already
partitioned across ranks. However, current FSDP auto-wrapping logic still
wraps these expert modules, causing parameters to be effectively tracked
and sharded twice.

This results in:

Increased peak GPU memory usage
Duplicated optimizer state
Additional communication overhead
Reduced throughput compared to EP-only runs

What this PR does

Skips FSDP auto-wrapping for MoE expert modules that are already managed
by Expert Parallelism
Adds an explicit ownership signal (expert_parallel_enabled) to MoE layers
to avoid heuristic or name-based checks
Preserves existing FSDP behavior for all non-expert parameters

Impact

Restores expected memory behavior (EP + FSDP ≤ EP-only)
Reduces redundant parameter metadata and optimizer state
Improves training throughput
No behavior change for non-MoE or non-EP models

Reproduction

This issue is reproducible on large MoE models and is reported in #2772.
The change removes the redundant sharding path responsible for the memory
regression.

Fixes #2772

copy-pr-bot · 2026-01-06T16:56:39Z

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.

zhujian19891203 · 2026-01-13T05:53:00Z

Hi, thank you for this work.
But I tried this PR, it can not run successfully. Because there are multiple code snippets like the following in megatron_fsdp.py and param_and_grad_buffer.py file, it will not filter expert params. For example, function _replace_param_with_distributed_if_needed in megatron_fsdp.py file, the assert will triggered.

CodersAcademy006 · 2026-01-14T05:59:02Z

@zhujian19891203 Thanks for testing this and for the detailed trace.

You’re right: while the current change correctly skips expert parameters during FSDP parameter grouping to avoid double sharding, the same exclusion needs to be applied in the FSDP replacement path.

In particular, _replace_param_with_distributed_if_needed still assumes all parameters are registered with FSDP, which is not true for EP-managed expert parameters. I will update the replacement logic to explicitly bypass expert parameters when expert_parallel_enabled=True, ensuring consistency across all FSDP code paths.

I’ll push a follow-up fix shortly

- Move nv-grouped-gemm from dev/lts extras to new moe extra - Users can now install megatron-core[dev] without build failures - Add comprehensive error messages with installation instructions - Update README with MoE dependencies documentation This resolves the installation failure when nv-grouped-gemm cannot build due to missing CUTLASS headers. Users who need MoE with grouped GEMM can now explicitly install it with megatron-core[moe], while others can install dev/lts extras without encountering build errors. Fixes NVIDIA#2541

…r RL workflows. Fixed imports, ensured native RL loop support, and validated no errors in touched files. Ready for review and CI on supported platforms.

…lag, gate all logic on --rl-amem-offload-during-rollout, encapsulate env setup, and update docs/examples per review

CodersAcademy006 · 2026-01-16T07:55:31Z

I can split AMem NCCL into a follow-up PR if reviewers prefer

shjwudp · 2026-01-16T09:00:21Z

partitioned across ranks. However, current FSDP auto-wrapping logic still

Hi @CodersAcademy006 — that's awesome! 🎉 Please open a separate PR for AMem NCCL integration. We’d also really appreciate it if you could include a brief design document outlining how this feature work in RL training and its expected performance impact.

@HaochenYuan is our RL expert here can help review your code.

…m-nccl flag, gate all logic on --rl-amem-offload-during-rollout, encapsulate env setup, and update docs/examples per review" This reverts commit ee9962c.

Remove all AMem NCCL-related code and documentation from this PR: - Delete megatron/core/amem_nccl.py - Delete examples/rl/train_with_amem.sh - Revert megatron/rl/rl_utils.py to upstream (remove AMem hooks) - Revert megatron/training/arguments.py to upstream (remove AMem flags) - Revert megatron/training/initialize.py to upstream (remove AMem setup) - Revert pyproject.toml to upstream (remove AMem dependencies) - Revert scripts/check_api_backwards_compatibility.py to upstream - Remove AMem documentation section from README.md This PR should focus only on FSDP/EP double-sharding fixes. AMem NCCL integration will be submitted as a separate PR as requested by reviewers.

CodersAcademy006 · 2026-01-16T10:52:49Z

I have made the separate pr's please review. #2979 #2980

- Add assert_not_fsdp_wrapped_ep_param() helper to enforce the invariant - Call assertion in _replace_param_with_distributed_if_needed() - Document invariant in MoE layer where expert_parallel_enabled is set This prevents future double-sharding regressions by making violations loud instead of silent.

CodersAcademy006 · 2026-01-16T11:00:18Z

Added a small invariant guard to ensure EP-owned parameters are never wrapped or replaced by FSDP in any code path. This should prevent future double-sharding regressions.

cspades · 2026-01-23T20:24:39Z

+        # Exclude expert-parallel params from FSDP bookkeeping
+        self.param_to_name = {}
+        self.raw_param = {}
+        from megatron.core.utils import is_ep_owned_param


Since Megatron-FSDP is also standalone pip-installable, mind guarding this with a try/except MCore import, or simply having these utilities live natively under Megatron-FSDP? (After all, this is an Megatron-FSDP-specific interaction with Megatron-LM.)

cspades

@CodersAcademy006 Can you elaborate a bit more on the "double sharding"? Megatron-FSDP should be able to shard expert parameters per local expert per Expert TP rank. Are you suggesting that FSDP sharding is useless here due to the memory inflation of using FSDP with EP? (I also wonder about the loss curves you're seeing - since I believe if you detach from FSDP, no reductions would be happening on the DP-Shard / intra-DP domain right?)

#2772 seems to be resolved by a TE issue as well, looking like some bug that un-distributed the heavyweight optimizer state (which coincidentally, becomes a major issue as model size increases), so wanted clarity on whether that problem has already been solved in a less radical way.

Could you provide more technical justification on the problem between the interaction of MoE layers in MLM with Megatron-FSDP? 🙏🏻 Or maybe this PR is just out-dated and the root cause of the memory issue was resolved.

CodersAcademy006 · 2026-01-24T05:40:39Z

hey @cspades Thanks for the review. This PR is not addressing the TE FusedAdam regression from #2772 (which was correctly fixed upstream), but instead enforces a correctness invariant between Expert Parallelism and Megatron-FSDP. EP-owned expert parameters are already partitioned and communicated by EP and must never be wrapped, replaced, or sharded again by FSDP. Without an explicit exclusion, FSDP still attempts to manage these parameters in multiple code paths, which is architecturally incorrect and can reintroduce double-sharding or invalid reductions in the future. This change makes EP ownership explicit and consistently respected across all FSDP paths, independent of optimizer implementation or current memory behavior.

Please correct me otherwise.

shjwudp · 2026-01-27T07:07:10Z

Hi @CodersAcademy006,

When enabling EP, expert parameters may still have replicas. If the DP world size > the EP world size, some expert parameters will have data-parallel neighbors—this setup is known as Expert Data Parallelism (EDP). Applying FSDP on top of EDP can still yield meaningful memory savings.

Conceptually, FSDP is transparent to all parameters. Each parameter maintains its original copy; we generally don't modify it directly—instead, we create shadow shards and unshard them as needed. The function that restores module parameters to their original form is located here: megatron_fsdp.py#L1159.

Regarding your point, I believe our aim should be to maintain broad compatibility between FSDP and DTensor—fixing issues in the ecosystem where needed—rather than avoiding either FSDP or DTensor altogether. This PR unintentionally cancels out FSDP’s effect on expert parameters, which isn’t ideal.

cspades · 2026-03-04T15:50:49Z

@CodersAcademy006 Just to conclude, maybe the ask is for better documentation on how to use EP-EDP-ETP distribution for MoE layers, or support for specific layers, with appropriate advice for how to actually reduce memory utilization? We definitely don't want users to try this code and come to the initial conclusion that the easiest way to get it to work is just to get rid of FSDP for MoE.

Otherwise, good to close this PR? I'll also try to improve the user experience with EP outside of Megatron-LM or custom training FW's, so hopefully all of this is better documented.

CodersAcademy006 · 2026-03-04T18:14:06Z

Thanks for the thorough review @cspades @shjwudp.

You're right — I was reasoning about the pure EP case and missed the
EDP scenario where DP world size > EP world size. In that setup, FSDP
over expert params is still meaningful and my exclusion logic would
incorrectly cancel those memory savings.

Given that #2772 was resolved upstream and the broader fix requires
DTensor/FSDP compatibility work rather than exclusion logic, I'll close
this PR. Happy to help with documentation on EP-EDP-ETP memory
behaviour if that's useful.

Also thank you for correcting my reasoning and helping me understand better with real exposure.

fsdp: avoid double sharding of MoE experts when EP is enabled

2583898

CodersAcademy006 requested review from a team as code owners January 6, 2026 16:56

Merge branch 'main' into fix/fsdp-ep-double-sharding

8fc0e45

github-actions Bot requested a review from Phlip79 January 6, 2026 16:56

github-actions Bot added the community-request label Jan 6, 2026

Skylion007 approved these changes Jan 6, 2026

View reviewed changes

Merge branch 'main' into fix/fsdp-ep-double-sharding

df99d31

yaox12 added Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. module: megatron-fsdp labels Jan 9, 2026

chtruong814 added the needs-follow-up Issue needs follow-up label Jan 11, 2026

CodersAcademy006 and others added 7 commits January 14, 2026 11:36

Resolve merge conflicts and ensure AMem NCCL integration is robust fo…

03cbcca

…r RL workflows. Fixed imports, ensured native RL loop support, and validated no errors in touched files. Ready for review and CI on supported platforms.

Merge branch 'main' into AMen-to-offloaf-NCCL

ff39441

Refactor AMem NCCL integration: remove redundant --rl-use-amem-nccl f…

ee9962c

…lag, gate all logic on --rl-amem-offload-during-rollout, encapsulate env setup, and update docs/examples per review

Merge branch 'main' into fix/fsdp-ep-double-sharding

d531165

Merge remote-tracking branch 'fork/fix/fsdp-ep-double-sharding'

6ac5a3c

refactor: use is_ep_owned_param helper for expert parallel param checks

a110379

CodersAcademy006 requested review from a team and pablo-garay as code owners January 16, 2026 06:40

CodersAcademy006 and others added 2 commits January 16, 2026 12:16

fix: move is_ep_owned_param to proper location and restore file header

dd4fdea

Merge branch 'main' into fix/fsdp-ep-double-sharding

56a650e

CodersAcademy006 added 2 commits January 16, 2026 15:04

Revert "Refactor AMem NCCL integration: remove redundant --rl-use-ame…

c9655e6

…m-nccl flag, gate all logic on --rl-amem-offload-during-rollout, encapsulate env setup, and update docs/examples per review" This reverts commit ee9962c.

cspades reviewed Jan 23, 2026

View reviewed changes

cspades requested review from a team and xuwchen January 23, 2026 20:33

cspades requested changes Jan 23, 2026

View reviewed changes

chtruong814 removed the needs-follow-up Issue needs follow-up label Jan 30, 2026

CodersAcademy006 closed this Mar 4, 2026

Navigation Menu

Uh oh!

fsdp: avoid double sharding of MoE experts when EP is enabled - #2833

fsdp: avoid double sharding of MoE experts when EP is enabled#2833
CodersAcademy006 wants to merge 15 commits into
NVIDIA:mainfrom
CodersAcademy006:fix/fsdp-ep-double-sharding

CodersAcademy006 commented Jan 6, 2026

Uh oh!

copy-pr-bot Bot commented Jan 6, 2026

Uh oh!

zhujian19891203 commented Jan 13, 2026 •

edited

Loading

Uh oh!

CodersAcademy006 commented Jan 14, 2026

Uh oh!

CodersAcademy006 commented Jan 16, 2026

Uh oh!

shjwudp commented Jan 16, 2026 •

edited

Loading

Uh oh!

CodersAcademy006 commented Jan 16, 2026

Uh oh!

CodersAcademy006 commented Jan 16, 2026

Uh oh!

cspades Jan 23, 2026 •

edited

Loading

Uh oh!

cspades left a comment •

edited

Loading

Uh oh!

CodersAcademy006 commented Jan 24, 2026

Uh oh!

shjwudp commented Jan 27, 2026

Uh oh!

cspades commented Mar 4, 2026 •

edited

Loading

Uh oh!

CodersAcademy006 commented Mar 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

Uh oh!

Conversation

CodersAcademy006 commented Jan 6, 2026

What this PR does

Impact

Reproduction

Uh oh!

copy-pr-bot Bot commented Jan 6, 2026

Uh oh!

zhujian19891203 commented Jan 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

CodersAcademy006 commented Jan 14, 2026

Uh oh!

CodersAcademy006 commented Jan 16, 2026

Uh oh!

shjwudp commented Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

CodersAcademy006 commented Jan 16, 2026

Uh oh!

CodersAcademy006 commented Jan 16, 2026

Uh oh!

cspades Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cspades left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

CodersAcademy006 commented Jan 24, 2026

Uh oh!

shjwudp commented Jan 27, 2026

Uh oh!

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

Uh oh!

Uh oh!

CodersAcademy006 commented Mar 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

zhujian19891203 commented Jan 13, 2026 •

edited

Loading

shjwudp commented Jan 16, 2026 •

edited

Loading

cspades Jan 23, 2026 •

edited

Loading

cspades left a comment •

edited

Loading

cspades commented Mar 4, 2026 •

edited

Loading