[feat] Hybrid model ep overlapping main by Wohox · Pull Request #4798 · NVIDIA/Megatron-LM

Wohox · 2026-05-14T10:32:38Z

What does this PR do ?

Summary

This MR adds grouped HybridStack support and extends EP-overlap scheduling/checkpoint compatibility to grouped hybrid layers.

Key changes:

Add bracketed HybridStack group syntax, e.g. [*-], [*E], M[M*]-.
Build bracketed groups as nested HybridStack instances.
Reject invalid recursion such as nested brackets inside a group.
Keep MoE constrained to the last symbol inside a group for overlap scheduling.
Rename/migrate TransformerSchedulePlan usage to HybridStackSchedulePlan, preserving backward-compatible aliases.
Add HybridStack fine-grained callables so grouped hybrid layers can participate in combined 1F1B / EP overlap.
Make grouped HybridStack checkpoints use Transformer-compatible logical layer keys.
Make HybridModel.sharded_state_dict() match GPT behavior by dropping empty output_layer._extra_state.

Performance Improvement Parity

With TP1CP1PP4VPP2EP64 on 16 layer DSv3 test, hybird model and gpt model achieve performance improvement with EP overlap and achieves throughput parity. w&b

Checkpoint Compatibility

Grouped HybridStack checkpoint keys now canonicalize to logical Transformer layer keys. For example, [*-] maps both attention and MLP weights to:

decoder.layers.N.self_attention.*
decoder.layers.N.mlp.*

instead of using separate physical layer ids. final_norm is also sharded as final_layernorm for Transformer compatibility.

Verified both directions:

Transformer checkpoint -> HybridModel: load succeeded, one train step completed, checkpoint saved, validation/test completed.
HybridModel checkpoint -> Transformer: load succeeded, one train step completed, checkpoint saved, validation/test completed.

Cross-load status:

STATUS save_final_hybrid 0
STATUS transformer_ckpt_to_hybrid 0
STATUS final_hybrid_ckpt_to_transformer 0

Validation

unit tests:

python -m pytest -q \
  tests/unit_tests/ssm/test_hybrid_layer_allocation.py \
  tests/unit_tests/ssm/test_hybrid_block.py \
  tests/unit_tests/models/test_hybrid_model.py::TestHybridModel::test_grouped_sharded_state_dict_uses_transformer_checkpoint_keys

Result:

86 passed, 63 warnings

Also ran:

git diff --check
DeepSeek-V3 deterministic Transformer/Hybrid baseline and EP-overlap smoke tests
Transformer VPP4 baseline/overlap comparison
strict model-key checkpoint cross-load checks using --dist-ckpt-strictness raise_unexpected

Notes

The checkpoint cross-load validation intentionally uses model-weight / finetune-style loading with optimizer and RNG state skipped. Full non-finetune resume across Transformer vs grouped HybridModel is still expected to hit config arg mismatches because Transformer uses logical num_layers=16, while grouped Hybrid uses physical symbols grouped into 16 logical layers.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

copy-pr-bot · 2026-05-14T10:32:41Z

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.

github-actions · 2026-05-14T10:32:47Z

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.

@Wohox

…unk) Carries over upstream commit f6ea23b from NVIDIA#4798: when a hybrid layer pattern places MTP in a post_process VPP chunk that holds no main HybridStack layers (e.g. trailing pipe before the MTP separator), the EP-overlap schedule never invokes ``_maybe_apply_final_norm`` on the main path, so the unnormalized hidden_states feed straight into the LM head and lm_loss diverges by ~10x. Fix in ``submodule_mtp_pre_dispatch_forward``: run the main decoder's ``final_norm`` just before ``torch.chunk``, gated on ``len(model.decoder.layers) == 0`` and ``isinstance(model, HybridModel)``. The HybridModel import is deferred inside the function so this file does not gain a module-level dependency on the hybrid package — PR 1 remains independently importable. Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Wohox

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>

@Wohox

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>

@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>

@Wohox

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>

@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>

