feat(training): wire generalized tensor parallelism

src/megatron/bridge/training/config.py

@@ -1400,6 +1400,17 @@ def validate(self) -> None:
         if hasattr(self.model, "finalize"):
             self.model.finalize()
 
+        from megatron.bridge.training.gtp import is_gtp_remat_active
+
+        if is_gtp_remat_active(self.model):
+            if self.dist.use_decentralized_pg:
+                raise ValueError(
+                    "GTP is not supported with dist.use_decentralized_pg=True. "
+                    "Set dist.use_decentralized_pg=False to use the standard MCore process-group runtime."
+                )
+            if self.ddp.average_in_collective:
+                raise ValueError("GTP requires ddp.average_in_collective=False.")
+
         self.logger.finalize()
         self.train.finalize()
         self.scheduler.finalize()

src/megatron/bridge/training/eval.py

@@ -34,6 +34,7 @@
 from megatron.bridge.training.callbacks import CallbackContext, CallbackManager, should_fire
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.forward_step_func_types import ForwardStepCallable
+from megatron.bridge.training.gtp import get_data_distribution_group
 from megatron.bridge.training.state import GlobalState
 from megatron.bridge.training.utils.mlflow_utils import _sanitize_mlflow_metrics
 from megatron.bridge.training.utils.pg_utils import get_pg_collection
