[dev]: Fix mHC boundaries in EP overlap schedule by jingqiny-99 · Pull Request #5471 · NVIDIA/Megatron-LM

jingqiny-99 · 2026-06-24T04:27:01Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do ?

What changed

Fix mHC boundary and recompute handling with EP overlap.
Keep mHC post-processing in the combine node by default.
Add --mhc-post-on-compute-stream to optionally schedule mHC post as a standalone compute-stream node.
Preserve PP, MTP, and FSDP compatibility.

mHC post and mHC recompute have no dedicated stream; they reuse the two EP-overlap streams: comp (get_comp_stream()) and comm (EP all-to-all / PP p2p, get_comm_stream()). All per-layer 1F1B scheduling is in TransformerLayerSchedulePlan.run (model_chunk_schedule_plan.py:264-342); combined_1f1b.py is unchanged vs dev.

mHC post — two modes, gated by mhc_post_on_compute_stream:

Default (False): post is fused inside the moe_combine node on the comm stream (the mhc_post node is a NoopScheduleNode). Forward runs in combine_fwd, backward in combine_bwd. (fine_grained_callables.py:703,725)

Compute-stream (True): post is a standalone node on the comp stream — post.F after combine_fwd (run:331), post.B before combine_bwd (run:297).

mHC recompute — only when recompute_granularity="selective" and "mhc" in recompute_modules. It is a ScheduleNode on the comp stream, created only on the last layer of each recompute group (model_chunk_schedule_plan.py:186-199). In backward it runs first for that layer, immediately before the combine/post backward (run:295-296), replaying the whole group's forward. Since backward walks layers in descending order, the group's last layer runs first, so the replay precedes the group's per-layer backward. NVTX: mhc/recompute/group_{i}/B.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

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 @NVIDIA/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.

copy-pr-bot · 2026-06-24T04:27:05Z

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.

jingqiny-99 · 2026-06-24T04:27:32Z

/claude strict-review

claude

Review Summary

CRITICAL: 0 | IMPORTANT: 2 | SUGGESTION: 4

This PR correctly fixes mHC (Hyper Connections) boundary handling in the EP overlap schedule by:

Extracting preprocess_for_layer_schedule / postprocess_for_layer_schedule from TransformerBlock.forward() into reusable methods
Wiring these into the fine-grained EP overlap callables (PreProcessNode, PostProcessNode, submodule_combine_forward)
Adding mHC-aware paths in submodule_attn_forward (hyper_connection + 6-tuple return) and submodule_combine_forward (post-MLP fusion via _forward_mhc_mlp_post)
Correctly threading mhc_multistream through MTP layers in the overlap schedule

Key verification points (all pass):

mHC contraction at decoder boundary: postprocess_for_layer_schedule correctly applies learned_output_contract + final_layernorm at the last decoder layer in the EP overlap schedule — this was the missing piece
MTP multi-stream threading: submodule_mtp_attn_forward correctly uses pre-contraction mhc_multistream for MTP input; submodule_mtp_postprocess_forward preserves multi-stream format between MTP layers while storing contracted tensors for the final concatenation
Residual handling: mHC residual from mlp_hyper_connection is correctly preserved through the fused norm path (not overwritten by TEFusedResidualRMSNorm output)
Stream safety: record_stream calls added for mlp_h_res and mlp_hc_h_post tensors crossing CUDA streams
Normal forward path equivalence: TransformerBlock.forward() now delegates to the same extracted methods — behavior is identical to the old inline code

IMPORTANT findings (advisory):

PostProcessNode MTP guard removal — the old not mtp_num_layers check was redundant (empty-decoder MTP stages always have final_layernorm=None), but the implicit invariant should be documented
MTP postprocess multi-stream logic — the next_hidden_states pattern is correct but subtle; a brief comment would help future readers

SUGGESTION findings:

postprocess_for_layer_schedule return type is polymorphic in a way that the return_mhc_multistream flag name doesn't fully communicate
is_mhc_layer detection via hasattr(layer, '_forward_mhc_mlp_post') is fragile; isinstance(layer, HyperConnectionTransformerLayer) would be more robust
len(forward_outputs) == 6 guard is defensive but unnecessary given CUDA graphs and EP overlap are mutually exclusive
_forward_mhc_mlp_post partially duplicates _forward_post_mlp_with_fused_hyper_connection

