diff --git a/examples/mimo/model_providers/nemotron_moe_vlm.py b/examples/mimo/model_providers/nemotron_moe_vlm.py
index 133bf9bc1d2..03013e73e85 100644
--- a/examples/mimo/model_providers/nemotron_moe_vlm.py
+++ b/examples/mimo/model_providers/nemotron_moe_vlm.py
@@ -12,6 +12,7 @@
 from examples.mimo.model_providers.radio_encoder import (
     RADIO_ENCODER_MODULE_NAME,
     _base_config,
+    _disable_gtp,
     _make_dense_non_hybrid,
     add_radio_encoder_args,
     radio_vision_config,
@@ -21,6 +22,7 @@
 from megatron.core.activations import squared_relu
 from megatron.core.hyper_comm_grid import HyperCommGrid
 from megatron.core.hyper_comm_grid import _is_process_group_member as is_process_group_member
+from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards
 from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec
 from megatron.core.models.mamba.mamba_model import MambaModel
 from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
@@ -105,6 +107,23 @@ def nemotron_language_config(
     config.expert_tensor_parallel_size = expt_tp_size
     config.tensor_model_parallel_size = tp_size
     config.pipeline_model_parallel_size = pp_size
+    config.tensor_parallel_num_weight_shards, config.gtp_weight_remat_size = (
+        resolve_tensor_parallel_weight_shards(
+            tp_size,
+            getattr(args, "tensor_parallel_num_weight_shards", None),
+            getattr(args, "gtp_weight_remat_size", 1),
+        )
+    )
+    (
+        config.expert_tensor_parallel_num_weight_shards,
+        config.expert_gtp_weight_remat_size,
+    ) = resolve_tensor_parallel_weight_shards(
+        expt_tp_size,
+        getattr(args, "expert_tensor_parallel_num_weight_shards", None),
+        getattr(args, "expert_gtp_weight_remat_size", 1),
+        shards_field="expert_tensor_parallel_num_weight_shards",
+        tp_field="expert_tensor_parallel_size",
+    )
     config.sequence_parallel = tp_size > 1
     config.position_embedding_type = "none"
     return config
@@ -142,6 +161,7 @@ def nemotron_projection_config(
     config.normalization = "RMSNorm"
     _make_dense_non_hybrid(config)  # Projection inherits no MoE/Mamba/hybrid settings.
     config.tensor_model_parallel_size = tp_size
+    _disable_gtp(config)
     config.sequence_parallel = False
     return config
 
diff --git a/examples/mimo/model_providers/radio_encoder.py b/examples/mimo/model_providers/radio_encoder.py
index 9e0591cc7e7..c301c0c29ca 100644
--- a/examples/mimo/model_providers/radio_encoder.py
+++ b/examples/mimo/model_providers/radio_encoder.py
@@ -83,6 +83,15 @@ def _make_dense_non_hybrid(config: TransformerConfig) -> None:
     config.use_fused_weighted_squared_relu = False
 
 
+def _disable_gtp(config: TransformerConfig) -> None:
+    """Keep encoder-side weights replicated across the language module's GTP axes."""
+    config.tensor_parallel_num_weight_shards = config.tensor_model_parallel_size
+    config.gtp_weight_remat_size = 1
+    expert_tp = config.expert_tensor_parallel_size or config.tensor_model_parallel_size
+    config.expert_tensor_parallel_num_weight_shards = expert_tp
+    config.expert_gtp_weight_remat_size = 1
+
+
 def radio_vision_config(args: argparse.Namespace, tp_size: int, pp_size: int) -> TransformerConfig:
     """RADIO vision config: stock from-args base + RADIO-specific overrides."""
     config = deepcopy(_base_config(args))
@@ -114,6 +123,7 @@ def radio_vision_config(args: argparse.Namespace, tp_size: int, pp_size: int) ->
     config.bf16 = bf16
     config.tensor_model_parallel_size = tp_size
     config.pipeline_model_parallel_size = pp_size
+    _disable_gtp(config)
     config.sequence_parallel = False
     return config
 
diff --git a/examples/mimo/pretrain_mimo.py b/examples/mimo/pretrain_mimo.py
index ee521b188a5..06219a286fb 100644
--- a/examples/mimo/pretrain_mimo.py
+++ b/examples/mimo/pretrain_mimo.py
@@ -48,14 +48,19 @@ def _parse_and_validate() -> argparse.Namespace:
     args = parse_args(extra_args_provider)
     validate_hetero_grid_args(args, args.world_size)
     physical_world_size = args.world_size
-    # Stock validate_args sets data_parallel_size = world_size // (tp*pp*cp); feed the
-    # language module's world (llm_dp; stock tp/pp/cp stay 1, MIMO parallelism is in --llm-*)
-    # so it yields llm_dp. The physical world incl. encoder ranks is restored below.
+    # Stock validation owns the derived training arguments. Give it the language module's
+    # parallel degrees and logical world so its DP/GTP accounting matches the explicit MIMO grid.
+    args.tensor_model_parallel_size = args.llm_tp
+    args.pipeline_model_parallel_size = args.llm_pp
+    args.context_parallel_size = args.llm_cp
+    args.expert_model_parallel_size = args.llm_ep
+    args.expert_tensor_parallel_size = args.llm_expt_tp or 1
     args.world_size = (
         args.llm_dp
-        * args.tensor_model_parallel_size
-        * args.pipeline_model_parallel_size
-        * args.context_parallel_size
+        * args.llm_tp
+        * args.gtp_weight_remat_size
+        * args.llm_pp
+        * args.llm_cp
     )
     try:
         validate_args(args, {"dataloader_type": "external"})
diff --git a/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh b/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh
index 787df50fef8..92bdce51a25 100755
--- a/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh
+++ b/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh
@@ -1,26 +1,78 @@
 #!/bin/bash
 
-# Run an eight-rank heterogeneous mock training loop with Nemotron6-MoE VLM 20L.
+# Run heterogeneous mock training with the Nemotron6-MoE VLM 20L recipe.
 
 set -euo pipefail
 
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+
 TRAIN_ITERS=${TRAIN_ITERS:-20}
 NUM_MICROBATCHES=${NUM_MICROBATCHES:-4}
 EVAL_INTERVAL=${EVAL_INTERVAL:-1}
 EVAL_ITERS=${EVAL_ITERS:-0}
-MICRO_BATCH_SIZE=1
-LLM_DP=2
-GLOBAL_BATCH_SIZE=$((MICRO_BATCH_SIZE * NUM_MICROBATCHES * LLM_DP))
+MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-1}
+NNODES=${NNODES:-1}
+NPROC_PER_NODE=${NPROC_PER_NODE:-8}
+ENCODER_TP=${ENCODER_TP:-2}
+ENCODER_DP=${ENCODER_DP:-2}
+LLM_OFFSET=${LLM_OFFSET:-$((ENCODER_TP * ENCODER_DP))}
+LLM_TP=${LLM_TP:-2}
+LLM_CP=${LLM_CP:-1}
+LLM_PP=${LLM_PP:-1}
+LLM_DP=${LLM_DP:-2}
+LLM_EP=${LLM_EP:-4}
+LLM_EXPT_TP=${LLM_EXPT_TP:-1}
+TENSOR_PARALLEL_NUM_WEIGHT_SHARDS=${TENSOR_PARALLEL_NUM_WEIGHT_SHARDS:-${LLM_TP}}
+EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS=${EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS:-${LLM_EXPT_TP}}
+
+if ((TENSOR_PARALLEL_NUM_WEIGHT_SHARDS % LLM_TP != 0)); then
+  echo "TENSOR_PARALLEL_NUM_WEIGHT_SHARDS must be divisible by LLM_TP" >&2
+  exit 2
+fi
+if ((EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS % LLM_EXPT_TP != 0)); then
+  echo "EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS must be divisible by LLM_EXPT_TP" >&2
+  exit 2
+fi
+
+GTP=$((TENSOR_PARALLEL_NUM_WEIGHT_SHARDS / LLM_TP))
+LLM_SIZE=$((LLM_TP * GTP * LLM_CP * LLM_PP * LLM_DP))
+EXPECTED_WORLD_SIZE=$((ENCODER_TP * ENCODER_DP + LLM_SIZE))
+WORLD_SIZE=$((NNODES * NPROC_PER_NODE))
+if ((WORLD_SIZE != EXPECTED_WORLD_SIZE)); then
+  echo "NNODES*NPROC_PER_NODE=${WORLD_SIZE}, but encoder+LLM grids require ${EXPECTED_WORLD_SIZE}" >&2
+  exit 2
+fi
+
+GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-$((MICRO_BATCH_SIZE * NUM_MICROBATCHES * LLM_DP * GTP))}
 TORCHRUN_LOG_DIR=${TORCHRUN_LOG_DIR:-"${PWD}/logs/torchrun-$(date +%Y%m%d_%H%M%S)-$$"}
 mkdir -p "${TORCHRUN_LOG_DIR}"
 
 TORCHRUN_ARGS=(
-  --standalone
-  --nproc-per-node 8
+  --nproc-per-node "${NPROC_PER_NODE}"
   --log-dir "${TORCHRUN_LOG_DIR}"
   --redirects 3
-  --tee 3
+  --tee 0:3
 )
+if ((NNODES == 1)); then
+  TORCHRUN_ARGS+=(--standalone)
+else
+  NODE_RANK=${NODE_RANK:-${SLURM_NODEID:-0}}
+  if [[ -z "${MASTER_ADDR:-}" ]]; then
+    if [[ -z "${SLURM_JOB_NODELIST:-}" ]]; then
+      echo "MASTER_ADDR or SLURM_JOB_NODELIST is required for a multi-node run" >&2
+      exit 2
+    fi
+    mapfile -t slurm_nodes < <(scontrol show hostnames "${SLURM_JOB_NODELIST}")
+    MASTER_ADDR=${slurm_nodes[0]}
+  fi
+  MASTER_PORT=${MASTER_PORT:-$((10000 + ${SLURM_JOB_ID:-0} % 50000))}
+  TORCHRUN_ARGS+=(
+    --nnodes "${NNODES}"
+    --node-rank "${NODE_RANK}"
+    --master-addr "${MASTER_ADDR}"
+    --master-port "${MASTER_PORT}"
+  )
+fi
 
 uv run --extra ssm python -m torch.distributed.run \
   "${TORCHRUN_ARGS[@]}" \
@@ -71,15 +123,17 @@ uv run --extra ssm python -m torch.distributed.run \
   --seq-length 8192 \
   --max-position-embeddings 8192 \
   --bf16 \
-  --encoder-tp 2 \
-  --encoder-dp 2 \
-  --llm-offset 4 \
-  --llm-tp 2 \
-  --llm-cp 1 \
-  --llm-pp 1 \
+  --encoder-tp "${ENCODER_TP}" \
+  --encoder-dp "${ENCODER_DP}" \
+  --llm-offset "${LLM_OFFSET}" \
+  --llm-tp "${LLM_TP}" \
+  --llm-cp "${LLM_CP}" \
+  --llm-pp "${LLM_PP}" \
   --llm-dp "${LLM_DP}" \
-  --llm-ep 4 \
-  --llm-expt-tp 1 \
+  --llm-ep "${LLM_EP}" \
+  --llm-expt-tp "${LLM_EXPT_TP}" \
+  --tensor-parallel-num-weight-shards "${TENSOR_PARALLEL_NUM_WEIGHT_SHARDS}" \
+  --expert-tensor-parallel-num-weight-shards "${EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS}" \
   --vocab-size 131072 \
   --micro-batch-size "${MICRO_BATCH_SIZE}" \
   --global-batch-size "${GLOBAL_BATCH_SIZE}" \
diff --git a/examples/mimo/training/args.py b/examples/mimo/training/args.py
index e62bc31967b..ad666a7aef3 100644
--- a/examples/mimo/training/args.py
+++ b/examples/mimo/training/args.py
@@ -8,6 +8,7 @@
 from typing import List
 
 from examples.mimo.training.topology import ModuleGridSpec
+from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards
 from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
 
 
@@ -61,6 +62,7 @@ def add_hetero_grid_args(parser: argparse.ArgumentParser) -> argparse.ArgumentPa
 
 def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tuple[int, int]:
     """Validate the disjoint hetero grid layout; returns ``(encoder_size, llm_size)``."""
+    resolve_hetero_gtp_args(args)
     if args.llm_cp != 1:
         raise ValueError("hetero MIMO training currently supports CP=1 only")
 
@@ -76,7 +78,7 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl
             f"--num-experts ({num_experts}) must be divisible by --llm-ep ({args.llm_ep})"
         )
 
-    llm_size = args.llm_tp * args.llm_cp * args.llm_pp * args.llm_dp
+    llm_size = args.llm_tp * args.gtp_weight_remat_size * args.llm_cp * args.llm_pp * args.llm_dp
 
     if args.llm_only:
         if getattr(args, "encoder_ddp_overlap", False):
@@ -94,12 +96,14 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl
             )
         return 0, llm_size
 
-    # Fan-out divisibility: the bridge splits (mbs * llm_dp) LLM lanes across
+    # Fan-out divisibility: the bridge splits (mbs * llm_dp * gtp) LLM lanes across
     # encoder_dp encoder lanes; the split must be exact.
-    if (args.micro_batch_size * args.llm_dp) % args.encoder_dp != 0:
+    llm_data_parallel_size = args.llm_dp * args.gtp_weight_remat_size
+    if (args.micro_batch_size * llm_data_parallel_size) % args.encoder_dp != 0:
         raise ValueError(
-            "--micro-batch-size * --llm-dp must be divisible by --encoder-dp "
-            f"(got {args.micro_batch_size} * {args.llm_dp} % {args.encoder_dp} != 0)"
+            "--micro-batch-size * --llm-dp * GTP must be divisible by --encoder-dp "
+            f"(got {args.micro_batch_size} * {args.llm_dp} * "
+            f"{args.gtp_weight_remat_size} % {args.encoder_dp} != 0)"
         )
 
     encoder_size = args.encoder_tp * args.encoder_dp
