[Split 2/N of #3430] feat(mHC): GPT model wiring by Connor-XY · Pull Request #4945 · NVIDIA/Megatron-LM

Connor-XY · 2026-05-22T17:15:12Z

What does this PR do?

Second of a sequence of split PRs derived from #3430 (mirror of upstream #2943). The goal of splitting is to keep each PR small and limit the number of CODEOWNERS reviewer groups required, making review faster and easier.

This PR contains the GPT model wiring layer of mHC — i.e., how the core mHC module (added in Split 1) gets selected by the GPT layer specs and reached from the training-time initialization path. No new mHC math lives in this PR; everything is plumbing.

Files

`megatron/core/models/gpt/gpt_layer_specs.py` — adds `enable_hyper_connection` kwarg to both the TE and local submodule builders; spec helper returns a `ModuleSpec` whose `module` is `HyperConnectionTransformerLayer` when the kwarg is set, and patches in `HyperConnectionModule` for `self_attention_hyper_connection` / `mlp_hyper_connection`.
`megatron/core/models/gpt/experimental_attention_variant_module_specs.py` — forwards the same kwarg through the experimental variant pathway.
`gpt_builders.py` — threads `enable_hyper_connection=config.enable_hyper_connections` from the runtime config into the spec builder.
`megatron/training/initialize.py` — enables NVTX range profiling early (before mHC modules import), so `@nvtx_decorator` captures the right value at decoration time.
`tests/unit_tests/models/test_gpt_layer_specs.py` — new test file: verifies the `enable_hyper_connection` branch produces a spec whose module is `HyperConnectionTransformerLayer`.
`tests/unit_tests/test_fp8_param.py` — minor mHC-related test wiring.

Reviewer groups touched: core-adlr, core-nemo, gpt, training-adlr, training-nemo.

Dependency

Depends on [Split 1/N of #3430] feat(mHC): basic pytorch implementation of manifold hyper connection #4531 (Split 1, core mHC transformer module). `gpt_layer_specs.py` here references `HyperConnectionTransformerLayer` and `HyperConnectionModule` symbols introduced in [Split 1/N of #3430] feat(mHC): basic pytorch implementation of manifold hyper connection #4531. CI on this PR will fail until [Split 1/N of #3430] feat(mHC): basic pytorch implementation of manifold hyper connection #4531 is merged into `main`.

Deferred to follow-up split PRs (NOT in this PR)

Split 3: pipeline-parallel support — `pipeline_parallel/schedules.py`, `tests/unit_tests/pipeline_parallel/test_pp_mhc_compatibility.py`
Split 4: fused mHC cuTile kernels — `fused_mhc_kernels.py`, kernel-side additions to `hyper_connection.py` / `transformer_config.py`, `tests/unit_tests/fusions/test_fused_mhc_kernels.py`
Split 5: functional-test recipe — `tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/*`, `tests/test_utils/recipes/h100/gpt.yaml`

Credit

All changes here are a re-slice of @jingqiny-99's work in #3430 (upstream #2943). No new functionality has been added, removed, or rewritten in this split — only repackaged for narrower review scope.

Pre-checks

Re-slice of an existing PR; no new feature
Original unit tests included for components touched
Linted with black + isort
PR opened as draft per repo contributing guidelines

copy-pr-bot · 2026-05-22T17:15:17Z

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

Third split of NVIDIA#3430 (mirror of upstream NVIDIA#2943). Extends the pipeline-parallel schedules so they handle the n-stream hidden-state shape used by mHC across PP boundaries, and adds the corresponding unit tests. Files: - megatron/core/pipeline_parallel/schedules.py: n-stream-aware tensor-shape computation in `get_tensor_shapes`, and the matching expand/contract handling at first/last PP stages. - tests/unit_tests/pipeline_parallel/test_pp_mhc_compatibility.py: coverage for `get_tensor_shapes`, `get_num_layers_to_build`, VPP edge cases, PP boundary shape consistency, and standalone embedding/loss layers under mHC. Depends on NVIDIA#4531 (Split 1) for the core mHC module symbols, and on NVIDIA#4945 (Split 2) for the spec-side `enable_hyper_connection` wiring used by the forward tests. Reviewer groups touched: core-adlr, core-nemo, pipeline-parallelism. Subsequent splits still pending: - Split 4: fused mHC cuTile kernels - Split 5: functional-test recipe Original work by @jingqiny-99 in NVIDIA#3430 (upstream NVIDIA#2943). Co-authored-by: jingqiny-99 <jingqiny@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com>

@jingqiny-99

Final split of NVIDIA#3430 (mirror of upstream NVIDIA#2943 + kernel-fusion NVIDIA#3828). Adds the end-to-end functional test for the mHC feature: Files: - tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/ model_config.yaml — TP=2, PP=2 GPT-3 mHC training recipe golden_values_dev_dgx_h100.json — golden metrics on dgx_h100 - tests/test_utils/recipes/h100/gpt.yaml — registers the new test case under the mr / mr-github scopes for dgx_h100. Golden values picked up from NVIDIA#4469 (dsv4 branch) where the recipe has been calibrated against the final fused-kernel-on configuration. Reviewers can recalibrate if needed once Splits 1-4 land. Depends on NVIDIA#4531 (Split 1), NVIDIA#4945 (Split 2), and NVIDIA#4947 (Split 4) — the recipe enables enable_hyper_connections via the GPT spec and exercises the fused kernel path, so all three feature splits must land first. Reviewer groups touched: ci. Original work by @jingqiny-99 in NVIDIA#3430 (upstream NVIDIA#2943). Co-authored-by: jingqiny-99 <jingqiny@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com>

@jingqiny-99

Third split of NVIDIA#3430 (mirror of upstream NVIDIA#2943). Extends the pipeline-parallel schedules so they handle the n-stream hidden-state shape used by mHC across PP boundaries, and adds the corresponding unit tests. Files: - megatron/core/pipeline_parallel/schedules.py: n-stream-aware tensor-shape computation in `get_tensor_shapes`, and the matching expand/contract handling at first/last PP stages. - tests/unit_tests/pipeline_parallel/test_pp_mhc_compatibility.py: coverage for `get_tensor_shapes`, `get_num_layers_to_build`, VPP edge cases, PP boundary shape consistency, and standalone embedding/loss layers under mHC. Depends on NVIDIA#4531 (Split 1) for the core mHC module symbols, and on NVIDIA#4945 (Split 2) for the spec-side `enable_hyper_connection` wiring used by the forward tests. Reviewer groups touched: core-adlr, core-nemo, pipeline-parallelism. Subsequent splits still pending: - Split 4: fused mHC cuTile kernels - Split 5: functional-test recipe Original work by @jingqiny-99 in NVIDIA#3430 (upstream NVIDIA#2943). Co-authored-by: jingqiny-99 <jingqiny@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

@jingqiny-99

Final split of NVIDIA#3430 (mirror of upstream NVIDIA#2943 + kernel-fusion NVIDIA#3828). Adds the end-to-end functional test for the mHC feature: Files: - tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/ model_config.yaml — TP=2, PP=2 GPT-3 mHC training recipe golden_values_dev_dgx_h100.json — golden metrics on dgx_h100 - tests/test_utils/recipes/h100/gpt.yaml — registers the new test case under the mr / mr-github scopes for dgx_h100. Golden values picked up from NVIDIA#4469 (dsv4 branch) where the recipe has been calibrated against the final fused-kernel-on configuration. Reviewers can recalibrate if needed once Splits 1-4 land. Depends on NVIDIA#4531 (Split 1), NVIDIA#4945 (Split 2), and NVIDIA#4947 (Split 4) — the recipe enables enable_hyper_connections via the GPT spec and exercises the fused kernel path, so all three feature splits must land first. Reviewer groups touched: ci. Original work by @jingqiny-99 in NVIDIA#3430 (upstream NVIDIA#2943). Co-authored-by: jingqiny-99 <jingqiny@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

@jingqiny-99

Third split of NVIDIA#3430 (mirror of upstream NVIDIA#2943). Extends the pipeline-parallel schedules so they handle the n-stream hidden-state shape used by mHC across PP boundaries, and adds the corresponding unit tests. Files: - megatron/core/pipeline_parallel/schedules.py: n-stream-aware tensor-shape computation in `get_tensor_shapes`, and the matching expand/contract handling at first/last PP stages. - tests/unit_tests/pipeline_parallel/test_pp_mhc_compatibility.py: coverage for `get_tensor_shapes`, `get_num_layers_to_build`, VPP edge cases, PP boundary shape consistency, and standalone embedding/loss layers under mHC. Depends on NVIDIA#4531 (Split 1) for the core mHC module symbols, and on NVIDIA#4945 (Split 2) for the spec-side `enable_hyper_connection` wiring used by the forward tests. Reviewer groups touched: core-adlr, core-nemo, pipeline-parallelism. Subsequent splits still pending: - Split 4: fused mHC cuTile kernels - Split 5: functional-test recipe Original work by @jingqiny-99 in NVIDIA#3430 (upstream NVIDIA#2943). Co-authored-by: jingqiny-99 <jingqiny@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

@jingqiny-99

Final split of NVIDIA#3430 (mirror of upstream NVIDIA#2943 + kernel-fusion NVIDIA#3828). Adds the end-to-end functional test for the mHC feature: Files: - tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/ model_config.yaml — TP=2, PP=2 GPT-3 mHC training recipe golden_values_dev_dgx_h100.json — golden metrics on dgx_h100 - tests/test_utils/recipes/h100/gpt.yaml — registers the new test case under the mr / mr-github scopes for dgx_h100. Golden values picked up from NVIDIA#4469 (dsv4 branch) where the recipe has been calibrated against the final fused-kernel-on configuration. Reviewers can recalibrate if needed once Splits 1-4 land. Depends on NVIDIA#4531 (Split 1), NVIDIA#4945 (Split 2), and NVIDIA#4947 (Split 4) — the recipe enables enable_hyper_connections via the GPT spec and exercises the fused kernel path, so all three feature splits must land first. Reviewer groups touched: ci. Original work by @jingqiny-99 in NVIDIA#3430 (upstream NVIDIA#2943). Co-authored-by: jingqiny-99 <jingqiny@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

@jingqiny-99

Second split of NVIDIA#3430 (mirror of upstream NVIDIA#2943). Wires the mHC core module added in Split 1 (NVIDIA#4531) into the GPT layer specs: - megatron/core/models/gpt/gpt_layer_specs.py: add enable_hyper_connection kwarg; build mHC-aware ModuleSpec selecting HyperConnectionTransformerLayer - megatron/core/models/gpt/experimental_attention_variant_module_specs.py: forward the kwarg through experimental specs - gpt_builders.py: thread enable_hyper_connection from config to spec - tests/unit_tests/models/test_gpt_layer_specs.py: coverage for the enable_hyper_connection branch - tests/unit_tests/test_fp8_param.py: minor mHC-related test wiring Depends on NVIDIA#4531 — this PR is mergeable only after Split 1 lands. Rebase note (onto main after NVIDIA#4531 / NVIDIA#6401 merged): the original megatron/training/initialize.py hunk (enable NVTX ranges at init when --profile is set) is dropped — main's nvtx_decorator now checks the flag at call time and training.py toggles it at profile_step_start under --nvtx-ranges. The fused-MLA (mla_down_proj_fusion) TE submodules added upstream are wired for mHC as well. Reviewer groups touched: core-adlr, core-nemo, gpt, training-adlr, training-nemo. Subsequent splits (still pending Split 1's merge): - Split 3: pipeline-parallel support - Split 4: fused mHC cuTile kernels - Split 5: functional-test recipe Original work by @jingqiny-99 in NVIDIA#3430 (upstream NVIDIA#2943). Co-authored-by: jingqiny-99 <jingqiny@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Split 1 (NVIDIA#4531) renamed the config fields to `enable_mhc_connections` / `mhc_num_residual_streams` per review. This split still read the old `config.enable_hyper_connections`, which would raise AttributeError once Split 1 lands. Sweep the spec builders, gpt_builders and tests onto the new names, and rename the `enable_hyper_connection` spec kwarg (plus the `enable_hc` local) to match, so the whole mHC surface uses one vocabulary. No behavior change beyond the rename. Signed-off-by: Yan Xu <yxu1@nvidia.com>

Connor-XY · 2026-09-10T18:33:05Z

included in other PR, no longer needed

This was referenced May 22, 2026

[Split 3/N of #3430] feat(mHC): pipeline-parallel support #4946

Closed

[Split 4/N of #3430] feat(mHC): fused cuTile kernels #4947

Closed

[Split 5/N of #3430] feat(mHC): functional-test recipe #4948

Closed

Connor-XY force-pushed the mhc-pr2-gpt-wiring branch from 7fdbed4 to b26c6fd Compare June 12, 2026 17:42

Connor-XY force-pushed the mhc-pr2-gpt-wiring branch from b26c6fd to e1a40d6 Compare June 12, 2026 17:48

sbhavani mentioned this pull request Jul 12, 2026

Feature Request: mHC (Manifold-Constrained Hyper-Connections) #2890

Closed

Connor-XY force-pushed the mhc-pr2-gpt-wiring branch from e1a40d6 to af41e41 Compare August 13, 2026 18:25

lhb8125 mentioned this pull request Aug 27, 2026

Megatron-LM Monthly Activity Report — August 2026 (Final, through August 31, 2026) #6921

Open

Connor-XY and others added 2 commits September 8, 2026 10:36

Connor-XY force-pushed the mhc-pr2-gpt-wiring branch from af41e41 to 1543c57 Compare September 8, 2026 20:30

Connor-XY closed this Sep 10, 2026

Navigation Menu

Repository navigation

[Split 2/N of #3430] feat(mHC): GPT model wiring - #4945

[Split 2/N of #3430] feat(mHC): GPT model wiring#4945
Connor-XY wants to merge 2 commits into
NVIDIA:mainfrom
Connor-XY:mhc-pr2-gpt-wiring

Connor-XY commented May 22, 2026

Uh oh!

copy-pr-bot Bot commented May 22, 2026

Uh oh!

Connor-XY commented Sep 10, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

Connor-XY commented May 22, 2026

What does this PR do?

Files

Dependency

Deferred to follow-up split PRs (NOT in this PR)

Credit

Pre-checks

Uh oh!

copy-pr-bot Bot commented May 22, 2026

Uh oh!

Connor-XY commented Sep 10, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant