[deepakn94]

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

What does this PR do?

MTP's input projections are the only weight-bearing modules in that path built without a pg_collection:

# multi_token_prediction.py, MHC path
"tp_group": pg_collection.tp if pg_collection is not None else None,
# and the eh_proj path
tp_group=pg_collection.tp if pg_collection is not None else None,

Everything else in the path passes the collection through: mtp_model_layer, and the block's layers. enorm, hnorm and final_layernorm take no group arguments at all, but they are 1-D norm weights and are not GTP-sharded anywhere.

ColumnParallelLinear and the TE equivalents read their GTP axis with resolve_gtp_remat_group(pg_collection, is_expert). With pg_collection=None that skips the attr in vars(...) branch and falls back to ProcessGroupCollection.use_mpu_process_groups(...):

• non-MIMO: the MPU globals are initialized, the fallback returns the real gtp_remat group, and the projection is sharded. The missing argument is invisible.
• MIMO: examples/mimo/pretrain_mimo.py passes skip_model_parallel_init=True, so the accessors run with check_initialized=False and return None. The projection is then silently built unsharded.

Same silent-None fallback as #7356, reached by a different route: there the top-level wrapper failed the vars() check, here the argument is simply never passed.

Impact

eh_proj is [2h, h]. At hidden_size=8192 that is ~134M parameters held as a full replica on every GTP peer rather than one shard, so at gtp_weight_remat_size=64 the layer costs 64 copies instead of one. e_proj and h_proj are [h, h] each and have the same problem on the MHC path.

Gradient accounting stays correct: the parameters are not tagged is_gtp_weight_remat, so param_is_not_gtp_duplicate counts them on GTP rank 0 only, which is the right answer for genuinely replicated tensors. The other visible consequence is that these projections' checkpoint shapes differ between a MIMO and a non-MIMO run at the same GTP size.

Fix

Pass pg_collection on both paths. All three implementations already accept it: ColumnParallelLinear, TEColumnParallelLinear and InferenceColumnParallelLinear.

The eh_proj line dates to 9a60a18ef7 ("ultra refit", 2026-03-19), before GTP under MIMO existed, and the MHC path inherited the same shape, so this looks like an oversight rather than a deliberate exclusion.

Tests

test_mtp_eh_proj_uses_custom_gtp_group in tests/unit_tests/generalized_tensor_parallel/test_gtp_custom_pgs.py, modelled on the neighbouring test_gdp_projections_use_custom_gtp_group. The MPU keeps one GTP pairing while the layer is handed the other, so reading the globals yields a valid-but-wrong group and the assertion on weight.group fails. The existing test_gtp_mtp.py cases cannot catch this, since they run with MPU globals initialized to the same topology.

The test covers the default eh_proj path only; mhc_enabled is off by default and its projections take the identical fix.

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter on my PR

black and isort clean. The test was not run locally; it needs 4 GPUs.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[nemo-automation-bot]

🔄 Merge queue validation started!

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