[yashaswikarnati]

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

What does this PR do?

Route explicit model-parallel process groups through GDP and latent-MoE distributed-checkpoint metadata without changing duplicated TELinear execution semantics.

MIMO models can own process groups that differ from the globally initialized MPU topology. This change:

• uses the GDP mixer's explicit TP and DP+CP groups for RMSNorm, base parameters, Conv1d, and child-module checkpoint wrappers;
• passes the owning TP group to latent-MoE projection linears;
• retains that group on duplicated TELinear only for checkpoint metadata, while Transformer Engine still receives tp_group=None, tp_size=1, and parallel_mode=None;
• preserves the current GTP in_proj checkpoint split and resharding path.

This is intentionally separate from LayerWise optimizer changes, MTP, optimizer policy, recipes, and padding-only checkpoint work.

Issue tracking

Linked issue: N/A

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

Validation:

• CHECK_ONLY=true uv run --no-sync bash tools/autoformat.sh (Black, isort, Pylint 10.00/10, and Ruff passed; mypy was unavailable and is non-blocking in the script)
• python3 -m compileall -q on all six changed Python files
• CW-DFW environment job 16387980: 1 H100; locked test environment preflight passed with PyTorch 2.13 nv26.06, Transformer Engine 2.16, pytest 9.1.1, and pytest-mock
• CW-DFW job 16388633: 1 node, 8 H100s, 8 ranks, current 26.06-py3 development image; all ranks passed the 6 focused GDP custom-group/TP-reshard, latent-MoE construction, and duplicated-TELinear execution/checkpoint tests

[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.

[Phlip79]

Why do we not just use dp_cp_group directly?

[yashaswikarnati]

This is just following the existing convention as other layers like Mamba, moe, TE linear layers. I think the above helper function is added ensure_metadata_has_dp_cp_group , so if hte group does not exist, it falls back to paralel state for backward compatibility

[fanshiqing]

the bottom line: we don't require tp_group for duplicated mode,
pls try to figure out a better fix if possible

[fanshiqing]

/claude review

[claude]

Light review — two points, both minor. The core change looks right: routing pg_collection.tp / metadata["dp_cp_group"] into the GDP wrappers, and the new replica_id / tp_group=None split in make_sharded_tensors_for_checkpoint_with_gtp_remat correctly makes a duplicated GTP weight describe GTP-only fragmentation while still electing exactly one TP replica as writer. Test coverage for the changed paths is good.
1. InferenceLinear is skipped by the is TELinear identity checks
megatron/core/transformer/moe/moe_layer.py:280 and :308 gate on linear_cls is TELinear. InferenceLinear subclasses TELinear and inherits TELinear.sharded_state_dict verbatim, so under transformer_impl="inference_optimized" neither gtp_replica_group nor _tp_group = pg_collection.tp is applied. TELinear.__init__ then falls back to get_tensor_model_parallel_group_if_none(None) (transformer_engine.py:1319), i.e. the MPU global TP group lands in the checkpoint replica coordinate — the exact failure this PR fixes for the non-inference path.
issubclass(linear_cls, TELinear) would cover both, but note that InferenceLinear.__init__ (megatron/core/tensor_parallel/inference_layers.py:83-100) accepts neither gtp_remat_group nor gtp_replica_group, and linear_gtp_kwargs already passes gtp_remat_group unconditionally — so this call site raises TypeError today regardless (pre-existing on main, not introduced here). Either way the is TELinear gate on gtp_replica_group is effectively dead code. Worth either forwarding both kwargs through InferenceLinear or leaving a comment that the inference-optimized latent-proj path is unsupported, so the gap is not silently inherited.
2. Stale docstring in make_sharded_tensors_for_checkpoint_with_gtp_remat
megatron/core/tensor_parallel/generalized_tensor_parallelism.py:2709-2711 still says: "GTP-sharded tensors elect their writer inside make_tp_sharded_tensor_for_checkpoint, always off their own stamped group."
After this change that only holds for tensors in the axis map. The new tp_axis is None branch passes an explicit replica_id=(0, tp_rank, dp_replica_rank), so the callee computes dp_replica_id = gtp_replica_rank(tensor) and discards it; the rank actually used comes from the callers gtp_replica_rank(first_gtp_param, intra_dp_cp_group) — which can be the intra_dp_cp_group override. Suggested amendment:

    ``intra_dp_cp_group`` optionally overrides the replica group behind the replica_ids of the
    GTP-REPLICATED entries here (bias / extra_state); it must EXCLUDE the gtp_remat axis. Left
    None, the group is taken from the params themselves. TP-sharded GTP tensors elect their
    writer inside :func:`make_tp_sharded_tensor_for_checkpoint`, off their own stamped group;
    GTP-sharded tensors absent from the axis map are TP-duplicated, so their writer is elected
    here instead (explicit ``replica_id`` folding in ``tp_rank``). See :func:`gtp_replica_rank`.

[fanshiqing]

Thanks for fixing this!

That's real bug after supporting --gtp-remat-opt-in-modules moe_latent_proj 2 weeks ago where the latent weights are DUPLICATED across TP group while are SHARDED across GTP groups!

[yashaswikarnati]

/ok to test 1633cb5

[nemo-automation-bot]

🔄 Merge queue validation started!

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