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

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>

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>

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>

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>