Enable GTP for heterogeneous MIMO training

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.

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
 

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
 

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"})

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}" \

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)

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(

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())

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,
     )