[fix] Cleanup usage of pg_collection in gtp

docs/api-guide/core/generalized_tensor_parallel.md

@@ -557,5 +557,6 @@ torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parall
 | `test_gtp_dcp.py` | DCP sharding metadata (§3.3): TP×GTP_remat offsets, pad reshard, `replica_id`, native-FP8 save/load. |
 | `test_gtp_muon_dcp.py` | Muon optimizer-state DCP roundtrip (§1.6): `replica_id` fold + native-FP8 backfill matching. |
 | `test_gtp_fp8_param_gather.py` | Native-FP8 GTP_remat (§1.3): fp8-vs-BF16 loss parity (TP1/TP2, MoE), post-save-spike guard. |
+| `test_gtp_custom_pgs.py` | `pg_collection` plumbing: a custom `gtp_remat` group (permuted ranks, same size) must give the same fwd/bwd results as the MPU groups — catches modules reading `parallel_state` instead of the collection passed to them. |
 
 All tests require ≥ 4 GPUs and TransformerEngine >= 2.19; they self-skip when those are unavailable. A green run (skips for unmet hardware/config are acceptable) is the minimum bar for any GTP_remat change.

megatron/core/distributed/finalize_model_grads.py

@@ -492,7 +492,10 @@ def _allreduce_non_tensor_model_parallel_grads(
 
 
 def _allreduce_replicated_grads_over_gtp_remat_group(
-    model: List[torch.nn.Module], calculate_per_token_loss: bool = False
+    model: List[torch.nn.Module],
+    gtp_remat_group: Optional[torch.distributed.ProcessGroup],
+    egtp_remat_group: Optional[torch.distributed.ProcessGroup],
+    calculate_per_token_loss: bool = False,
 ):
     """Complete the gtp_remat / egtp_remat axis reduction for replicated parameters.
 
@@ -510,12 +513,6 @@ def _allreduce_replicated_grads_over_gtp_remat_group(
 
     No-op when GTP_remat is inactive (group size <= 1).
     """
-    pg_collection = ProcessGroupCollection.use_mpu_process_groups(
-        required_pgs=["gtp_remat", "expt_gtp_remat"]
-    )
-    gtp_remat_group = pg_collection.gtp_remat
-    egtp_remat_group = pg_collection.expt_gtp_remat
-
     dense_active = gtp_remat_group is not None and gtp_remat_group.size() > 1
     expert_active = egtp_remat_group is not None and egtp_remat_group.size() > 1
     if not dense_active and not expert_active:
@@ -604,12 +601,29 @@ def finalize_model_grads(
         # Full DP x CP x gtp_remat group: num_tokens (the per-token-loss divisor below) counts the
         # gtp_remat peers' distinct tokens. Falls back to replicate dp_cp when gtp is inactive.
         dp_cp_group = getattr(pg_collection, 'dp_cp_gtp_remat', None) or pg_collection.dp_cp
+        gtp_remat_group = getattr(pg_collection, 'gtp_remat', None)
+        egtp_remat_group = getattr(pg_collection, 'expt_gtp_remat', None)
     else:
         tp_group = parallel_state.get_tensor_model_parallel_group()
         pp_group = parallel_state.get_pipeline_model_parallel_group()
         embd_group = parallel_state.get_embedding_group(check_initialized=False)
         pos_emb_group = parallel_state.get_position_embedding_group(check_initialized=False)
         dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=True)
+        gtp_remat_group = parallel_state.get_gtp_weight_remat_group(check_initialized=False)
+        egtp_remat_group = parallel_state.get_expert_gtp_weight_remat_group(check_initialized=False)
+
+    # A missing group would silently skip the gtp_remat-axis reduction below and train on
+    # wrong gradients, so fail loudly whenever the config says the axis is active.
+    for axis, group, axis_size in (
+        ('gtp_remat', gtp_remat_group, config.gtp_weight_remat_size),
+        ('expt_gtp_remat', egtp_remat_group, config.expert_gtp_weight_remat_size),
+    ):
+        if axis_size > 1:
+            found = 'None' if group is None else f'a size-{group.size()} group'
+            assert group is not None and group.size() == axis_size, (
+                f"{axis} is enabled (size={axis_size}) but pg_collection provides {found}. "
+                f"Pass a pg_collection carrying `{axis}` to finalize_model_grads."
+            )
 
     # Fence the current stream against all GTP backward grad work before the DP gradient sync.
     if config.gtp_weight_remat_size > 1 or config.expert_gtp_weight_remat_size > 1:
@@ -646,7 +660,10 @@ def finalize_model_grads(
         )
     _allreduce_non_tensor_model_parallel_grads(model, config, tp_group)
     _allreduce_replicated_grads_over_gtp_remat_group(
-        model, calculate_per_token_loss=config.calculate_per_token_loss
+        model,
+        gtp_remat_group,
+        egtp_remat_group,
+        calculate_per_token_loss=config.calculate_per_token_loss,
     )
     if config.timers is not None:
         config.timers('non-tensor-parallel-grads-all-reduce').stop()

megatron/core/extensions/transformer_engine.py

@@ -33,7 +33,7 @@
     get_tensor_model_parallel_world_size,
     model_parallel_is_initialized,
 )