@@ -141,7 +142,11 @@ def evaluate(
     eval_micro_batch_size = state.cfg.validation.eval_micro_batch_size
     # MegatronMIMO has heterogeneous per-module DP groups and intentionally owns
     # global-batch accounting through the container-level DP size.
-    eval_data_parallel_size = state.cfg.data_parallel_size if is_multimodule else pg_collection.dp.size()
+    eval_data_parallel_size = (
+        state.cfg.data_parallel_size
+        if is_multimodule
+        else get_data_distribution_group(pg_collection, state.cfg.model).size()
+    )
     eval_num_microbatches = eval_batch_size // (eval_micro_batch_size * eval_data_parallel_size)
 
     if is_multimodule and not isinstance(p2p_communicator, MultiModulePipelineCommunicator):
@@ -290,7 +295,9 @@ def evaluate(
                 if is_multimodule:
                     dp_cp_group = pg_collection.get_language_model_collection().dp_cp
                 else:
-                    dp_cp_group = pg_collection.dp_cp
+                    dp_cp_group = get_data_distribution_group(
+                        pg_collection, state.cfg.model, with_context_parallel=True
+                    )
 
                 for key in loss_dicts[0].keys():
                     if key not in total_loss_dict:

src/megatron/bridge/training/gtp.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generalized Tensor Parallelism helpers for the standard Bridge runtime."""
+
+from typing import Any
+
+import torch
+from megatron.core import parallel_state
+from megatron.core.process_groups_config import ProcessGroupCollection
+
+
+def get_transformer_config(model_config: Any) -> Any:
+    """Return the MCore transformer config nested in a Bridge model config."""
+    model_fields = getattr(type(model_config), "__dataclass_fields__", {})
+    if "transformer" in model_fields:
+        return model_config.transformer
+    return model_config
+
+
+def is_gtp_remat_active(model_config: Any) -> bool:
+    """Return whether dense or expert GTP weight rematerialization is enabled."""
+    transformer_config = get_transformer_config(model_config)
+    dense_size = getattr(transformer_config, "gtp_weight_remat_size", 1)
+    expert_size = getattr(transformer_config, "expert_gtp_weight_remat_size", 1)
+    return any(isinstance(size, int) and size > 1 for size in (dense_size, expert_size))
+
+
+def configure_gtp_remat(model_config: Any) -> None:
+    """Configure process-global GTP state before constructing model modules."""
+    if not is_gtp_remat_active(model_config):
+        return
+
+    transformer_config = get_transformer_config(model_config)
+    from megatron.core.tensor_parallel import gtp_api
+
+    if not gtp_api.HAVE_GTP:
+        raise RuntimeError("GTP requires TransformerEngine >= 2.19.")
+
+    gtp_api.configure_gtp_remat_from_recipe(
+        fp4=transformer_config.fp4 is not None,
+        fp8_recipe=transformer_config.fp8_recipe,
+        fp8=transformer_config.fp8 is not None,
+        calculate_per_token_loss=transformer_config.calculate_per_token_loss,
+    )
+
+
+def classify_gtp_remat_chains(model: list[torch.nn.Module], model_config: Any) -> None:
+    """Classify all model chunks after distributed wrapping and before first forward."""
+    if not is_gtp_remat_active(model_config):
+        return
+
+    transformer_config = get_transformer_config(model_config)
+    from megatron.core.tensor_parallel import gtp_api
+
+    gtp_api.classify_gtp_remat_chains(
+        model,
+        cuda_graph_modules=transformer_config.cuda_graph_modules,
+        moe_shared_expert_overlap=transformer_config.moe_shared_expert_overlap,
+        cuda_graph_impl=transformer_config.cuda_graph_impl,
+    )
+
+
+def get_data_distribution_group(
+    pg_collection: ProcessGroupCollection,
+    model_config: Any,
+    *,
+    with_context_parallel: bool = False,
+) -> torch.distributed.ProcessGroup:
+    """Return the group spanning every rank that consumes distinct input data."""
+    if not is_gtp_remat_active(model_config):
+        return pg_collection.dp_cp if with_context_parallel else pg_collection.dp
+    if with_context_parallel:
+        return pg_collection.dp_cp_gtp_remat
+    return parallel_state.get_data_parallel_group(with_gtp_remat=True)

src/megatron/bridge/training/initialize.py

@@ -50,6 +50,7 @@
 from megatron.bridge.models.hybrid.hybrid_builder import HybridModelConfig
 from megatron.bridge.models.transformer_config import TransformerConfig, _set_moe_expert_tensor_parallel_default
 from megatron.bridge.training.config import ConfigContainer, DistributedInitConfig, RerunStateMachineConfig, RNGConfig
+from megatron.bridge.training.gtp import is_gtp_remat_active
 from megatron.bridge.training.utils.pg_utils import DistTrainProcessGroupCollection
 from megatron.bridge.utils.common_utils import (
     get_local_rank_preinit,
@@ -759,6 +760,18 @@ def _initialize_distributed(
         if dist_config.use_decentralized_pg or dist_config.distributed_backend == "nccl":
             raise RuntimeError("Cannot initialize parallel groups with no CUDA devices available (device_count=0)")
 
+    if dist_config.use_decentralized_pg and is_gtp_remat_active(model_config):
+        raise NotImplementedError(
+            "GTP is not supported with dist.use_decentralized_pg=True. "
+            "Use the standard MCore process-group runtime by setting dist.use_decentralized_pg=False."
+        )
+
+    if is_gtp_remat_active(model_config):
+        from megatron.core.tensor_parallel.gtp_api import HAVE_GTP
+
+        if not HAVE_GTP:
+            raise RuntimeError("GTP requires TransformerEngine >= 2.19.")
+
     if dist_config.use_decentralized_pg:
         # Use HyperCommGrid to create local parallel groups passed through functions
         # instead of relying on mcore's global parallel state (mpu) variables.
@@ -812,6 +825,8 @@ def _initialize_distributed(
                 expert_model_parallel_size=model_config.expert_model_parallel_size,
                 num_distributed_optimizer_instances=num_distributed_optimizer_instances,
                 expert_tensor_parallel_size=model_config.expert_tensor_parallel_size,
+                gtp_remat_size=model_config.gtp_weight_remat_size,
+                expert_gtp_remat_size=model_config.expert_gtp_weight_remat_size,
                 distributed_timeout_minutes=dist_config.distributed_timeout_minutes,
                 nccl_communicator_config_path=dist_config.nccl_communicator_config_path,
                 order="tp-cp-ep-dp-pp" if not dist_config.use_tp_pp_dp_mapping else "tp-cp-ep-pp-dp",

src/megatron/bridge/training/setup.py

@@ -55,6 +55,11 @@
 )
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.fsdp_compat import MEGATRON_FSDP_TYPES
+from megatron.bridge.training.gtp import (
+    classify_gtp_remat_chains,
+    configure_gtp_remat,
+    get_data_distribution_group,
+)
 from megatron.bridge.training.initialize import initialize_megatron, set_jit_fusion_options
 from megatron.bridge.training.optim import (
     memory_efficient_fp32_optimizer_state_loading,
@@ -500,7 +505,7 @@ def modelopt_pre_wrap_hook(model):
         train_state=state.train_state,
         model_length=len(model),
         train_valid_test_datasets_provider=train_valid_test_datasets_provider,
-        dp_group=pg_collection.dp,
+        dp_group=get_data_distribution_group(pg_collection, cfg.model),
         eval_dp_group=state._eval_pgs.dp if state._eval_pgs is not None else None,
     )
     timers("train/valid/test-data-iterators-setup").stop()
@@ -584,10 +589,11 @@ def _register_setup_pre_wrap_hook(
 def _build_distributed_model(cfg: ConfigContainer, pg_collection: ProcessGroupCollection) -> list[MegatronModule]:
     """Build distributed model from either ModelConfig or ModelProviderMixin."""
     model_config = cfg.model
+    configure_gtp_remat(model_config)
     if isinstance(model_config, ModelConfig):
         builder_cls = model_config.get_builder_cls()
         builder = builder_cls(model_config)
-        return builder.build_distributed_models(
+        model = builder.build_distributed_models(
             pg_collection=pg_collection,
             ddp_config=cfg.ddp,
             overlap_param_gather_with_optimizer_step=cfg.optimizer.overlap_param_gather_with_optimizer_step,
@@ -596,14 +602,16 @@ def _build_distributed_model(cfg: ConfigContainer, pg_collection: ProcessGroupCo
             data_parallel_random_init=cfg.rng.data_parallel_random_init,
         )
     else:
-        return model_config.provide_distributed_model(
+        model = model_config.provide_distributed_model(
             ddp_config=cfg.ddp,
             use_megatron_fsdp=cfg.dist.use_megatron_fsdp,
             use_torch_fsdp2=cfg.dist.use_torch_fsdp2,
             overlap_param_gather_with_optimizer_step=cfg.optimizer.overlap_param_gather_with_optimizer_step,
             data_parallel_random_init=cfg.rng.data_parallel_random_init,
             pg_collection=pg_collection,
         )
+    classify_gtp_remat_chains(model, model_config)
+    return model
 
 
 def _update_model_config_funcs(

src/megatron/bridge/training/train.py

@@ -74,6 +74,7 @@
 from megatron.bridge.training.eval import evaluate_and_print_results
 from megatron.bridge.training.forward_step_func_types import ForwardStepCallable
 from megatron.bridge.training.fsdp_compat import MEGATRON_FSDP_TYPES
+from megatron.bridge.training.gtp import get_data_distribution_group
 from megatron.bridge.training.initialize import destroy_global_state
 from megatron.bridge.training.nvrx_straggler import (
     check_nvrx_straggler_detection,
@@ -336,7 +337,8 @@ def train(
     start_iteration = global_state.train_state.step
     print_rank_0(f"Starting training loop at iteration {start_iteration}")
     p2p_communicator = P2PCommunicator(pp_group=pg_collection.pp, config=model_config)
-    dp_size = pg_collection.dp.size()
+    data_distribution_group = get_data_distribution_group(pg_collection, config.model)
+    dp_size = data_distribution_group.size()
     # Anchor for interval-average throughput logging: training_log reports the FLOPS
     # performed over each logging interval as the delta of
     # floating_point_operations_so_far. Seed it with the current cumulative (0 fresh,
@@ -588,7 +590,7 @@ def train(
             global_state,
             data_parallel_size=dp_size,
             vp_size=config.model.virtual_pipeline_model_parallel_size,
-            dp_group=pg_collection.dp,
+            dp_group=data_distribution_group,
             include_vision_patch_stats=True,
             include_cross_attention_stats=hasattr(
                 config.model, "_get_num_floating_point_operations_with_runtime_stats"
@@ -993,7 +995,7 @@ def train_step(
                 # there is one dict per microbatch. in new reporting, we average
                 # over the total number of tokens across the global batch.
                 val = torch.vstack(val).sum(dim=0)
-                dp_cp_group = pg_collection.dp_cp
+                dp_cp_group = get_data_distribution_group(pg_collection, cfg.model, with_context_parallel=True)
                 torch.distributed.all_reduce(val, group=dp_cp_group)
                 loss_reduced[key] = val[0] / val[1]
             elif val[0].numel() == 1:
@@ -1576,7 +1578,7 @@ def _should_skip_and_handle_iteration(
 
     # Update step and sample counters
     global_state.train_state.step += 1
-    dp_size = pg_collection.dp.size()
+    dp_size = get_data_distribution_group(pg_collection, cfg.model).size()
     batch_size = dp_size * cfg.train.micro_batch_size * get_num_microbatches()
     global_state.train_state.consumed_train_samples += batch_size
     global_state.train_state.skipped_train_samples += batch_size