[HollowMan6]

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

What does this PR do?

Add the mHC (Manifold-Constrained Hyper-Connections) numerics used by GLM-5.3-Flash as opt-in TransformerConfig flags, and support full activation recompute with mHC-wrapped hybrid layers.

Split out of #7054 (GLM-5.3-Flash support) as discussed there; #7054 now carries only the KDA layers. Sibling PRs: #7522 (KPool DSA indexer) and #7523 (blockwise FP8 wgrad guard). All are independent of each other and based on main.

Changes

• mhc_norm_eps_inside_sqrt: compute the mapping norm as rsqrt(mean(x²) + layernorm_epsilon) instead of ‖x‖ / √K + 1e-6.
• mhc_keep_mappings_in_fp32: keep h_pre / h_post / h_res and the residual-stream mixing in FP32 until the final cast to the activation dtype.
• mhc_learned_output_contract (default True, the current behaviour): with False, the residual streams are contracted by their mean instead of the learned hc_head_* weights, in both the HybridStack post-process and the MTP layer (GLM-5.3-Flash has no learned head).
• Full activation recompute with mHC-wrapped hybrid layers: HyperConnectionHybridLayer sets supports_hybrid_recompute_kwargs, recompute.py forwards the matching kwargs (including packed_sequence_cp_metadata) without importing hybrid_block, and the NotImplementedError in TransformerConfig.__post_init__ is removed.
• use_fused_mhc raises ValueError together with either new precision flag (the fused kernels do not implement them).

All new flags default to the current behaviour. GLM-5.3-Flash uses mHC with 4 residual streams and 20 Sinkhorn iterations and has no learned output head, i.e. mhc_norm_eps_inside_sqrt=True, mhc_keep_mappings_in_fp32=True, mhc_learned_output_contract=False.

Motivation

GLM-5.3-Flash's reference implementation computes the mHC mapping norm with the epsilon inside the square root, keeps the mappings and the stream mixing in FP32, and contracts the residual streams by their mean. Matching this is needed for the Megatron model to reproduce the released checkpoint's numerics. Full recompute with mHC removes the NotImplementedError that previously blocked full activation recompute for mHC models.

Testing

• tests/unit_tests/transformer/test_mhc_precision.py (new): mappings, aggregation and mixing against a pure-torch reference in both precision modes; fused mHC rejects the new flags.
• tests/unit_tests/models/test_hybrid_mhc.py: full recompute (uniform / block × fp32 / bf16 + FP32 mixing / bf16 fused) matches forward and backward without recompute; a wrapped residual layer does not double-count its input.
• tests/unit_tests/transformer/test_hyper_connection_recompute.py: drop the obsolete "mHC + full recompute is rejected" case.

pytest tests/unit_tests/transformer/test_mhc_precision.py \
       tests/unit_tests/models/test_hybrid_mhc.py \
       tests/unit_tests/transformer/test_hyper_connection_recompute.py

⚠️ 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: Related to #6392

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• 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"

1. When your PR is ready, click Ready for Review.
2. 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]

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.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Phlip79]

Can you add some tests with/without mhc_learned_output_contract?

[HollowMan6]

Can you add some tests with/without mhc_learned_output_contract?

Thank you @Phlip79 ! I added coverage with mhc_learned_output_contract enabled and disabled, including outputs, gradients, checkpoint state, and full recomputation. Please let me know if there are any other comments, thanks!

[Phlip79]

/ok to test 9004813

[HollowMan6]

@Phlip79 Your last commit introduced some merge conflicts, would you mind retriggering the CI again? Thanks!

[Phlip79]

/ok to test d35a58d

[Phlip79]

@Phlip79 Your last commit introduced some merge conflicts, would you mind retriggering the CI again? Thanks!

Sorry about that, thanks for your quick commits!

[Phlip79]

/ok to test 106db94

[kevalmorabia97]

LGTM from the Model Optimizer side. Tested 106db944 (with #7054 merged in) and NVIDIA-NeMo/Megatron-Bridge#6044 in a NeMo 26.10 container on GB300: ModelOpt GPTQ PTQ, distillation with save + resume, and quantized HF export on a 4-layer GLM-5.3-Flash, plus full-size GLM-5.3-Flash QAD on 32 nodes (EP32 x CP4, packed SFT) with distributed-optimizer save + resume, matching the loss and grad norms of our earlier full PTQ + QAD + vLLM eval study on older heads of this stack.

NVIDIA-NeMo/Megatron-Bridge#6044 and our GLM-5.3-Flash support (NVIDIA/Model-Optimizer#2539) are waiting on this PR, so we'd appreciate a merge when possible. Thanks @HollowMan6!

[nemo-automation-bot]

🔄 Merge queue validation started!

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

[kevalmorabia97]

The merge-queue run failed because of this PR's changes, not flaky tests. The numbers changed in the way the PR intends, but the golden values weren't updated (the PR doesn't touch any golden_values files).

Why it's this PR: the merge-group commit is ed9c2c8 (main) plus this PR only. On the same base, the merge-queue runs for #4743 (run) and #7043 (run) passed all three of these tests.

Failed job	What failed	Cause in this PR
hybrid/hybrid_mhc_tp2_pp2_mtp	DETERMINISTIC lm loss / mtp_1 loss / num-zeros differ from step 6 on (largest relative error about 4.7e-5). The APPROXIMATE check (rtol 5%) passes.	Intended mHC numerics changes: the final stream contraction now accumulates in fp32 (x_streams.float().mean(dim=2)), mHC mappings are kept in fp32, and mHC recompute now replays under the original autocast state. Both tests run --recompute-modules mhc.
hybrid/hybrid_mhc_tp1_pp2_vpp2_empty_mtp	Same pattern from step 8 on (largest relative error about 7.9e-5). APPROXIMATE passes.	Same as above.
moe/deepseek_proxy_mfsdp_v2_ep2_ep_overlap	mem-allocated-bytes is 38% lower than golden at every step (543.5 MB vs 884.3 MB). Loss matches exactly, and mem-max-allocated-bytes passes.	token_dispatcher.py now sets self.probs = None in combine_postprocess, freeing the router probs after the last use. The APPROXIMATE check (rtol 5%) fails on a decrease too.
Nemo_CICD_Test	Aggregate job that fails because the jobs above failed.	—

Suggested fix: update golden_values_dev_dgx_h100.json for these three test cases with the values from this run, along with any other platform golden files these tests have (e.g. GB200), then requeue.

[HollowMan6]

Thank you @kevalmorabia97 ! I've updated the golden values, would you help with retriggering the CI and the merge queue? Thanks! cc: @Phlip79 @guihong-nv @dimapihtar

[kevalmorabia97]

/ok to test 0741315

[kevalmorabia97]

Checked 0741315: the three updated golden files match the values from merge-queue run 40308355574 exactly, and the other 15 commits are unchanged from 106db944 after the rebase. LGTM, ready to requeue.

[nemo-automation-bot]

🔄 Merge queue validation started!

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