[fix] Training wire-up for hybrid EP-overlap (4/4 of #4798) - #4944
Draft
Connor-XY wants to merge 12 commits into
Draft
[fix] Training wire-up for hybrid EP-overlap (4/4 of #4798)#4944Connor-XY wants to merge 12 commits into
Connor-XY wants to merge 12 commits into
Conversation
Connor-XY
force-pushed
the
pr4798-4-training-wireup
branch
8 times, most recently
from
June 3, 2026 19:21
0d4be06 to
24c7f80
Compare
72 tasks
Connor-XY
force-pushed
the
pr4798-4-training-wireup
branch
from
June 29, 2026 16:29
24c7f80 to
73a523e
Compare
Connor-XY
force-pushed
the
pr4798-4-training-wireup
branch
5 times, most recently
from
July 20, 2026 17:14
9801d4d to
24f94a8
Compare
Connor-XY
force-pushed
the
pr4798-4-training-wireup
branch
2 times, most recently
from
July 27, 2026 23:33
b026532 to
200c4e5
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>
Update training.py for the hybrid EP-overlap path: - Switch the MoE per-layer logging import from the deprecated ``mamba_hybrid_layer_allocation`` module to the new ``hybrid.hybrid_layer_allocation`` module, and pass only the MAIN-pattern MoE count to ``track_moe_metrics`` (the function adds ``mtp_num_layers`` internally; the old ``get_hybrid_layer_counts`` aggregated MTP, which double-counted and inflated the divisor). - When ``--use-megatron-fsdp`` is combined with ``--overlap-moe-expert-parallel-comm`` on a hybrid model, wrap the megatron-FSDP DP class with ``fsdp_unit_modules=[HybridStack]`` so grouped HybridStack participates correctly in the schedule plan. Part 4/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>
Connor-XY
force-pushed
the
pr4798-4-training-wireup
branch
from
July 28, 2026 00:22
200c4e5 to
d3bb55a
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 4 of 4 splitting #4798 by @Wohox and @Connor-XY. Original changes by @Wohox and @Connor-XY.
Summary
Two small unrelated training-script wire-ups for the hybrid EP-overlap path:
megatron.core.ssm.mamba_hybrid_layer_allocationto the newmegatron.core.models.hybrid.hybrid_layer_allocation, and pass only the MAIN-pattern MoE count totrack_moe_metrics(the function addsmtp_num_layersinternally; the oldget_hybrid_layer_countsaggregated MTP, which double-counted and inflated the divisor).--use-megatron-fsdp+--overlap-moe-expert-parallel-commare set on a hybrid model, wrapmegatron_FSDPwithfsdp_unit_modules=[HybridStack]so grouped HybridStack participates correctly in the schedule plan.Touches one file (
megatron/training/training.py): +19 / −5.Why this slice
Touches 2 reviewer groups:
training-adlr,training-nemo. Splitting this out keeps the training reviewers off the much larger feature PR (#4942) and lets this small change land in parallel.Dependencies
HybridStackfrommegatron.core.models.hybrid.hybrid_blockandparse_hybrid_patternetc. frommegatron.core.models.hybrid.hybrid_layer_allocation, both reworked in [feat] HybridStack grouped syntax + checkpoint compat + EP-overlap (2/4 of #4798) #4942.Validation
Validated by @Wohox as part of #4798's integrated training run.
Issue tracking
Linked issue: part of #4798.
Pre-checks
🤖 Generated with Claude Code