-from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.process_groups_config import ProcessGroupCollection, resolve_gtp_remat_group
 from megatron.core.quantization.quant_config import QuantizationConfig
 from megatron.core.quantization.utils import get_quant_config_or_none
 from megatron.core.tensor_parallel.layers import (
@@ -1362,10 +1362,13 @@ def __init__(
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
         stride: int = 1,
         name: str | None = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         """
         Args:
             name (str | None): module instance name passed top-down from its paranet module
+            pg_collection (ProcessGroupCollection | None): process groups used by this layer.
+                Falls back to the MPU global process groups when not given.
         """
         if not HAVE_TE:
             raise ImportError(
@@ -1457,10 +1460,7 @@ def __init__(
             ), "Must have at least TE version 2.3 or higher to use symmetric memory all reduce"
             extra_kwargs["symmetric_ar_type"] = self.config.symmetric_ar_type
 
-        pg_collection = ProcessGroupCollection.use_mpu_process_groups(
-            required_pgs=["gtp_remat", "expt_gtp_remat"]
-        )
-        gtp_remat_group = pg_collection.expt_gtp_remat if is_expert else pg_collection.gtp_remat
+        gtp_remat_group = resolve_gtp_remat_group(pg_collection, is_expert)
         self.stride = stride
 
         self.te_quant_params: Optional[TEQuantizationParams] = None
@@ -1621,10 +1621,13 @@ def __init__(
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
         stride: int = 1,
         name: str | None = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         """
         Args:
             name (str | None): module instance name passed top-down from its paranet module
+            pg_collection (ProcessGroupCollection | None): process groups used by this layer.
+                Falls back to the MPU global process groups when not given.
         """
         if not HAVE_TE:
             raise ImportError(
@@ -1639,10 +1642,7 @@ def __init__(
         world_size = get_pg_size(tp_group)
         rank = get_pg_rank(tp_group)
         self.stride = stride
-        pg_collection = ProcessGroupCollection.use_mpu_process_groups(
-            required_pgs=["gtp_remat", "expt_gtp_remat"]
-        )
-        gtp_remat_group = pg_collection.expt_gtp_remat if is_expert else pg_collection.gtp_remat
+        gtp_remat_group = resolve_gtp_remat_group(pg_collection, is_expert)
 
         super().__init__(
             input_size=input_size,
@@ -1882,10 +1882,13 @@ def __init__(
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
         name: str | None = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         """
         Args:
             name (str | None): module instance name passed top-down from its paranet module
+            pg_collection (ProcessGroupCollection | None): process groups used by this layer.
+                Falls back to the MPU global process groups when not given.
         """
         if not HAVE_TE:
             raise ImportError(
@@ -1899,10 +1902,7 @@ def __init__(
             )
         tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
         self._tp_group = tp_group
-        pg_collection = ProcessGroupCollection.use_mpu_process_groups(
-            required_pgs=["gtp_remat", "expt_gtp_remat"]
-        )
-        gtp_remat_group = pg_collection.expt_gtp_remat if is_expert else pg_collection.gtp_remat
+        gtp_remat_group = resolve_gtp_remat_group(pg_collection, is_expert)
 
         super().__init__(
             input_size=input_size,

megatron/core/models/common/embeddings/language_model_embedding.py

@@ -6,6 +6,7 @@
 from torch import Tensor
 
 from megatron.core import tensor_parallel
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import get_tensor_model_parallel_group_if_none, nvtx_decorator
@@ -24,6 +25,7 @@ class LanguageModelEmbedding(MegatronModule):
         num_tokentypes (int): Set to 0 without binary head, and 2 with a binary head. Defaults to 0.
         scatter_to_sequence_parallel (bool): Set to False to disable scatter of embedding
             across sequence parallel region. Defaults to True.
+        pg_collection (ProcessGroupCollection, optional): Process groups used by the embedding.
     """
 
     def __init__(
@@ -35,6 +37,7 @@ def __init__(
         num_tokentypes: int = 0,
         scatter_to_sequence_parallel: bool = True,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         super().__init__(config=config)
 
@@ -60,6 +63,7 @@ def __init__(
             reduce_scatter_embeddings=self.reduce_scatter_embeddings,
             config=self.config,
             tp_group=self.tp_group,
+            pg_collection=pg_collection,
         )
 
         # Position embedding (serial).

megatron/core/models/hybrid/hybrid_model.py

@@ -239,6 +239,7 @@ def __init__(
                 position_embedding_type=position_embedding_type,
                 scatter_to_sequence_parallel=scatter_embedding_sequence_parallel,
                 tp_group=self.pg_collection.tp,
+                pg_collection=self.pg_collection,
             )
 
         # MLA (also used by DeepSeek Sparse Attention) uses its own decoupled RoPE, therefore we do
@@ -322,6 +323,7 @@ def __init__(
                 skip_weight_param_allocation=self.pre_process
                 and self.share_embeddings_and_output_weights,
                 tp_group=self.pg_collection.tp,
+                pg_collection=self.pg_collection,
             )
 
         if self.pre_process or self.post_process or self.mtp_process:

megatron/core/pipeline_parallel/schedules.py

@@ -690,6 +690,15 @@ def _build_default_pg_collection() -> ProcessGroupCollection:
     pg_collection.dp = parallel_state.get_data_parallel_group(
         with_context_parallel=False, partial_data_parallel=False
     )
+    # gtp_remat axis: consumers read these with getattr and silently skip the gtp_remat
+    # reduction when absent, so populate them even when GTP_remat is inactive.
+    pg_collection.gtp_remat = parallel_state.get_gtp_weight_remat_group(check_initialized=False)
+    pg_collection.expt_gtp_remat = parallel_state.get_expert_gtp_weight_remat_group(
+        check_initialized=False
+    )
+    pg_collection.dp_cp_gtp_remat = parallel_state.get_data_parallel_group(
+        with_context_parallel=True, partial_data_parallel=False
+    )
     return pg_collection
 
 
@@ -1614,7 +1623,7 @@ def forward_backward_helper_wrapper(
                 recv_next = True
                 if is_pp_last_stage(p2p_communicator.pp_group):
                     recv_next = False
-                (input_tensor, output_tensor_grad) = (
+                input_tensor, output_tensor_grad = (
                     p2p_communicator.send_forward_backward_recv_forward_backward(
                         output_tensor,
                         input_tensor_grad,
@@ -1678,7 +1687,7 @@ def forward_backward_helper_wrapper(
                 if is_pp_last_stage(p2p_communicator.pp_group):
                     recv_next = False
 
-                (bwd_recv_buffer[-1], bwd_wait_handles) = (
+                bwd_recv_buffer[-1], bwd_wait_handles = (
                     p2p_communicator.send_backward_recv_backward(
                         input_tensor_grad,
                         recv_next=recv_next,
@@ -1831,7 +1840,7 @@ def pp_post_backward(input_tensor_grad, vp_stage=None):
                     backward_k, forward=False
                 )
 
-                (bwd_recv_buffer[backward_k % bwd_recv_buffer_size], bwd_wait_handles) = (
+                bwd_recv_buffer[backward_k % bwd_recv_buffer_size], bwd_wait_handles = (
                     p2p_communicator.send_backward_recv_backward(
                         input_tensor_grad,
                         recv_next=recv_next,
@@ -1904,7 +1913,7 @@ def pp_post_backward(input_tensor_grad, vp_stage=None):
                 recv_prev = False
 
             # Communicate tensors.
-            (input_tensor, output_tensor_grad) = (
+            input_tensor, output_tensor_grad = (
                 p2p_communicator.send_forward_backward_recv_forward_backward(
                     output_tensor,
                     input_tensor_grad,

megatron/core/process_groups_config.py

@@ -689,6 +689,31 @@ def setup_process_groups_for_ddp(
             return result
 
 
+def resolve_gtp_remat_group(
+    pg_collection: Optional["ProcessGroupCollection"], is_expert: bool
+) -> Optional[torch.distributed.ProcessGroup]:
+    """Resolve the gtp_remat / expt_gtp_remat group for a weight-owning module.
+
+    Prefers the group carried by ``pg_collection``; falls back to the MPU globals when the
+    caller passed no collection, or one predating the gtp_remat fields. The fallback keeps
+    pre-pg_collection callers working — a collection that does carry the field is always
+    honored, including when it holds a custom (non-MPU) group.
+
+    Args:
+        pg_collection: Collection supplied by the caller, or None.
+        is_expert: Select the expert axis (``expt_gtp_remat``) instead of the dense one.
+    """
+    attr = 'expt_gtp_remat' if is_expert else 'gtp_remat'
+    # `vars()`, not hasattr: __getattr__ makes hasattr always True, so the fallback below
+    # would be unreachable.
+    if pg_collection is not None and attr in vars(pg_collection):
+        return getattr(pg_collection, attr)
+    mpu_pgs = ProcessGroupCollection.use_mpu_process_groups(
+        required_pgs=['gtp_remat', 'expt_gtp_remat']
+    )
+    return getattr(mpu_pgs, attr)
+
+
 @dataclass
 class MultiModuleProcessGroupCollection:
     """Process group collection for multi-module pipelines.

megatron/core/ssm/mamba_mixer.py

@@ -291,6 +291,7 @@ def __init__(
             is_expert=False,
             tp_comm_buffer_name="fc1",
             tp_group=self.pg_collection.tp,
+            pg_collection=self.pg_collection,
             name=(name + f".in_proj") if name is not None else None,
         )
         # in_proj packs [z, x, B, C, dt] into one ColumnParallelLinear.  Each
@@ -442,6 +443,7 @@ def __init__(
             is_expert=False,
             tp_comm_buffer_name="fc2",
             tp_group=self.pg_collection.tp,
+            pg_collection=self.pg_collection,
             name=(name + f".out_proj") if name is not None else None,
         )
 

megatron/core/tensor_parallel/inference_layers.py

@@ -20,6 +20,7 @@
 from megatron.core.inference.quantization.utils import mm_mxfp8
 from megatron.core.inference.symmetric_memory import SymmetricMemoryManager
 from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.process_groups_config import ProcessGroupCollection, resolve_gtp_remat_group
 from megatron.core.tensor_parallel.mappings import (
     gather_from_tensor_model_parallel_region,
     reduce_scatter_to_sequence_parallel_region,
@@ -91,6 +92,7 @@ def __init__(
         symmetric_ar_type: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
         name: str | None = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine"
         super().__init__(
@@ -107,6 +109,8 @@ def __init__(
             symmetric_ar_type=symmetric_ar_type,
             tp_group=tp_group,
             name=name,
+            # TELinear takes the resolved group rather than the collection.
+            gtp_remat_group=resolve_gtp_remat_group(pg_collection, is_expert),
         )
 
     def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, None]:
@@ -139,6 +143,7 @@ def __init__(
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
         name: str | None = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine"
         super().__init__(
@@ -155,6 +160,7 @@ def __init__(
             tp_comm_buffer_name=tp_comm_buffer_name,
             tp_group=tp_group,
             name=name,
+            pg_collection=pg_collection,
         )
         self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
         self.tp_size = dist.get_world_size(self.tp_group)
@@ -268,6 +274,7 @@ def __init__(
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
         name: str | None = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine"
         super().__init__(
@@ -284,6 +291,7 @@ def __init__(
             tp_comm_buffer_name=tp_comm_buffer_name,
             tp_group=tp_group,
             name=name,
+            pg_collection=pg_collection,
         )
         self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
         self.tp_size = dist.get_world_size(self.tp_group)
@@ -366,6 +374,7 @@ def __init__(
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
         name: str | None = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine"
         super().__init__(
@@ -380,6 +389,7 @@ def __init__(
             tp_comm_buffer_name=tp_comm_buffer_name,
             tp_group=tp_group,
             name=name,
+            pg_collection=pg_collection,
         )
         self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
         self.tp_size = dist.get_world_size(self.tp_group)