@Wohox

…unk) Carries over upstream commit f6ea23b from NVIDIA#4798: when a hybrid layer pattern places MTP in a post_process VPP chunk that holds no main HybridStack layers (e.g. trailing pipe before the MTP separator), the EP-overlap schedule never invokes ``_maybe_apply_final_norm`` on the main path, so the unnormalized hidden_states feed straight into the LM head and lm_loss diverges by ~10x. Fix in ``submodule_mtp_pre_dispatch_forward``: run the main decoder's ``final_norm`` just before ``torch.chunk``, gated on ``len(model.decoder.layers) == 0`` and ``isinstance(model, HybridModel)``. The HybridModel import is deferred inside the function so this file does not gain a module-level dependency on the hybrid package — PR 1 remains independently importable. Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Wohox

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>

@Wohox

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>

@Wohox

Move model-agnostic schedule-plan helpers out of gpt/fine_grained_callables.py into megatron/core/models/common/ so non-GPT models (HybridStack) can build the same combined-1F1B / EP-overlap schedule plans. No behavior change for GPT/MTP. - Add megatron/core/models/common/{utils.py, fine_grained_callables.py, model_chunk_schedule_plan.py} with the shared abstractions. - Reduce megatron/core/models/gpt/fine_grained_callables.py to GPT-specific pieces; reuse the new common base classes. - Switch GraphableMegatronModule.init_backward_dw_wrapper to import _BackwardDWWrapper from the new common module. - Relax combined_forward_backward_step's GPTModel-only assert to a duck-type on build_schedule_plan so any model implementing it can participate. Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Wohox

…unk) Carries over upstream commit f6ea23b from NVIDIA#4798: when a hybrid layer pattern places MTP in a post_process VPP chunk that holds no main HybridStack layers (e.g. trailing pipe before the MTP separator), the EP-overlap schedule never invokes ``_maybe_apply_final_norm`` on the main path, so the unnormalized hidden_states feed straight into the LM head and lm_loss diverges by ~10x. Fix in ``submodule_mtp_pre_dispatch_forward``: run the main decoder's ``final_norm`` just before ``torch.chunk``, gated on ``len(model.decoder.layers) == 0`` and ``isinstance(model, HybridModel)``. The HybridModel import is deferred inside the function so this file does not gain a module-level dependency on the hybrid package — PR 1 remains independently importable. Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Wohox

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>

@Wohox

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>

@Wohox

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>

@Wohox

Move model-agnostic schedule-plan helpers out of gpt/fine_grained_callables.py into megatron/core/models/common/ so non-GPT models (HybridStack) can build the same combined-1F1B / EP-overlap schedule plans. No behavior change for GPT/MTP. - Add megatron/core/models/common/{utils.py, fine_grained_callables.py, model_chunk_schedule_plan.py} with the shared abstractions. - Reduce megatron/core/models/gpt/fine_grained_callables.py to GPT-specific pieces; reuse the new common base classes. - Switch GraphableMegatronModule.init_backward_dw_wrapper to import _BackwardDWWrapper from the new common module. - Relax combined_forward_backward_step's GPTModel-only assert to a duck-type on build_schedule_plan so any model implementing it can participate. Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Wohox

…unk) Carries over upstream commit f6ea23b from NVIDIA#4798: when a hybrid layer pattern places MTP in a post_process VPP chunk that holds no main HybridStack layers (e.g. trailing pipe before the MTP separator), the EP-overlap schedule never invokes ``_maybe_apply_final_norm`` on the main path, so the unnormalized hidden_states feed straight into the LM head and lm_loss diverges by ~10x. Fix in ``submodule_mtp_pre_dispatch_forward``: run the main decoder's ``final_norm`` just before ``torch.chunk``, gated on ``len(model.decoder.layers) == 0`` and ``isinstance(model, HybridModel)``. The HybridModel import is deferred inside the function so this file does not gain a module-level dependency on the hybrid package — PR 1 remains independently importable. Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Wohox

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>

@Wohox

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>

@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>

