[feat] FSDP support for HybridStack EP-overlap (3/4 of #4798) by Connor-XY · Pull Request #4943 · NVIDIA/Megatron-LM

Connor-XY · 2026-05-22T17:06:35Z

Review this follow-up's changes against #4942: scoped diff. The GitHub base is pull-request/4942, verified at the current parent head; this scoped diff excludes the parent changes.

Before merging #4942, retarget this PR to main so GitHub preserves the dependent review discussion.

What does this PR do?

Add Megatron FSDP v1 support for bracketed HybridStack groups in the 1F1B EP-overlap schedule. Each bracket group becomes one FSDP unit; the outer decoder stack is excluded.

Apply the same instance filter to parameter buckets and lifecycle hooks, so separate groups never share the outer stack's bucket identity. Preserve the existing constructor's positional arguments. Regression coverage includes grouped, mixed, and ungrouped unit selection, plus three-step loss/parameter parity against non-overlapped FSDP for grouped attention and Mamba patterns.

Dependencies and review scope

Depends on [feat] HybridStack grouped syntax and EP-overlap (2/4 of #4798) #4942; [refactor] Common combined-1F1B schedule-plan base (1/4 of #4798) #4941 has merged.
[fix] Training wire-up for hybrid EP-overlap (4/4 of #4798) #4944 wires these FSDP units into the training launcher.
FSDP slice: adapter, FSDP core/bucketing, and two test files. The MTP slot-boundary fix now belongs to [feat] HybridStack grouped syntax and EP-overlap (2/4 of #4798) #4942.
Original feature split from [feat] Hybrid model ep overlapping main #4798 by @Wohox and @Connor-XY.

Validation

Refreshed on current main through #4942. Black, isort, repository-pinned pylint, Ruff, and diff checks pass for this slice. A limited CPU harness using PyTorch 2.14.0 passes all six unit-selection cases and fails the filtered case when the bucketing fix is removed. It executes the source grouping function with ordinary CPU parameters; this does not validate CUDA/FSDP collectives.

An internal GB200 run tested the integrated stack at #6960 c331df5ef3f2d6e057ad11cee28cf6341d9cdd8b plus #7192 d2385590f5ed4e0189e19e434d262ed5c7ac02e7 on four GB200 GPUs (PyTorch 2.12.0a0+0291f960b6.nv26.04.48445190, CUDA 13.2, Transformer Engine 2.14.0+f031cf87). All 18 FSDP tests passed on every rank: 12 three-step loss/final-parameter parity cases across attention/Mamba/MLA, alltoall/flex dispatchers, shared experts off/on, plus six unit/bucket tests. This validates the integrated source combination; standalone GitHub CI has been triggered at this exact head.

copy-pr-bot · 2026-05-22T17:06: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.

Connor-XY · 2026-07-27T23:38:15Z

Rebased onto the updated #4942 branch to pick up the two review fixes there (ffcae6d, 08e141e): the output_processor branch in HybridModel._postprocess, and gating the final_norm → final_layernorm sharded-key rename on bracketed-group patterns so non-grouped hybrid checkpoints keep loading. No changes to this PR's own commits.

Connor-XY · 2026-07-28T00:32:29Z

Rebased onto current main via the updated #4942 branch. No conflicts in this PR's own commits — the conflicts were all in #4942's files against the MLA-in-HybridModel port (#4452); details there.

Connor-XY · 2026-09-30T04:18:17Z

/ok to test 166d10e

wujingyue · 2026-09-30T06:04:18Z

            ddp_config, ["optim_grads_params"]
        ):
-            supported_fsdp_unit_modules = [TransformerLayer, MoETransformerLayer, MambaLayer]
+            from megatron.core.models.hybrid.hybrid_block import HybridStack


Done in 1715f7e: HybridStack is now imported at the top of mcore_fsdp_adapter.py. It doesn't create an import cycle. I checked fresh-process imports of megatron.core, the adapter, hybrid_block/hybrid_model and megatron.training, plus a full unit-test collection, on our container.

wujingyue · 2026-09-30T06:05:33Z

                fsdp_unit_modules=self.fsdp_unit_modules,
+                fsdp_unit_filter=_fsdp_unit_filter,


These two now overlap quite a bit. Keep only fsdp_unit_filter because it's more general?

Agreed. In 1715f7e fsdp_unit_filter is now the only thing MegatronFSDP and BucketingPolicy use to pick FSDP units. It's a standalone predicate now, not a post-filter on the class match. I kept fsdp_unit_modules only as shorthand, lowered to an equivalent FSDPUnitTypeFilter, because it's the public megatron-fsdp / fully_shard argument (it also takes class-path strings) and many callers use it. Passing both now raises. The MCore adapter passes just fsdp_unit_filter: its unit types minus the outer HybridStack of a grouped pattern.

many callers use it

How bad is it really? We could provide a converter from fsdp_module_types to a filter to ease the migration.

Rewrite comments that only made sense relative to the diff or to a past incident: drop a line-number reference in favor of the function name, state the recompute-hook invariant instead of the symptom it caused, and describe the checkpoint keys each pattern publishes instead of calling them historical or pre-existing. Signed-off-by: Yan Xu <yxu1@nvidia.com>

The schedule plan checked static config on every build_schedule_plan() call, i.e. every microbatch, walked all modules each time, and mixed assert with ValueError. Move the checks to HybridModel.__init__ and raise ValueError: - CUDA graphs (the message no longer claims the check is specific to grouped patterns; it applies to any HybridModel with EP overlap), - mHC connections and MoE shortcut connections, - mixed context-parallel layouts, checked on every HybridStack including nested bracketed groups. Drop the hash-routed MoE and wide-residual checks: TransformerConfig already rejects both with overlap_moe_expert_parallel_comm. Signed-off-by: Yan Xu <yxu1@nvidia.com>

Read known attributes directly and pass arguments unconditionally: - HybridStack sets final_norm to None when it has no final norm (like TransformerBlock.final_layernorm), so the final_layernorm alias and the overlap callables read it directly. The flextron hook checks for a None final norm instead of a missing attribute. - The schedule plan dispatches on type: HybridStack for the layer-type symbol, and TransformerLayer / MultiTokenPredictionLayer for the layer-level quantization context. Its imports move to module level. - The MoE preprocess slot always sets shared_expert_output; the combine slot reads it and mlp_norm_manager directly. The Mamba pre-layer call no longer probes the chunk state for an inference context the training schedule never sets. - HybridStack.forward drops rotary_pos_cos / rotary_pos_sin / rotary_pos_cos_sin, which nothing passes in, and HybridModel._postprocess drops is_spec_decode, which no caller passes. - A bracketed group stack never applies full recomputation itself; gate on the static is_layer_group_stack instead of threading the private _checkpointed_forward_in_parent flag through forward. checkpointed_forward identifies group layers with isinstance(layer, HybridStack) and passes input_ids, packed_seq_params_by_layout and cp_layout_plan unconditionally. - MambaLayer.backward_dw and MambaMixer.backward_dw call their delayed wgrad unconditionally. Building the overlap callables with delay_wgrad_compute now raises for a Mamba layer whose mixer cannot defer its wgrad (e.g. GatedDeltaProductMixer) instead of silently skipping those weight gradients. Signed-off-by: Yan Xu <yxu1@nvidia.com>

@Wohox

Adjust the mcore-FSDP adapter and the megatron-FSDP core so HybridStack (including nested grouped HybridStack instances) is a valid FSDP unit and participates in the EP-overlap schedule plan. Add the ``test_fsdp_hybrid_overlap`` integration test exercising the FSDP + grouped HybridModel forward/backward path. Part 3/4 of splitting NVIDIA#4798 (original changes by @Wohox). Depends on the HybridStack changes in part 2/4 (#TBD). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Signed-off-by: Yan Xu <yxu1@nvidia.com>

Address review feedback on the FSDP HybridStack support: - Import HybridStack at module level in the MCore FSDP adapter instead of inside __init__. - Make fsdp_unit_filter the one mechanism that selects FSDP units. MegatronFSDP and BucketingPolicy now consult only the filter; the public fsdp_unit_modules argument is kept as shorthand for a class-based filter (FSDPUnitTypeFilter), and passing both raises a ValueError. - The MCore adapter passes only fsdp_unit_filter, built from its unit types, and leaves out the outer HybridStack of a grouped hybrid pattern. Signed-off-by: Yan Xu <yxu1@nvidia.com>

Connor-XY · 2026-09-30T23:00:51Z

/ok to test 1715f7e

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch 8 times, most recently from 6fc6e32 to 2e09a54 Compare June 3, 2026 19:21

Victarry mentioned this pull request Jun 10, 2026

[ROADMAP][2026 Q2] Megatron Core MoE Roadmap #4815

Closed

44 of 72 tasks

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch from 2e09a54 to 1f3b873 Compare June 29, 2026 16:29

sbhavani mentioned this pull request Jul 6, 2026

[ROADMAP][2026 Q3] Megatron Core Roadmap #5676

Open

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch 5 times, most recently from 36b00db to f06507c Compare July 20, 2026 17:14

cuichenx mentioned this pull request Jul 20, 2026

[bug] GPT-OSS HybridModel lacks EP-overlap schedule support after migration NVIDIA-NeMo/Megatron-Bridge#4975

Open

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch 2 times, most recently from e378675 to 0b44daf Compare July 27, 2026 23:33

Connor-XY mentioned this pull request Jul 27, 2026

[feat] HybridStack grouped syntax and EP-overlap (2/4 of #4798) #4942

Open

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch from 0b44daf to bc5a0f7 Compare July 28, 2026 00:22

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch from bc5a0f7 to 9a766fc Compare September 8, 2026 20:30

NVSkySC assigned Connor-XY Sep 29, 2026

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch 2 times, most recently from 9a766fc to e7ad92c Compare September 29, 2026 22:23

This was referenced Sep 29, 2026

[fix] Training wire-up for hybrid EP-overlap (4/4 of #4798) #4944

Open

Test MLA grouped HybridStack FSDP EP-overlap #6960

Open

Connor-XY marked this pull request as ready for review September 29, 2026 22:32

Connor-XY requested review from a team as code owners September 29, 2026 22:32

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch from b95f1a3 to e74158c Compare September 29, 2026 22:38

github-actions Bot added the complexity: high label Sep 29, 2026

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch from e74158c to 166d10e Compare September 29, 2026 22:43

Connor-XY changed the base branch from main to pull-request/4942 September 30, 2026 04:19

copy-pr-bot Bot deployed to test September 30, 2026 04:19 Active

copy-pr-bot Bot deployed to public September 30, 2026 04:22 Active

wujingyue reviewed Sep 30, 2026

View reviewed changes

Connor-XY and others added 8 commits September 30, 2026 15:16

Keep grouped HybridStack FSDP buckets aligned with lifecycle units

b0789b9

Signed-off-by: Yan Xu <yxu1@nvidia.com>

Resolve optimizer DTensors to original FSDP bucket parameters in test

5da94bb

Signed-off-by: Yan Xu <yxu1@nvidia.com>

Enable deterministic Mamba scan in hybrid FSDP parity test

349dc21

Signed-off-by: Yan Xu <yxu1@nvidia.com>

Connor-XY force-pushed the pr4798-3-fsdp-hybrid branch from 166d10e to 1715f7e Compare September 30, 2026 23:00

copy-pr-bot Bot deployed to test September 30, 2026 23:01 Active

copy-pr-bot Bot deployed to public September 30, 2026 23:04 Active

Navigation Menu

[feat] FSDP support for HybridStack EP-overlap (3/4 of #4798) - #4943

[feat] FSDP support for HybridStack EP-overlap (3/4 of #4798)#4943
Connor-XY wants to merge 8 commits into
NVIDIA:pull-request/4942from
Connor-XY:pr4798-3-fsdp-hybrid

Connor-XY commented May 22, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented May 22, 2026

Uh oh!

Connor-XY commented Jul 27, 2026

Uh oh!

Connor-XY commented Jul 28, 2026

Uh oh!

Connor-XY commented Sep 30, 2026

Uh oh!

wujingyue Sep 30, 2026

Uh oh!

Connor-XY Sep 30, 2026

Uh oh!

wujingyue Sep 30, 2026

Uh oh!

Connor-XY Sep 30, 2026

Uh oh!

wujingyue Oct 1, 2026

Uh oh!

Connor-XY commented Sep 30, 2026

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		fsdp_unit_modules=self.fsdp_unit_modules,
		fsdp_unit_filter=_fsdp_unit_filter,

Conversation

Connor-XY commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Dependencies and review scope

Validation

Uh oh!

copy-pr-bot Bot commented May 22, 2026

Uh oh!

Connor-XY commented Jul 27, 2026

Uh oh!

Connor-XY commented Jul 28, 2026

Uh oh!

Connor-XY commented Sep 30, 2026

Uh oh!

wujingyue Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

Connor-XY Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

wujingyue Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

Connor-XY Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

wujingyue Oct 1, 2026

Choose a reason for hiding this comment

Uh oh!

Connor-XY commented Sep 30, 2026

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Connor-XY commented May 22, 2026 •

edited

Loading