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