@Wohox

Move model-agnostic schedule-plan helpers out of gpt/fine_grained_callables.py into megatron/core/models/common/ so non-GPT models (HybridStack) can build the same combined-1F1B / EP-overlap schedule plans. No behavior change for GPT/MTP. - Add megatron/core/models/common/{utils.py, fine_grained_callables.py, model_chunk_schedule_plan.py} with the shared abstractions. - Reduce megatron/core/models/gpt/fine_grained_callables.py to GPT-specific pieces; reuse the new common base classes. - Switch GraphableMegatronModule.init_backward_dw_wrapper to import _BackwardDWWrapper from the new common module. - Relax combined_forward_backward_step's GPTModel-only assert to a duck-type on build_schedule_plan so any model implementing it can participate. Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Wohox

…unk) Carries over upstream commit f6ea23b from NVIDIA#4798: when a hybrid layer pattern places MTP in a post_process VPP chunk that holds no main HybridStack layers (e.g. trailing pipe before the MTP separator), the EP-overlap schedule never invokes ``_maybe_apply_final_norm`` on the main path, so the unnormalized hidden_states feed straight into the LM head and lm_loss diverges by ~10x. Fix in ``submodule_mtp_pre_dispatch_forward``: run the main decoder's ``final_norm`` just before ``torch.chunk``, gated on ``len(model.decoder.layers) == 0`` and ``isinstance(model, HybridModel)``. The HybridModel import is deferred inside the function so this file does not gain a module-level dependency on the hybrid package — PR 1 remains independently importable. Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…ment Cherry-pick of NVIDIA#4798 onto xren's branch auto-merged hybrid_layer_allocation.py in a semantically-broken state: select_pipeline_segment_with_logical_offset (PR) forwards tp_group/dp_cp_group into select_pipeline_segment, but the merged select_pipeline_segment kept xren's older signature without those params -> TypeError at HybridModel init. Restore the PR's params + forwarding to log_on_each_pipeline_stage. Function now matches NVIDIA#4798 verbatim. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

@Wohox

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>

@Wohox

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>

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>

@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>

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>

@Wohox

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>

Wohox requested review from a team as code owners May 14, 2026 10:32

svcnvidia-nemo-ci marked this pull request as draft May 14, 2026 10:32

Wohox marked this pull request as ready for review May 20, 2026 04:12

Wohox requested review from a team as code owners May 20, 2026 04:12

svcnvidia-nemo-ci requested a review from a team May 20, 2026 04:12

svcnvidia-nemo-ci added the complexity: high label May 20, 2026

Wohox self-assigned this May 20, 2026

Phlip79 marked this pull request as draft May 28, 2026 00:35

Connor-XY force-pushed the hybrid_model_ep_overlapping_main branch from a744f05 to 0d4be06 Compare June 3, 2026 16:51

Connor-XY and others added 5 commits June 3, 2026 11:48

fix: address PR1 review comments

89e85d6

fix: address follow-up PR1 comments

c151346

fix: address Keshav review comments

eb4c325

Connor-XY and others added 8 commits June 3, 2026 12:19

fix: apply black formatting to MTP unit test

21484bc

fix: apply black formatting to hybrid callables

91e86bb

Connor-XY force-pushed the hybrid_model_ep_overlapping_main branch from 0d4be06 to 24c7f80 Compare June 3, 2026 19:21

Victarry mentioned this pull request Jun 10, 2026

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

Open

71 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[feat] Hybrid model ep overlapping main#4798

[feat] Hybrid model ep overlapping main#4798
Wohox wants to merge 13 commits into
NVIDIA:mainfrom
Wohox:hybrid_model_ep_overlapping_main

Wohox commented May 14, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented May 14, 2026

Uh oh!

github-actions Bot commented May 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Wohox commented May 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Summary

Performance Improvement Parity

Checkpoint Compatibility

Validation

Notes

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented May 14, 2026

Uh oh!

github-actions Bot commented May 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Wohox commented May 14, 2026 •

edited

Loading