Overall risk: Low. The refactoring correctly unifies boundary logic between the normal and EP-overlap forward paths. The core fix (mHC contraction + final layernorm at decoder boundary in the overlap schedule) is sound. Tests cover the new boundary helpers and integration with PreProcessNode/PostProcessNode. LGTM.

jingqiny-99 · 2026-06-24T07:08:05Z

/claude strict-review

claude

Review Summary

CRITICAL: 0 | IMPORTANT: 0 | SUGGESTION: 3

Thorough review of all 5 changed files (396 additions, 102 deletions). This PR fixes mHC (multi-stream Hyper Connections) boundaries in the EP overlap schedule by:

Extracting shared boundary logic from TransformerBlock.forward() into preprocess_for_layer_schedule() and postprocess_for_layer_schedule(), enabling reuse by both the normal forward path and the EP overlap schedule's PreProcessNode/PostProcessNode/submodule_combine_forward.
Handling mHC tensors through the EP overlap split graph: mlp_h_res and mlp_hc_h_post are properly detached, recorded to streams, and cleaned up in submodule_attn_forward and submodule_combine_forward. The new _forward_mhc_mlp_post() method correctly applies the fused H_res/H_post/BDA without the recompute manager (consistent with EP overlap not supporting mHC recompute).
MTP + mHC in EP overlap: mhc_multistream is correctly saved from postprocess_for_layer_schedule and used by MTP layers to receive pre-contraction multi-stream input, while _postprocess contracts for mtp_hidden_states.

Verification

Backward compatibility: Non-mHC models follow identical code paths—is_mhc_layer = is_moe and isinstance(layer, HyperConnectionTransformerLayer) correctly gates all new logic.
Data flow: Traced mHC tensor flow through attention HC → mlp_hyper_connection → detach → dispatch → experts → combine → _forward_mhc_mlp_post → postprocess_for_layer_schedule. All tensor shapes and residual connections are consistent with the non-EP-overlap HyperConnectionTransformerLayer.forward().
Unused variables: All new identifiers (mhc_multistream, is_mhc_layer, mlp_h_res/mlp_hc_h_post on layer_state, _forward_mhc_mlp_post_core, _forward_mhc_mlp_post, preprocess_for_layer_schedule, postprocess_for_layer_schedule) have meaningful read/use paths verified via grep.
Process groups: No new direct parallel_state.get_*_group() calls added in megatron/core production code.
Tests: Good coverage of boundary helpers and schedule node integration with mock-based unit tests.

Well-structured PR. LGTM.

jingqiny-99 · 2026-06-25T07:24:45Z

/claude strict-review

claude · 2026-06-25T07:37:30Z

[IMPORTANT Correctness] CUDA graph (moe_router) + EP overlap + mHC return mismatch

After this PR, submodule_attn_forward in fine_grained_callables.py distinguishes mHC from non-mHC layers: for is_mhc_layer=True it unpacks 6 values from forward_outputs (line 605-613), but HyperConnectionTransformerLayer._te_cuda_graph_replay_impl (in transformer_layer.py:2321-2325) still returns only 4 values on the EP-overlap early-return path — it drops mlp_h_res and mlp_hc_h_post even though they were extracted from the graph output at lines 2282-2284.

If someone enables CUDA graph (moe_router) + EP overlap + mHC, this will crash with not enough values to unpack.

Before this PR, mHC + EP overlap was silently incorrect (missing hyper-connection processing). This PR correctly fixes the non-CUDA-graph path. To also cover the CUDA-graph path, either:

Option A — extend the return in HyperConnectionTransformerLayer._te_cuda_graph_replay_impl:2325:

if self.config.overlap_moe_expert_parallel_comm:
    probs, routing_map = self.mlp.route(hidden_states)
    hidden_states, probs = self.mlp.preprocess(hidden_states, probs, routing_map)
    nvtx_range_pop(suffix="mlp")
    return residual, hidden_states, probs, shared_expert_output, mlp_h_res, mlp_hc_h_post