@@ -134,8 +138,10 @@ def build_module_grid_specs(
         cp=args.llm_cp,
         pp=args.llm_pp,
         ep=args.llm_ep,
+        gtp_remat=args.gtp_weight_remat_size,
         rank_offset=args.llm_offset,
         expt_tp=args.llm_expt_tp or 1,
+        expt_gtp_remat=args.expert_gtp_weight_remat_size,
     )
 
     if args.llm_only:
@@ -154,6 +160,27 @@ def build_module_grid_specs(
     return [encoder_grid_spec, language_grid_spec]
 
 
+def resolve_hetero_gtp_args(args: argparse.Namespace) -> None:
+    """Resolve dense and expert GTP degrees against the language grid's TP axes."""
+    args.tensor_parallel_num_weight_shards, args.gtp_weight_remat_size = (
+        resolve_tensor_parallel_weight_shards(
+            args.llm_tp,
+            getattr(args, "tensor_parallel_num_weight_shards", None),
+            getattr(args, "gtp_weight_remat_size", 1),
+        )
+    )
+    (
+        args.expert_tensor_parallel_num_weight_shards,
+        args.expert_gtp_weight_remat_size,
+    ) = resolve_tensor_parallel_weight_shards(
+        args.llm_expt_tp or 1,
+        getattr(args, "expert_tensor_parallel_num_weight_shards", None),
+        getattr(args, "expert_gtp_weight_remat_size", 1),
+        shards_field="expert_tensor_parallel_num_weight_shards",
+        tp_field="expert_tensor_parallel_size",
+    )
+
+
 def _num_experts(args: argparse.Namespace) -> int:
     """Resolve MoE expert count from the stock --num-experts arg."""
     value = getattr(args, "num_experts", None)
diff --git a/examples/mimo/training/data.py b/examples/mimo/training/data.py
index 9139711662b..986c521874e 100644
--- a/examples/mimo/training/data.py
+++ b/examples/mimo/training/data.py
@@ -291,8 +291,12 @@ def build_train_valid_test_data_loaders(
         raise ValueError(f"unsupported dataset provider: {args.dataset_provider}")
 
     encoder_name = _encoder_name(topology)
-    if encoder_name is not None and (args.micro_batch_size * args.llm_dp) % args.encoder_dp:
-        raise ValueError("micro_batch_size * llm_dp must be divisible by encoder_dp")
+    llm_data_parallel_size = args.llm_dp * args.gtp_weight_remat_size
+    if (
+        encoder_name is not None
+        and (args.micro_batch_size * llm_data_parallel_size) % args.encoder_dp
+    ):
+        raise ValueError("micro_batch_size * llm_dp * GTP must be divisible by encoder_dp")
 
     language_grid = topology.grids[MIMO_LANGUAGE_MODULE_KEY]
     language_pgc = topology.module_pgs[MIMO_LANGUAGE_MODULE_KEY]
@@ -312,7 +316,7 @@ def build_train_valid_test_data_loaders(
     if encoder_needs_data and language_needs_data:
         raise ValueError("the external DataLoader adapter requires non-colocated module grids")
     if encoder_needs_data:
-        encoder_mbs = args.micro_batch_size * args.llm_dp // args.encoder_dp
+        encoder_mbs = args.micro_batch_size * llm_data_parallel_size // args.encoder_dp
         return _build_split_loaders(
             args,
             batch_size=encoder_mbs,
@@ -340,7 +344,8 @@ def _build_split_loaders(
     encoder_name: Optional[str],
 ) -> tuple[DataLoader, DataLoader, DataLoader]:
     """Build split-local datasets with deterministic module/DP/split seeds."""
-    base_seed = args.seed + module_seed_offset + get_pg_rank(pg_collection.dp)
+    data_group = pg_collection.dp_cp_gtp_remat or pg_collection.dp
+    base_seed = args.seed + module_seed_offset + get_pg_rank(data_group)
     common = _mock_loader_kwargs(args, encoder_name)
     return tuple(
         _build_mock_vlm_dataloader(
diff --git a/examples/mimo/training/grad_sync.py b/examples/mimo/training/grad_sync.py
index 2a06a4b8188..2b04b89c30f 100644
--- a/examples/mimo/training/grad_sync.py
+++ b/examples/mimo/training/grad_sync.py
@@ -62,7 +62,7 @@ def _is_pg_member(pg) -> bool:
 
 def _is_token_source_rank(language_pg) -> bool:
     """Whether this rank is on the LLM (last PP stage, TP rank 0) coordinate that sums
-    the global token count over DP/CP.
+    the global token count over DP/CP/GTP.
 
     Sourcing from this single coordinate avoids double-counting across TP/PP replicas.
     The _is_pg_member guards short-circuit encoder-grid ranks (non-member pp/tp groups)
@@ -101,18 +101,18 @@ def _token_source_global_rank(language_grid) -> int:
 def _global_token_count(num_tokens, language_pg, src_global_rank) -> float:
     """Total non-padded tokens in the global batch, visible on every rank.
 
-    Only the LLM token-source rank computes the count by summing over the LLM DP/CP
+    Only the LLM token-source ranks compute the count by summing over the LLM DP/CP/GTP
     group; it then broadcasts that N_global from its global rank to every rank in the
     world (including the non-colocated encoder grid, where ``language_pg`` is None) so
     both modules divide by the same per-token mean.
     """
     global_num_tokens = torch.zeros(1, dtype=torch.float32, device="cuda")
     if _is_token_source_rank(language_pg):
-        # Collective over DP/CP: every (pp_last, tp0) rank participates so the all-reduce
-        # does not hang; only DP/CP rank 0 keeps the result and is the broadcast root.
+        # Collective over every data lane: all (pp_last, tp0) GTP peers participate.
+        data_group = language_pg.dp_cp_gtp_remat or language_pg.dp_cp
         token_count = num_tokens.to(dtype=torch.float32).sum().view(1)
-        dist.all_reduce(token_count, group=language_pg.dp_cp, op=dist.ReduceOp.SUM)
-        if dist.get_rank(group=language_pg.dp_cp) == 0:
+        dist.all_reduce(token_count, group=data_group, op=dist.ReduceOp.SUM)
+        if dist.get_rank(group=data_group) == 0:
             global_num_tokens.copy_(token_count)
     dist.broadcast(global_num_tokens, src=src_global_rank)
     return float(global_num_tokens.item())
diff --git a/examples/mimo/training/runtime.py b/examples/mimo/training/runtime.py
index 6a5235aba13..8b46f1fb5f2 100644
--- a/examples/mimo/training/runtime.py
+++ b/examples/mimo/training/runtime.py
@@ -40,7 +40,7 @@ def configure_module_rng(
     so disjoint modules (and stages) get independent RNG state. Caller invokes once per active
     module on this rank.
     """
-    for _required in ("pp", "dp", "tp", "ep", "expt_tp"):
+    for _required in ("pp", "dp", "tp", "ep", "expt_tp", "gtp_remat", "expt_gtp_remat"):
         assert (
             getattr(pg_collection, _required, None) is not None
         ), f"pg_collection passed to configure_module_rng must define {_required}"
@@ -52,6 +52,8 @@ def configure_module_rng(
         tp_group=pg_collection.tp,
         ep_group=pg_collection.ep,
         etp_group=pg_collection.expt_tp,
+        gtp_remat_group=pg_collection.gtp_remat,
+        egtp_remat_group=pg_collection.expt_gtp_remat,
     )
 
 
diff --git a/examples/mimo/training/topology.py b/examples/mimo/training/topology.py
index 60e473f58ff..351d6de3dcc 100644
--- a/examples/mimo/training/topology.py
+++ b/examples/mimo/training/topology.py
@@ -33,23 +33,28 @@ class ModuleGridSpec:
     cp: int = 1
     pp: int = 1
     ep: int = 1
+    gtp_remat: int = 1
     rank_offset: int = 0
     # Experts default to TP=1 (set explicitly for MoE); intentionally not Megatron's etp=tp default.
     expt_tp: int = 1
+    expt_gtp_remat: int = 1
     dp: int = field(init=False)
     expt_dp: int = field(init=False)
 
     def __post_init__(self) -> None:
-        dense = self.tp * self.cp * self.pp
+        dense = self.tp * self.gtp_remat * self.cp * self.pp
         if self.num_ranks % dense != 0:
             raise ValueError(
-                f"num_ranks ({self.num_ranks}) must be divisible by tp*cp*pp ({dense})"
+                "num_ranks "
+                f"({self.num_ranks}) must be divisible by tp*gtp_remat*cp*pp ({dense})"
             )
         self.dp = self.num_ranks // dense
-        expert = self.expt_tp * self.ep * self.pp
+        expert = self.expt_tp * self.ep * self.expt_gtp_remat * self.pp
         if self.num_ranks % expert != 0:
             raise ValueError(
-                f"num_ranks ({self.num_ranks}) must be divisible by expt_tp*ep*pp ({expert})"
+                "num_ranks "
+                f"({self.num_ranks}) must be divisible by "
+                f"expt_tp*ep*expt_gtp_remat*pp ({expert})"
             )
         self.expt_dp = self.num_ranks // expert
 
@@ -118,27 +123,45 @@ def create_topology(specs: list[ModuleGridSpec]) -> HeteroTopology:
 def _build_grid(spec: ModuleGridSpec) -> HyperCommGrid:
     """Create a dense grid plus its expert view and the process groups MIMO needs."""
     grid = HyperCommGrid(
-        shape=[spec.tp, spec.cp, spec.dp, spec.pp],
-        dim_names=["tp", "cp", "dp", "pp"],
+        shape=[spec.tp, spec.gtp_remat, spec.cp, spec.dp, spec.pp],
+        dim_names=["tp", "gtp_remat", "cp", "dp", "pp"],
         rank_offset=spec.rank_offset,
         backend="nccl",
     )
     # Expert factorization over the same rank span; pp is shared with the base view.
     grid.register_view(
         _EXPERT_VIEW,
-        shape=[spec.expt_tp, spec.ep, spec.expt_dp, spec.pp],
-        dim_names=["expt_tp", "ep", "expt_dp", "pp"],
+        shape=[spec.expt_tp, spec.ep, spec.expt_gtp_remat, spec.expt_dp, spec.pp],
+        dim_names=["expt_tp", "ep", "expt_gtp_remat", "expt_dp", "pp"],
         shared_dims=["pp"],
     )
 
     try:
         for dims in (
-            ["tp"], ["cp"], ["pp"], ["dp"],
-            ["dp", "cp"], ["tp", "cp"], ["tp", "pp"],
-            ["tp", "dp"], ["tp", "dp", "cp"], ["tp", "cp", "dp", "pp"],
+            ["tp"],
+            ["gtp_remat"],
+            ["cp"],
+            ["pp"],
+            ["dp"],
+            ["dp", "cp"],
+            ["gtp_remat", "dp", "cp"],
+            ["tp", "cp"],
+            ["tp", "gtp_remat", "pp"],
+            ["tp", "gtp_remat", "dp"],
+            ["tp", "gtp_remat", "dp", "cp"],
+            ["tp", "gtp_remat", "cp", "dp", "pp"],
         ):
             grid.create_pg(dims)
-        for dims in (["ep"], ["expt_tp"], ["expt_dp"], ["expt_tp", "ep"], ["expt_tp", "ep", "pp"]):
+        for dims in (
+            ["ep"],
+            ["expt_tp"],
+            ["expt_gtp_remat"],
+            ["expt_dp"],
+            ["expt_tp", "ep"],
+            ["expt_tp", "ep", "pp"],
+            ["expt_tp", "ep", "expt_gtp_remat", "pp"],
+            ["expt_gtp_remat", "expt_dp"],
+        ):
             grid.create_pg(dims, view=_EXPERT_VIEW)
     except Exception:
         grid.destroy()
@@ -193,18 +216,27 @@ def pg_collection_from_grid(
     pgc.pp = grid.get_pg("pp")
     pgc.dp = grid.get_pg("dp")
     pgc.dp_cp = grid.get_pg(["dp", "cp"])
+    pgc.dp_cp_gtp_remat = grid.get_pg(["gtp_remat", "dp", "cp"])
     pgc.intra_dp_cp = pgc.dp_cp
+    pgc.gtp_remat = grid.get_pg("gtp_remat")
     pgc.tp_cp = grid.get_pg(["tp", "cp"])
-    pgc.tp_dp = grid.get_pg(["tp", "dp"])
-    pgc.tp_dp_cp = grid.get_pg(["tp", "dp", "cp"])
-    pgc.mp = grid.get_pg(["tp", "pp"])
-    pgc.intra_dist_opt = grid.get_pg(["tp", "cp", "dp", "pp"])
+    pgc.tp_dp = grid.get_pg(["tp", "gtp_remat", "dp"])
+    pgc.tp_dp_cp = grid.get_pg(["tp", "gtp_remat", "dp", "cp"])
+    pgc.mp = grid.get_pg(["tp", "gtp_remat", "pp"])
+    pgc.intra_dist_opt = grid.get_pg(["tp", "gtp_remat", "cp", "dp", "pp"])
     pgc.ep = grid.get_pg("ep", view=_EXPERT_VIEW)
     pgc.expt_tp = grid.get_pg("expt_tp", view=_EXPERT_VIEW)
+    pgc.expt_gtp_remat = grid.get_pg("expt_gtp_remat", view=_EXPERT_VIEW)
     pgc.expt_dp = grid.get_pg("expt_dp", view=_EXPERT_VIEW)
+    pgc.expt_dp_gtp_remat = grid.get_pg(
+        ["expt_gtp_remat", "expt_dp"], view=_EXPERT_VIEW
+    )
     pgc.intra_expt_dp = pgc.expt_dp
     pgc.tp_ep = grid.get_pg(["expt_tp", "ep"], view=_EXPERT_VIEW)
     pgc.tp_ep_pp = grid.get_pg(["expt_tp", "ep", "pp"], view=_EXPERT_VIEW)
+    pgc.tp_ep_pp_with_egtp_remat = grid.get_pg(
+        ["expt_tp", "ep", "expt_gtp_remat", "pp"], view=_EXPERT_VIEW
+    )
     pgc.embd = None
     pgc.pos_embd = None
     if is_language:
diff --git a/megatron/core/models/mimo/model/base.py b/megatron/core/models/mimo/model/base.py
index b226b5c1e4b..b60b7734cdd 100644
--- a/megatron/core/models/mimo/model/base.py
+++ b/megatron/core/models/mimo/model/base.py
@@ -549,6 +549,10 @@ def _attach_modality_split_sizes(
             language_grid = grid_map[MIMO_LANGUAGE_MODULE_KEY]
             encoder_dp = encoder_grid.shape[encoder_grid.dim_names.index("dp")]
             language_dp = language_grid.shape[language_grid.dim_names.index("dp")]
+            if "gtp_remat" in language_grid.dim_names:
+                language_dp *= language_grid.shape[
+                    language_grid.dim_names.index("gtp_remat")
+                ]
             assert encoder_dp <= language_dp, (
                 f"Bridge fan-out split metadata with non-uniform per-sample sizes "
                 f"requires encoder DP <= LM DP (got encoder='{encoder_name}' "
diff --git a/megatron/core/models/mimo/optimizer.py b/megatron/core/models/mimo/optimizer.py
index 598f6c883af..d2e7f18c4b5 100644
--- a/megatron/core/models/mimo/optimizer.py
+++ b/megatron/core/models/mimo/optimizer.py
@@ -16,6 +16,7 @@
 from megatron.core.optimizer.optimizer import MegatronOptimizer
 from megatron.core.optimizer.optimizer_config import OptimizerConfig
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.utils import unwrap_model
 
 if TYPE_CHECKING:
     from megatron.core.hyper_comm_grid import HyperCommGrid
@@ -339,29 +340,6 @@ def _get_replica_id(pg_collection: Optional[ProcessGroupCollection]) -> tuple:
     return (pg_collection.tp.rank(), pg_collection.pp.rank(), pg_collection.dp.rank())
 
 
-_EXPERT_VIEW = "expert"
-
-
-def _get_pg_collection_for_optimizer(grid) -> ProcessGroupCollection:
-    """Derive the optimizer's ProcessGroupCollection from a populated HyperCommGrid.
-
-    Dense groups come from the base view; expert-parallel groups (tp_ep_pp, expt_dp) come from
-    the grid's dedicated expert view -- expert parallelism is always factored into a separate
-    view (expt_tp/ep/expt_dp), never the base view. All groups must be pre-created on the grid.
-    """
-    pg = ProcessGroupCollection()
-    pg.dp = grid.get_pg("dp")
-    pg.dp_cp = grid.get_pg(["dp", "cp"])
-    pg.tp = grid.get_pg("tp")
-    pg.pp = grid.get_pg("pp")
-    pg.mp = grid.get_pg(["tp", "pp"])
-    pg.tp_ep_pp = grid.get_pg(["expt_tp", "ep", "pp"], view=_EXPERT_VIEW)
-    pg.expt_dp = grid.get_pg("expt_dp", view=_EXPERT_VIEW)
-    # Distributed-optimizer grad-stats group spans the dense shards (mirrors the topology PGC).
-    pg.intra_dist_opt = grid.get_pg(["tp", "cp", "dp", "pp"])
-    return pg
-
-
 def get_mimo_optimizer(mimo_model: "MimoModel", config: OptimizerConfig) -> MimoOptimizer:
     """Create optimizer for MimoModel with heterogeneous parallelism."""
     from megatron.core.optimizer import get_megatron_optimizer
@@ -386,7 +364,10 @@ def get_mimo_optimizer(mimo_model: "MimoModel", config: OptimizerConfig) -> Mimo
                 module = mimo_model.modality_submodules[module_name]
 
             if module is not None:
-                pg_collection = _get_pg_collection_for_optimizer(grid)
+                pg_collection = getattr(unwrap_model(module), 'pg_collection', None)
+                assert pg_collection is not None, (
+                    f"Module '{module_name}' must own a ProcessGroupCollection before optimizer setup"
+                )
                 assert (
                     not hasattr(module, 'ddp_config')
                     or module.ddp_config is None
diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py
index a9619ea4819..22aaf515a8f 100644
--- a/megatron/core/tensor_parallel/random.py
+++ b/megatron/core/tensor_parallel/random.py
@@ -454,6 +454,8 @@ def model_parallel_cuda_manual_seed(
     gtp_remat_rank: Optional[int] = None,
     egtp_remat_rank: Optional[int] = None,
     force_reset_rng: bool = False,
+    gtp_remat_world_size: Optional[int] = None,
+    egtp_remat_world_size: Optional[int] = None,
 ):
     """Initialize model parallel cuda seed.
 
@@ -482,6 +484,10 @@ def model_parallel_cuda_manual_seed(
         gtp_remat_rank = get_gtp_weight_remat_rank()
     if egtp_remat_rank is None:
         egtp_remat_rank = get_expert_gtp_weight_remat_rank()
+    if gtp_remat_world_size is None:
+        gtp_remat_world_size = get_gtp_weight_remat_world_size()
+    if egtp_remat_world_size is None:
+        egtp_remat_world_size = get_expert_gtp_weight_remat_world_size()
     # 2718 is just for fun and any POSITIVE value will work.
     offset = seed + 2718
     tensor_model_parallel_seed = offset + tp_rank
@@ -506,10 +512,10 @@ def model_parallel_cuda_manual_seed(
     # must draw DIFFERENT values (everything above is identical across peers by design). The 65536
     # stride keeps these disjoint from the tp/ep/etp seeds. Added only when the axis is active, so
     # non-GTP runs keep a byte-identical tracker set (and checkpoint rng payload).
-    if get_gtp_weight_remat_world_size() > 1:
+    if gtp_remat_world_size > 1:
         gtp_remat_seed = tensor_model_parallel_seed + 65536 * (1 + gtp_remat_rank)
         _CUDA_RNG_STATE_TRACKER.add(_GTP_REMAT_RNG_TRACKER_NAME, gtp_remat_seed)
-    if get_expert_gtp_weight_remat_world_size() > 1:
+    if egtp_remat_world_size > 1:
         egtp_remat_seed = expert_parallel_seed + 32768 + 65536 * (1 + egtp_remat_rank)
         _CUDA_RNG_STATE_TRACKER.add(_EXPERT_GTP_REMAT_RNG_TRACKER_NAME, egtp_remat_seed)
 
diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py
index 87d6aa65b03..ed8cf3cdf6c 100644
--- a/megatron/training/initialize.py
+++ b/megatron/training/initialize.py
@@ -26,7 +26,13 @@
 from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
     enable_batch_invariant_mode,
 )
-from megatron.core.utils import get_pg_rank, get_te_version, is_te_min_version, is_torch_min_version
+from megatron.core.utils import (
+    get_pg_rank,
+    get_pg_size,
+    get_te_version,
+    is_te_min_version,
+    is_torch_min_version,
+)
 from megatron.training import (
     get_adlr_autoresume,
     get_args,
@@ -426,10 +432,12 @@ def _set_random_seed(
     tp_group: Optional[torch.distributed.ProcessGroup] = None,
     ep_group: Optional[torch.distributed.ProcessGroup] = None,
     etp_group: Optional[torch.distributed.ProcessGroup] = None,
+    gtp_remat_group: Optional[torch.distributed.ProcessGroup] = None,
+    egtp_remat_group: Optional[torch.distributed.ProcessGroup] = None,
 ):
     """Set random seed for reproducability.
 
-    The optional pp/dp/tp/ep/etp groups let a caller without an initialized mpu
+    The optional parallel groups let a caller without an initialized mpu
     (e.g. a disjoint-grid run) supply the parallel ranks explicitly; each falls
     back to the mpu group when None.
     """
@@ -448,6 +456,16 @@ def _set_random_seed(
             tp_rank = get_pg_rank(tp_group) if tp_group is not None else None
             ep_rank = get_pg_rank(ep_group) if ep_group is not None else None
             etp_rank = get_pg_rank(etp_group) if etp_group is not None else None
+            gtp_remat_rank = get_pg_rank(gtp_remat_group) if gtp_remat_group is not None else None
+            egtp_remat_rank = (
+                get_pg_rank(egtp_remat_group) if egtp_remat_group is not None else None
+            )
+            gtp_remat_world_size = (
+                get_pg_size(gtp_remat_group) if gtp_remat_group is not None else None
+            )
+            egtp_remat_world_size = (
+                get_pg_size(egtp_remat_group) if egtp_remat_group is not None else None
+            )
             tensor_parallel.model_parallel_cuda_manual_seed(
                 seed,
                 te_rng_tracker,
@@ -456,6 +474,10 @@ def _set_random_seed(
                 tp_rank=tp_rank,
                 ep_rank=ep_rank,
                 etp_rank=etp_rank,
+                gtp_remat_rank=gtp_remat_rank,
+                egtp_remat_rank=egtp_remat_rank,
+                gtp_remat_world_size=gtp_remat_world_size,
+                egtp_remat_world_size=egtp_remat_world_size,
             )
     else:
         raise ValueError("Seed ({}) should be a positive integer.".format(seed_))
diff --git a/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py b/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py
index fb923c2b37d..cee45406b39 100644
--- a/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py
+++ b/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py
@@ -62,6 +62,33 @@ def test_canonical_layout_validates_and_maps_specs():
     assert language_grid_spec.expt_tp == 1
 
 
+def test_gtp_layout_validates_and_maps_weight_shard_axes():
+    args = _layout_8gpu_20l(
+        llm_dp=1,
+        llm_ep=2,
+        tensor_parallel_num_weight_shards=4,
+        expert_tensor_parallel_num_weight_shards=2,
+    )
+
+    assert validate_hetero_grid_args(args, WORLD_SIZE_8) == (4, 4)
+    assert args.gtp_weight_remat_size == 2
+    assert args.expert_gtp_weight_remat_size == 2
+
+    _, language_grid_spec = build_module_grid_specs(
+        args, WORLD_SIZE_8, encoder_module_name="radio_encoder"
+    )
+    assert language_grid_spec.gtp_remat == 2
+    assert language_grid_spec.dp == 1
+    assert language_grid_spec.expt_gtp_remat == 2
+    assert language_grid_spec.expt_dp == 1
+
+
+def test_weight_shards_must_divide_language_tp():
+    args = _layout_8gpu_20l(tensor_parallel_num_weight_shards=3)
+    with pytest.raises(ValueError, match="must be divisible"):
+        validate_hetero_grid_args(args, WORLD_SIZE_8)
+
+
 def test_overlapping_spans_raise():
     # llm-offset 2 makes llm ranks {2,3,4,5} overlap encoder ranks {0,1,2,3}.
     args = _layout_8gpu_20l(llm_offset=2)
diff --git a/tests/unit_tests/test_mimo_hetero_topology.py b/tests/unit_tests/test_mimo_hetero_topology.py
index 7d68e107bf7..95ba6c5c7d4 100644
--- a/tests/unit_tests/test_mimo_hetero_topology.py
+++ b/tests/unit_tests/test_mimo_hetero_topology.py
@@ -25,6 +25,21 @@ def _specs():
     ]
 
 
+def _gtp_specs():
+    return [
+        ModuleGridSpec(name=ENCODER, num_ranks=4, tp=2, rank_offset=0),
+        ModuleGridSpec(
+            name=MIMO_LANGUAGE_MODULE_KEY,
+            num_ranks=4,
+            tp=2,
+            gtp_remat=2,
+            ep=2,
+            expt_gtp_remat=2,
+            rank_offset=4,
+        ),
+    ]
+
+
 class TestModuleGridSpecResolution:
     def test_derived_dims_resolve_to_concrete_ints(self):
         # num_ranks=4,tp=2 with default expt_tp=1: dp=2, expt_dp=4.
@@ -87,6 +102,33 @@ def test_pgc_group_sizes(self):
         finally:
             topo.destroy()
 
+    def test_gtp_pgc_group_sizes(self):
+        topo = create_topology(_gtp_specs())
+        try:
+            rank = dist.get_rank()
+            pgc = (
+                topo.module_pgs[ENCODER]
+                if rank < 4
+                else topo.module_pgs[MIMO_LANGUAGE_MODULE_KEY]
+            )
+            if rank < 4:
+                assert pgc.gtp_remat.size() == 1
+                assert pgc.expt_gtp_remat.size() == 1
+            else:
+                assert pgc.tp.size() == 2
+                assert pgc.gtp_remat.size() == 2
+                assert pgc.dp.size() == 1
+                assert pgc.dp_cp_gtp_remat.size() == 2
+                assert pgc.mp.size() == 4
+                assert pgc.expt_tp.size() == 1
+                assert pgc.ep.size() == 2
+                assert pgc.expt_gtp_remat.size() == 2
+                assert pgc.expt_dp.size() == 1
+                assert pgc.expt_dp_gtp_remat.size() == 2
+                assert pgc.tp_ep_pp_with_egtp_remat.size() == 4
+        finally:
+            topo.destroy()
+
     def test_embedding_groups(self):
         # Language grid is tp=2,pp=2 at ranks 4-7: each PP group is [first,last] (size 2),
         # so first/last-stage ranks get a 2-rank .embd and the first stage gets .pos_embd.