Wire CP into MIMO

examples/mimo/model_providers/llava_avlm.py

@@ -157,7 +157,10 @@ def model_provider_llava_avlm(
     # Create MIMO model
     cp_group = pg_collection.cp if pg_collection is not None else None
     tp_group = pg_collection.tp if pg_collection is not None else None
-    mimo_model = MimoModel(mimo_model_config, cp_group=cp_group, tp_group=tp_group)
+    tp_cp_group = pg_collection.tp_cp if pg_collection is not None else None
+    mimo_model = MimoModel(
+        mimo_model_config, cp_group=cp_group, tp_group=tp_group, tp_cp_group=tp_cp_group
+    )
     print("*"*100)
     print_mimo_structure(mimo_model)
     print("*"*100)
@@ -184,4 +187,4 @@ def model_provider_llava_avlm(
         for param in module.parameters():
             param.requires_grad = False
 
-    return mimo_model
+    return mimo_model

examples/mimo/model_providers/llava_vlm.py

@@ -129,7 +129,10 @@ def model_provider_llava_vlm(
     # Create MIMO model
     cp_group = pg_collection.cp if pg_collection is not None else None
     tp_group = pg_collection.tp if pg_collection is not None else None
-    mimo_model = MimoModel(mimo_model_config, cp_group=cp_group, tp_group=tp_group)
+    tp_cp_group = pg_collection.tp_cp if pg_collection is not None else None
+    mimo_model = MimoModel(
+        mimo_model_config, cp_group=cp_group, tp_group=tp_group, tp_cp_group=tp_cp_group
+    )
     print("*"*100)
     print_mimo_structure(mimo_model)
     print("*"*100)

examples/mimo/model_providers/mock.py

@@ -97,6 +97,9 @@ def model_provider_mock_vlm_single_encoder(
     # Create MIMO model
     cp_group = pg_collection.cp if pg_collection is not None else None
     tp_group = pg_collection.tp if pg_collection is not None else None
-    mimo_model = MimoModel(mimo_model_config, cp_group=cp_group, tp_group=tp_group)
+    tp_cp_group = pg_collection.tp_cp if pg_collection is not None else None
+    mimo_model = MimoModel(
+        mimo_model_config, cp_group=cp_group, tp_group=tp_group, tp_cp_group=tp_cp_group
+    )
 
     return mimo_model

examples/mimo/model_providers/nemotron_moe_vlm.py

@@ -134,6 +134,7 @@ def nemotron_language_config(
 ) -> TransformerConfig:
     """Nemotron6-MoE language config: stock from-args base + model-specific overrides."""
     config = deepcopy(_base_config(args))
+    gtp_remat_num_sequence_shards = config.gtp_remat_num_sequence_shards
     # Code-only fields + hetero parallelism pins.
     config.variable_seq_lengths = True
     config.expert_model_parallel_size = ep_size
@@ -158,6 +159,7 @@ def nemotron_language_config(
         )
     )
     config.sequence_parallel = tp_size > 1
+    config.tensor_parallel_num_sequence_shards = tp_size * gtp_remat_num_sequence_shards
     config.position_embedding_type = "none"
     return config
 
@@ -202,6 +204,7 @@ def nemotron_projection_config(
     if base_config is None:
         _disable_gtp(config)
     config.sequence_parallel = False
+    config.tensor_parallel_num_sequence_shards = None
     return config
 
 
@@ -243,7 +246,9 @@ def language_model_spec(
         pp_rank = 0
         pp_size = get_grid_dim_size(llm_grid, "pp")
         tp_size = get_grid_dim_size(llm_grid, "tp")
-        cp_size = get_grid_dim_size(llm_grid, "cp")
+        cp_size = get_grid_dim_size(llm_grid, "cp") * get_grid_dim_size(
+            llm_grid, "gtp_remat_sequence_shards"
+        )
         ep_size = getattr(args, "mimo_llm_ep", 1)
         expt_tp_size = getattr(args, "mimo_llm_expt_tp", None) or 1
     else:

examples/mimo/model_providers/radio_encoder.py

@@ -109,6 +109,7 @@ def _disable_gtp(config: TransformerConfig) -> None:
     """Keep this module replicated across any LLM GTP axes."""
     config.tensor_parallel_num_weight_shards = config.tensor_model_parallel_size
     config.gtp_weight_remat_size = 1
+    config.tensor_parallel_num_sequence_shards = None
     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

examples/mimo/train.py

@@ -53,6 +53,27 @@
     "llava_avlm": llava_avlm_dataloader_provider,
 }
 
+
+def _canonicalize_packed_sequence_metadata(batch: Dict[str, Any]) -> Dict[str, Any]:
+    """Convert the legacy Energon packing dictionary to canonical MIMO batch fields."""
+    packing_kwargs = batch.pop("packing_kwargs", None)
+    if packing_kwargs is None:
+        return batch
+
+    cu_seqlens = packing_kwargs["cu_seqlens_q"].to(dtype=torch.int32)
+    batch["cu_seqlens"] = cu_seqlens.unsqueeze(0) if cu_seqlens.dim() == 1 else cu_seqlens
+
+    cu_seqlens_padded = packing_kwargs.get("cu_seqlens_q_padded")
+    if cu_seqlens_padded is not None:
+        cu_seqlens_padded = cu_seqlens_padded.to(dtype=torch.int32)
+        batch["cu_seqlens_padded"] = (
+            cu_seqlens_padded.unsqueeze(0) if cu_seqlens_padded.dim() == 1 else cu_seqlens_padded
+        )
+
+    batch["max_seqlen"] = packing_kwargs["max_seqlen_q"].to(dtype=torch.int32).reshape(1)
+    return batch
+
+
 def add_mimo_args(parser):
     """Add MIMO-specific arguments to the parser."""
     group = parser.add_argument_group('MIMO', 'MIMO specific arguments')
@@ -131,12 +152,14 @@ def get_batch(data_iterator: Iterator[Dict[str, Any]]):
     # loss_mask: Optional[torch.Tensor] = None,
     # labels: Optional[torch.Tensor] = None,
     # modality_inputs: Optional[Dict[str, Dict[str, Any]]] = None,
-    # packing_kwargs: Optional[dict] = None,
+    # cu_seqlens: Optional[torch.Tensor] = None,
+    # cu_seqlens_padded: Optional[torch.Tensor] = None,
+    # max_seqlen: Optional[torch.Tensor] = None,
 
     # For the modality inputs, the keys can be arbitrary
     # so we do a broadcast of the schema followed by a broadcast of the actual data
     # check broadcast_nested_data_batch for more details
-    batch = broadcast_nested_data_batch(data)
+    batch = _canonicalize_packed_sequence_metadata(broadcast_nested_data_batch(data))
 
     return batch
 

examples/mimo/training/args.py

@@ -8,8 +8,12 @@
 from typing import List
 
 from examples.mimo.training.topology import ModuleGridSpec
+from examples.mimo.utils.hetero import get_language_batch_parallel_size
 from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards
 from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
+from megatron.core.tensor_parallel.gtp_parallel_layout import (
+    resolve_tensor_parallel_sequence_shards,
+)
 
 
 def add_hetero_grid_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
@@ -98,9 +102,9 @@ 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)``."""
-    gtp_weight_remat_size, _ = resolve_hetero_gtp_degrees(args)
     if args.mimo_llm_cp < 1:
         raise ValueError("--mimo-llm-cp must be positive")
+    llm_data_parallel_size = get_language_batch_parallel_size(args)
 
     if getattr(args, "mimo_encoder_ddp_overlap", False) and not getattr(
         args, "overlap_grad_reduce", False
@@ -120,12 +124,13 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl
             f"--mimo-llm-ep ({args.mimo_llm_ep})"
         )
 
+    gtp_weight_remat_size, _ = resolve_hetero_gtp_degrees(args)
     llm_size = (
         args.mimo_llm_tp
-        * gtp_weight_remat_size
         * args.mimo_llm_cp
         * args.mimo_llm_pp
         * args.mimo_llm_dp
+        * gtp_weight_remat_size
     )
 
     if args.mimo_llm_only:
@@ -147,13 +152,12 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl
 
     # Fan-out divisibility: the bridge splits every LLM data lane across
     # mimo_encoder_dp encoder lanes; the split must be exact.
-    llm_data_parallel_size = args.mimo_llm_dp * gtp_weight_remat_size
     if (args.micro_batch_size * llm_data_parallel_size) % args.mimo_encoder_dp != 0:
         raise ValueError(
-            "--micro-batch-size * --mimo-llm-dp * GTP must be divisible by "
+            "--micro-batch-size * language sample lanes must be divisible by "
             "--mimo-encoder-dp "
-            f"(got {args.micro_batch_size} * {args.mimo_llm_dp} * "
-            f"{gtp_weight_remat_size} % {args.mimo_encoder_dp} != 0)"
+            f"(got {args.micro_batch_size} * {llm_data_parallel_size} "
+            f"% {args.mimo_encoder_dp} != 0)"
         )
 
     encoder_size = args.mimo_encoder_tp * args.mimo_encoder_dp
@@ -182,6 +186,12 @@ def build_module_grid_specs(
     """Map grid args to the ModuleGridSpec list create_topology consumes."""
     encoder_size, llm_size = validate_hetero_grid_args(args, world_size)
     gtp_weight_remat_size, expert_gtp_weight_remat_size = resolve_hetero_gtp_degrees(args)
+    _, gtp_remat_num_sequence_shards = resolve_tensor_parallel_sequence_shards(
+        args.mimo_llm_tp,
+        getattr(args, "tensor_parallel_num_sequence_shards", None),
+        gtp_weight_remat_size,
+        getattr(args, "sequence_parallel", False),
+    )
 
     language_grid_spec = ModuleGridSpec(
         name=MIMO_LANGUAGE_MODULE_KEY,
@@ -191,6 +201,7 @@ def build_module_grid_specs(
         pp=args.mimo_llm_pp,
         ep=args.mimo_llm_ep,
         gtp_remat=gtp_weight_remat_size,
+        gtp_remat_num_sequence_shards=gtp_remat_num_sequence_shards,
         rank_offset=args.mimo_llm_offset,
         expt_tp=args.mimo_llm_expt_tp or 1,
         expt_gtp_remat=expert_gtp_weight_remat_size,

examples/mimo/training/builder.py

@@ -120,6 +120,7 @@ def build_model(
             mimo_config,
             cp_group=active_pg.cp if is_language else None,
             tp_group=active_pg.tp if is_language else None,
+            tp_cp_group=active_pg.tp_cp if is_language else None,
         )
 
     def build_distributed_models(

examples/mimo/training/data.py

@@ -12,10 +12,10 @@
 from torch.utils.data import DataLoader, Dataset
 
 from examples.mimo.training.topology import HeteroTopology
+from examples.mimo.utils.hetero import get_data_lane_rank, get_language_batch_parallel_size
 from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage
-from megatron.core.utils import get_pg_rank
 from megatron.training.global_vars import get_run_config
 
 _ENCODER_SEED_OFFSET = 10_000
@@ -292,13 +292,13 @@ def build_train_valid_test_data_loaders(
         raise ValueError(f"unsupported dataset provider: {args.dataset_provider}")
 
     encoder_name = _encoder_name(topology)
-    llm_data_parallel_size = args.mimo_llm_dp * args.gtp_weight_remat_size
+    llm_data_parallel_size = get_language_batch_parallel_size(args)
     if (
         encoder_name is not None
         and (args.micro_batch_size * llm_data_parallel_size) % args.mimo_encoder_dp
     ):
         raise ValueError(
-            "micro_batch_size * mimo_llm_dp * GTP must be divisible by mimo_encoder_dp"
+            "micro_batch_size * language sample lanes must be divisible by mimo_encoder_dp"
         )
 
     language_grid = topology.grids[MIMO_LANGUAGE_MODULE_KEY]
@@ -348,14 +348,7 @@ def _build_split_loaders(
 ) -> tuple[DataLoader, DataLoader, DataLoader]:
     """Build split-local datasets with deterministic module/DP/split seeds."""
     cfg = get_run_config()
-    data_group = pg_collection.dp_cp_gtp_remat or pg_collection.dp
-    lane_rank = get_pg_rank(data_group)
-    if pg_collection.dp_cp_gtp_remat is not None:
-        # The combined group orders CP first (fastest), then GTP and DP.
-        # CP replicas consume the same full batch before the model shards it;
-        # GTP and DP remain distinct data lanes, matching the bridge topology.
-        lane_rank //= pg_collection.cp.size()
-    base_seed = cfg.rng.seed + module_seed_offset + lane_rank
+    base_seed = cfg.rng.seed + module_seed_offset + get_data_lane_rank(pg_collection)
     common = _mock_loader_kwargs(args, encoder_name)
     return tuple(
         _build_mock_vlm_dataloader(