Option B — add a guard assertion:

if self.config.overlap_moe_expert_parallel_comm:
    assert not self.config.enable_hyper_connections, (
        "CUDA graph (moe_router) + EP overlap + mHC is not yet supported."
    )
    ...

Option A is preferable if the rest of the pipeline handles it correctly (which it should, given this PR's changes to the combine path). Option B is a simpler stopgap.

claude · 2026-06-25T07:37:50Z

Code Review Summary

PR: Fix mHC boundaries in EP overlap schedule

Findings: CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 2

Most impactful finding

[IMPORTANT Correctness] HyperConnectionTransformerLayer._te_cuda_graph_replay_impl returns 4 values on the EP-overlap path but the new is_mhc_layer branch in submodule_attn_forward expects 6. This introduces a crash for the CUDA graph (moe_router) + EP overlap + mHC configuration. See detailed comment above for fix options.

Overall assessment

Low risk. This is a well-structured refactoring that correctly consolidates boundary logic (mHC expand/contract + final layernorm) into shared preprocess_for_layer_schedule / postprocess_for_layer_schedule helpers on TransformerBlock, then uses them from both the normal forward() path and the EP overlap scheduled path.

Key observations from the review:

The refactoring preserves identical behavior for the normal forward path (verified by diff tracing).
The mHC support in EP overlap (submodule_attn_forward, submodule_combine_forward) correctly mirrors the normal HyperConnectionTransformerLayer._forward_mlp data flow: mlp_hyper_connection → layernorm → route/dispatch → combine → _forward_mhc_mlp_post.
The residual handling for mHC + fused residual norm (TEFusedResidualRMSNorm) is correct: line 577 preserves the multi-stream residual from mlp_hyper_connection instead of using the fused norm's residual.
The mhc_multistream flow for MTP in the EP overlap schedule correctly saves the pre-contraction multi-stream and passes it to MTP layers, matching the normal forward path.
The _forward_mhc_mlp_post / _forward_mhc_mlp_post_core extraction is clean — the surrounding checkpoint management and norm offload are correctly handled by the EP overlap schedule's submodule_moe_forward and submodule_combine_forward respectively.
Unit tests cover the boundary helpers and pre/post process nodes with good assertions on shape and reference identity.

jingqiny-99 · 2026-07-08T08:31:11Z

/claude strict-review

claude · 2026-07-08T08:40:47Z

Strict Review Summary

Findings by severity — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1

Overall assessment: Low risk

This PR fixes mHC (hyper-connection) boundary and recompute handling under EP-overlap scheduling, and adds an opt-in --mhc-post-on-compute-stream node. It is unusually well-scoped and well-tested for a change touching the fine-grained overlap scheduler.

Strengths verified during review:

Correctness / numerical equivalence is covered directly: TestMhcA2AOverlapNumerics compares the overlap schedule (and the interleaved two-in-flight-plan case) against the eager hook path for forward output and every parameter gradient, across the merged / compute-stream × recompute / no-recompute matrix. This is the right guard for a scheduler change.
Backward compatibility is preserved carefully: the new config field mhc_post_on_compute_stream defaults to False, keeping the existing combine-stream topology unchanged; the block-boundary refactor extracts preprocess_for_layer_schedule / postprocess_for_layer_schedule and the normal forward() delegates to them, so the non-scheduled path is behavior-identical (verified by diff). New __post_init__ guards fail loudly for the unsupported CUDA-graph + EP-overlap + mHC-recompute combination rather than silently misbehaving.
Recompute state machine: CheckpointManager.recompute_now() / discard_all_outputs() are guarded by _recomputed / _outputs_discarded idempotency flags, raise if replayed before discard, and per-group managers are kept independent across in-flight microbatches (asserted in tests). The _unified_recompute_hook now delegates to recompute_now(), unifying the hook-driven and scheduler-driven paths.
PP / MTP / FSDP compatibility is exercised: boundary-helper tests confirm first-PP-stage expansion, non-first-stage pass-through, pre-final-layernorm contraction, and empty-decoder handling; MTP retains the optional mHC-post callable slot and correctly threads mhc_multistream pre-contraction state.
No new direct parallel_state.get_*_group() reads were introduced in megatron/core.

Only finding (non-blocking): mhc_recompute_group_start / mhc_recompute_group_end are added to extra_args and asserted in a test but have no runtime consumer — see inline comment. Purely a cleanliness/drift concern.

Recommend addressing the minor cleanup at the author's discretion; nothing here blocks merge from a correctness or performance standpoint.

jingqiny-99 · 2026-07-09T08:47:00Z

/ok to test 0bc168c

Signed-off-by: Jingqin Yang <jingqiny@nvidia.com>

Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

Signed-off-by: Jingqin Yang <jingqiny@nvidia.com>

…line.svg fix: delete useless figs generated by AI Signed-off-by: Jingqin Yang <jingqiny@nvidia.com>

The fine-grained EP-overlap schedule applied the decoder block boundary (mHC output contraction + final layer norm, i.e. postprocess_for_layer_schedule) only inside the MoE combine / mHC-post callables. For a mixed pattern whose final decoder layer is dense (for example moe_layer_freq=[1, 0]) the terminal node is the dense MLP, so the boundary was skipped and an uncontracted [s, b, n*h] tensor reached GPT postprocessing without learned_output_contract or the final layer norm. Factor the boundary into finalize_decoder_layer_output() and attach it to whichever node is terminal for the layer (mhc_post, combine, or the dense mlp), keeping the math independent of layer type and mHC-post placement. The boundary's pre-contraction mHC multi-stream (consumed by MTP) is now detached at its producer so MTP backward cannot leave its schedule node and traverse the decoder mHC graph out of order (which could trigger a second-backward error after saved tensors are freed, or bypass the point where the contracted and MTP gradient branches merge). The producer node's backward_impl reconnects the accumulated gradient, matching the existing residual / mlp_h_res / mlp_hc_h_post bridge. Add mHC + EP-overlap regression tests comparing scheduled vs eager forward output and parameter gradients: a dense final layer (moe_layer_freq=[1, 0]) and MTP (mtp_num_layers=1, default mtp_detach_heads=False), each for merged and compute-stream mHC-post, plus a two-in-flight MTP variant that checks the MTP side gradient stays bound to the correct microbatch. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

…d tolerance The mHC HyperConnection aggregate() returns fp32 (for mixing-weight numerical stability) without downcasting. A dense final layer feeds that straight into a low-precision TE MLP, which raises "Data types for parameters must match", so moe_layer_freq=[1, 0] with mHC could not run in eager or the fine-grained schedule, independent of the decoder-boundary finalization. Restore the activation dtype for dense (non-MoE) mHC layers only, leaving the MoE routing dtype path untouched. Also relax the gradient tolerance in the two-in-flight MTP test: the tied word-embedding gradient is reduced in a different bf16 order under interleaving than under sequential eager execution. The MTP-specific parameter gradients are bitwise-identical, confirming the MTP side gradient stays bound to the correct microbatch. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

The previous guard compared the aggregated MLP input against the layer's input dtype, but the mHC residual stream is already fp32 by the time it reaches a dense layer (pre_mlp_layernorm is IdentityOp there), so the guard was a no-op and moe_layer_freq=[1, 0] with mHC still raised the TE "Data types for parameters must match" error. Cast the dense MLP input to params_dtype instead. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

… path The flag toggled MLP-side mHC post-processing between the default merged path (inside the MoE combine node, on the communication stream) and a separate compute-stream node. The separate path had no demonstrated benefit, so drop the flag and the entire separate-node machinery, keeping only the merged default. This removes the config field (and its auto-generated --mhc-post-on-compute-stream CLI argument) plus its validation, the use_separate_mhc_post branch, the mhc_post forward slot and schedule node in both the transformer and MTP callable builders and the model chunk schedule plan, and all flag-specific tests. MLP-side mHC post-processing continues to run inline in the combine node via submodule_mhc_post_forward. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

Signed-off-by: Jingqin Yang <jingqiny@nvidia.com>

Under hyper-connections MultiTokenPredictionLayer builds separate e_proj/h_proj and sets eh_proj to None, so the unconditional eh_proj append in build_mtp_layer_callables placed None in the attn node's delayed-wgrad list (AttributeError in backward_dw once delay_wgrad_compute is enabled) and left the e_proj/h_proj weight gradients without a delayed-wgrad trigger. Select [e_proj, h_proj] when hyper-connections are enabled and keep [eh_proj] otherwise. The unconditional append predates this PR; the mHC MTP layout it introduces is what made it reachable. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

…tion CUDA graph - Add a numerical test for mHC + MTP + EP overlap + delay_wgrad_compute that executes the delayed-wgrad backward path (crashes without the e_proj/h_proj callable fix and catches silently deferred wgrads via gradient-key equality). - Parameterize both MTP schedule tests over recompute=[False, True] to exercise mHC + MTP + EP overlap + mhc recompute. - Add a full-iteration CUDA-graph capture/replay test (no mhc recompute) proving the new config guard admits CG + EP overlap and that the scheduled step is capturable and numerically faithful; MoE runs drop_and_pad since the dropless dispatcher's D2H splits sync is illegal during capture. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

The test builds TransformerConfig(bf16=True) without params_dtype (fp32 default) and then converts the model with .bfloat16(), a combination real training cannot produce (validate_args derives params_dtype from --bf16). The inconsistency defeats the dense-mHC params_dtype cast in _forward_mlp, so TE rejects the fp32 activations against bf16 params. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

With delay_wgrad_compute, every TE linear built from the shared config defers its weight gradient until an explicit backward_dw() call, but DSv4HybridSelfAttention.backward_dw() and MLASelfAttention.backward_dw() never traversed core_attention. The CSA compressor/indexer linears (linear_wkv, linear_wgate, linear_wq_b, linear_weights_proj — six per ratio-4 layer counting the indexer's own compressor) and the DSA indexer linears (linear_wq_b, linear_wk, linear_weights_proj) therefore kept their weight gradients deferred forever: the parameters silently never trained and their wgrad activation stashes were retained. Add backward_dw() to Compressor, CSAIndexer, CompressedSparseAttention (None-guards mirroring __init__ optionality), DSAIndexer, and DSAttention, and flush core_attention from both attention classes. The MLA call is guarded because standard MLA cores (TEDotProductAttention) own no linears and define no backward_dw; the DSv4 hybrid core is invariantly CompressedSparseAttention, so its call is unguarded. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

Single-GPU module tests: forward + backward with delay_wgrad_compute leaves every nested TE linear's weight grad deferred (grad is None); attention.backward_dw() must flush all of them — six CSA linears (compressor, indexer, and the indexer's own compressor) for the dsv4_hybrid variant and three DSA indexer linears for the dsa variant. dsa_indexer_loss_coeff=1.0 wires real gradients into the indexer linears, whose only gradient source is the indexer loss. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

The recompute-group counter restarts per module (decoder and mtp both build layer schedule plans), so fold a module tag into the NVTX label to keep the two apart in profiles. Also trim the viewless-tensor comment in preprocess_for_layer_schedule, which still referenced a transpose/contiguous sequence and an 'else' branch from its original location in forward(). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

jingqiny-99 · 2026-07-19T23:24:25Z

/ok to test 7b548c7

…erlap-dev Resolve conflict in fine_grained_callables.py: keep both dev's ncclep backend flag and the mHC layer checks in build_transformer_layer_callables. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>

jingqiny-99 · 2026-07-21T07:13:54Z

/ok to test b5d6a28

Signed-off-by: Jingqin Yang <jingqiny@nvidia.com>

jingqiny-99 · 2026-07-22T03:23:54Z

/ok to test db98830

svcnvidia-nemo-ci · 2026-07-22T13:39:36Z

🔄 Merge queue validation started!

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

jingqiny-99 changed the title ~~Fix mHC boundaries in EP overlap schedule~~ [dev]: Fix mHC boundaries in EP overlap schedule Jun 24, 2026