[feat] FSDP support for HybridStack EP-overlap (3/4 of #4798) - #4943
Draft
Connor-XY wants to merge 11 commits into
Draft
[feat] FSDP support for HybridStack EP-overlap (3/4 of #4798)#4943Connor-XY wants to merge 11 commits into
Connor-XY wants to merge 11 commits into
Conversation
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
8 times, most recently
from
June 3, 2026 19:21
6fc6e32 to
2e09a54
Compare
72 tasks
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
from
June 29, 2026 16:29
2e09a54 to
1f3b873
Compare
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
5 times, most recently
from
July 20, 2026 17:14
36b00db to
f06507c
Compare
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
2 times, most recently
from
July 27, 2026 23:33
e378675 to
0b44daf
Compare
3 tasks
Contributor
Author
|
Rebased onto the updated #4942 branch to pick up the two review fixes there (ffcae6d, 08e141e): the |
Add bracketed HybridStack group syntax (e.g. ``[*-]``, ``M[M*]-``) with nested HybridStack instances, rejecting invalid recursion. Migrate grouped HybridStack checkpoints to Transformer-compatible logical layer keys and make ``HybridModel.sharded_state_dict()`` drop the empty ``output_layer._extra_state`` to match GPT behavior. Extend EP-overlap scheduling to HybridStack: add the hybrid fine-grained callables and ``HybridStackModelChunkSchedulePlan``, expose ``HybridModel.build_schedule_plan`` and add the ``return_schedule_plan`` path in ``pretrain_hybrid.py``. Add Mamba ``backward_dw`` so the hybrid schedule node can register Mamba pre-layer weight grads alongside attention and GDN pre-layers. Fix the MoE TopKRouter MTP layer-number indexing when the MTP block wraps a HybridStack so the aux-loss tracker is not indexed past its size. Part 2/4 of splitting NVIDIA#4798 (original changes by @Wohox). Depends on the common combined-1F1B refactor in part 1/4 (#TBD). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Carries over upstream commit ce6e229 from NVIDIA#4798: in HybridStack's ``_run_moe_combine`` (A2A overlap path), ``layer._forward_post_mlp`` registers a second ``discard_output_and_register_recompute`` hook on ``mlp_output_with_bias[0]``. The hook fires during combine_bwd's autograd backward and triggers the LN recompute ahead of mlp_bwd / pre_dispatch_bwd. In bracketed-hybrid logical layers (``[*E]``), this corrupts gradients in attention's autograd chain (grad_norm explodes from iter 2). Fix: stop calling ``_forward_post_mlp`` from ``_run_moe_combine``; inline the ``bda + offload_mlp_norm + make_viewless_tensor`` steps directly, mirroring GPT's ``submodule_combine_forward``. The first recompute hook on ``expert_output`` (registered in ``_run_moe_experts``) already fires the LN recompute in mlp_bwd, so the second hook is redundant. Part 2/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
The recent merge of origin/main introduced `name=(name + f".layers.{i}")`
into every layer-type branch of HybridStack's build loop, but didn't change
the local loop header `for layer_type in self.layer_type_list:` to surface
`i`. Result: `NameError: name 'i' is not defined` at HybridStack init for
all hybrid runs (GPT path unaffected).
Trigger: any hybrid_stack_spec model crashes on init, including the 16-node
Bug 2a repro and the 8-node GPT-vs-Hybrid perf comparison runs.
Fix: convert the loop to `for i, layer_type in enumerate(...)`. Keep the
existing `physical_layer_offset` counter (used for FP8/FP4 contexts and
`layer_number`) because bracket groups count >1 physical layer per logical
entry — these are separate from the logical index `i` used for module names.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Co-authored-by: Pingtian Li <pingtianl@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Publish speculative-decoding hidden states through the inference context, including the fixed buffer used by block-scope CUDA graphs. Preserve the canonical inference-mode check and the inputs required to derive RL MTP labels. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
…essor Two compatibility fixes from the strict review on NVIDIA#4942: - HybridModel._postprocess accepted output_processor / output_processor_context but silently discarded them, so a caller that wired an output hook through PostProcessNode would get the default logits/loss path with no error. Implement the same early-return branch GPTModel._postprocess has. - The final-norm sharded key rename (final_norm -> final_layernorm) was unconditional, so it also changed the keys of non-grouped hybrid models whose existing dist checkpoints were saved under final_norm. Gate it on a new transformer_sharded_keys flag that HybridModel derives from the full layer pattern, so only bracketed-group models (whose logical layers map one-to-one onto transformer layers, which is what the GPT cross-load compatibility is for) get the transformer-style key. Signed-off-by: Yan Xu <yxu1@nvidia.com>
test_group_sharded_state_dict_uses_logical_layer_keys builds a HybridStack directly rather than through HybridModel, so it has to set transformer_sharded_keys itself now that the final-norm key rename is gated. Add the mirror-image case asserting the default keeps final_norm. Signed-off-by: Yan Xu <yxu1@nvidia.com>
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>
When the EP overlap schedule runs an MTP layer, `submodule_mtp_pre_dispatch_forward` stashes `torch.chunk(hidden_states, ...)` into ``node.chunk_state.mtp_hidden_states`` in the pre_dispatch slot, and ``submodule_mtp_postprocess_forward`` later does ``torch.cat(mtp_hidden_states, ...)`` in the mtp_post_process slot before feeding the LM head. Because ``chunk_state`` is a plain Python container shared across slots, the chunks carry their original grad_fn across the slot boundary — sidestepping the implicit ``detach()`` that ``ScheduleNode._forward`` applies to slot inputs. When the Bug 1 fix (commit f6ea23b) added ``final_norm(hidden_states)`` ahead of the chunk on HybridModel-with-empty-decoder VPP chunks, the chunks' ``SplitBackward → final_norm`` chain became reachable from two independent ``run_backward`` calls: post_process → mtp_post_process traverses it via the cat, and pre_dispatch's own backward traverses it via the MTP forward chain (`chunks[offset]` is the same Tensor as ``mtp_hidden_states[offset]``). ``final_norm`` is a ``TENorm`` and therefore goes through TE's modular OpFuser, whose backward consumes ``ctx.tensor_objects`` and sets it to ``None``. The second traversal then trips the guard in ``transformer_engine/pytorch/quantized_tensor.py:restore_from_func_ctx`` and raises ``AttributeError: ctx must have .tensor_objects to restore saved tensors`` — observed on every rank of the PP stage that owns MTP on the DeepSeek-V3-Proxy-Hybrid-NoMLA 8-node EP-overlap run. GPT does not hit the same error because: - the Bug 1 final_norm branch is gated on ``isinstance(model, HybridModel)`` so GPT never inserts an OpFuser node into the MTP pre_dispatch slot's autograd chain, - GPT's mixed-VPP layout puts ``final_layernorm`` inside the *last decoder* layer's combine slot (see ``submodule_combine_forward``), and the ``ScheduleNode._forward`` input ``.detach()`` between that combine slot and MTP's pre_dispatch keeps the chunks' grad_fn rooted at a slot leaf rather than at ``final_layernorm`` itself. Fix: route the stored chunks through ``node.detach`` so the cross-slot view is a list of leaves. ``node.detach`` records the originals in ``before_detached`` and the detached copies in ``self.detached``, so ``TransformerLayerNode.backward_impl`` keeps pulling the LM-head-side grad (accumulated on the detached leaves by mtp_post_process / post_process backward) back into pre_dispatch's ``run_backward(outputs + before_detached, ...)`` call — gradient flow stays mathematically equivalent, just no longer shared across slots. ``hidden_states = chunks[offset]`` (the live tensor) remains the MTP input so the in-slot forward chain (eh_proj → attention → ...) still propagates grads correctly to the rest of the slot's graph. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-authored-by: Pingtian Li <pingtianl@nvidia.com>
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
from
July 28, 2026 00:22
0b44daf to
bc5a0f7
Compare
Contributor
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Part 3 of 4 splitting #4798 by @Wohox and @Connor-XY. Original changes by @Wohox and @Connor-XY.
Summary
Wire the megatron-FSDP adapter and core so HybridStack (including nested grouped HybridStack instances) is a valid FSDP unit and participates in the EP-overlap schedule plan.
megatron/core/distributed/fsdp/mcore_fsdp_adapter.py: adapter changes to allow HybridStack as an FSDP unit.megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py: core changes accepting grouped HybridStack instances.tests/unit_tests/a2a_overlap/test_fsdp_hybrid_overlap.pyexercising the FSDP + grouped HybridModel forward/backward path.Why this slice
Touches 3 reviewer groups:
core-adlr,core-nemo,megatron-fsdp. Keeping FSDP separate from the hybrid feature PR avoids pulling the megatron-fsdp reviewers into the much larger #4942.Dependencies
fsdp_unit_modules=[HybridStack], which needs the HybridStack grouped-syntax changes.Validation
Validated by @Wohox as part of #4798's integrated EP-overlap smoke tests and the new
test_fsdp_hybrid_overlapintegration test.Issue tracking
Linked issue: part of #4798.
Pre-checks
test_fsdp_hybrid_overlap)🤖 Generated with Claude Code