Add mHC support for HybridModel on dev

[Connor-XY]

Ports the HybridModel/HybridStack portion of #4469 (which targeted dsv4) to dev. The mHC TransformerLayer foundation is already on dev; this PR layers the hybrid integration on top.

Summary

• Adds HyperConnectionHybridLayer, a layer-boundary mHC wrapper for HybridStack layers. It aggregates n streams to the wrapped layer's input, computes the layer delta, and feeds it through the n-stream BDA.
• Wraps standard HybridStack layers when config.enable_hyper_connections=True.
• Expands hidden states to multi-residual-stream form at the first PP stage and contracts back at the last.
• Adds _compute_mhc_block_end_plan, _build_mhc_recompute_layer_plan, and _finalize_mhc_recompute_layer on HybridStack, mirroring the caching pattern already used by TransformerBlock.
• Threads a _called_from_hybrid_mhc_wrapper kwarg through TransformerLayer.forward so the wrapper can drive a standard TransformerLayer past the existing mHC-must-use-HyperConnectionTransformerLayer guard.
• Adds HybridModel / HybridStack tests covering Mamba / attention / MLP / GDN / DSA layer types and the HC recompute plan, plus a dummy HybridModel mHC forward/backward test.

Notes

• Uses dev's fused_h_res_h_post_bda method name (PR Add mHC support for HybridModel on dsv4 #4469 renamed this method on its branch; the signature is unchanged so the rename is unnecessary on dev).
• Intentionally skipped from Add mHC support for HybridModel on dsv4 #4469: _should_call_local_cudagraph / __call__ override on HybridStack, and the CudaGraphScope.full_iteration clause in the local cudagraph guard. Dev refactored cudagraph entry in [Dev] Refactor CUDA graph API: decompose cuda_graph_scope into full_iteration impl, inference scope, and per-layer capture modules #4293 and no longer inherits GraphableMegatronModule, so those bits don't apply.

Validation

• python3 -m py_compile on all changed files
• uv run isort / uv run black clean

GPU validation (e.g. the test_dummy_hybrid_model_with_hyper_connections_forward_backward test that ran on the dsv4 branch) is still TODO on this dev port.

🤖 Generated with Claude Code

[copy-pr-bot]

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.

[Connor-XY]

/claude review

[claude]

LGTM

[Phlip79]

/claude review

[claude]

LGTM

[Phlip79]

/ok to test 656f1ec

[hxbai]

/claude strict-review

[claude]

[IMPORTANT Correctness] mhc_multistream is not saved before contraction, so MTP depth layers receive degraded single-stream data.

TransformerBlock.forward (lines 948-965) saves the pre-contraction n-stream tensor as mhc_multistream and returns it alongside hidden_states so that GPTModel can pass it to MultiTokenPredictionBlock. HybridStack.forward contracts without saving the pre-contraction tensor, and HybridModel.forward never propagates mhc_multistream.

When enable_hyper_connections=True + mtp_num_layers > 0 + HybridModel, MultiTokenPredictionBlock.forward receives mhc_multistream=None and falls back to feeding contracted [s, b, C] hidden states to its depth layers (see multi_token_prediction.py line ~1721-1726), losing per-stream diversity that TransformerBlock preserves.

Suggested fix: Mirror the TransformerBlock pattern — save hidden_states before learned_output_contract, return a tuple (hidden_states, mhc_multistream) when MTP is configured, and update HybridModel.forward to unpack and propagate mhc_multistream to the MTP block.

[claude]

[SUGGESTION Simplification] This bare assert is stripped by python -O / PYTHONOPTIMIZE=1. If a future inner layer violates the shape contract, the subtraction on line 180 would silently broadcast instead of failing with a clear message.

I see Megatron-LM uses assert extensively in similar paths, so this is consistent with codebase convention. But given this specific assertion guards against silent n-stream corruption (the subtraction could broadcast rather than crash), a conditional raise would be more robust:

if layer_output.shape != aggregated.shape:
    raise RuntimeError(
        "HyperConnectionHybridLayer requires inner layers to preserve "
        f"hidden-state shape. Got {tuple(layer_output.shape)} from inner layer "
        f"vs {tuple(aggregated.shape)} input — layer must add its own residual."
    )

[claude]

Code Review Summary

CRITICAL: 0 | IMPORTANT: 1 | SUGGESTION: 1

Important Findings

1. [IMPORTANT Correctness] mhc_multistream not propagated for MTP (hybrid_block.py:593-601)
   • HybridStack.forward contracts n-stream hidden states without saving the pre-contraction tensor. HybridModel doesn't propagate mhc_multistream to MultiTokenPredictionBlock.
   • Impact: When enable_hyper_connections=True + mtp_num_layers > 0, MTP depth layers receive degraded single-stream data instead of the full n-stream tensor that TransformerBlock preserves.
   • Fix: Mirror TransformerBlock's pattern — save pre-contraction hidden states, return (hidden_states, mhc_multistream), and update HybridModel.forward to unpack and propagate it.

Suggestions

2. [SUGGESTION Simplification] Shape assertion uses bare assert (hybrid_block.py:169)
   • The shape check guarding against silent n-stream corruption is stripped by python -O. A conditional raise RuntimeError(...) would be more robust. Consistent with codebase convention but this specific invariant protects against silent broadcast corruption rather than just a programmer error.

Verified Non-Issues

The following candidates were investigated and refuted:

• _called_from_hybrid_mhc_wrapper leaking into CUDA graph path → bool is a supported ArgMetadata type and the value is always consistent from the wrapper
• hc_head_* params remaining float32 → pre-existing TransformerBlock pattern; learned_output_contract explicitly casts to fp32 internally
• tp_group fallback for MambaLayer wrapping → pre-existing MambaLayer behavior; the fallback is intentional and returns the same group
• Double mHC wrapping with HyperConnectionTransformerLayer → built-in hybrid specs never use this class; would require deliberate custom construction
• MoE router distribution change → architecturally equivalent to HyperConnectionTransformerLayer path in TransformerBlock

Overall Assessment

Low risk. The core mHC integration — HyperConnectionHybridLayer wrapper, delta computation, dtype handling, recompute plan, contraction/expansion, and sharded_state_dict — is well-implemented and correctly mirrors the TransformerBlock/HyperConnectionTransformerLayer patterns. The fp32 residual connection handling, the boundary BDA manager skipping, and the pipeline parallel shape management are all correct. The one actionable finding (MTP multistream propagation) affects a specific feature combination that the PR description notes as not yet GPU-validated. Test coverage is solid for the core paths.

🤖 Generated with Claude Code

[Victarry]

Hi，@Connor-XY, Thanks for this PR! It look good to me overall. Please also consider the above comment from Claude Code.
#4949 (comment)

[Phlip79]

/ok to test 7f8faa0

[Phlip79]

Hi，@Connor-XY, Thanks for this PR! It look good to me overall. Please also consider the above comment from Claude Code. #4949 (comment)

I resolved this because @Connor-XY is OOO this week.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26530436719