[data] fix: preserve sparse MIMO modality DP ownership

src/megatron/bridge/data/megatron_mimo/collate.py

@@ -9,6 +9,9 @@
 import torch
 
 
+_MIMO_SAMPLE_INDICES_KEY = "_mimo_sample_indices"
+
+
 def megatron_mimo_collate_fn(
     batch: List[Dict[str, Any]],
     modality_names: List[str],
@@ -65,6 +68,7 @@ def megatron_mimo_collate_fn(
 
     # Collate modality inputs
     modality_inputs: Dict[str, Dict[str, Any]] = {}
+    modality_sample_indices: Dict[str, Dict[str, torch.Tensor]] = {}
 
     for modality_name in modality_names:
         # Collect all tensors for this modality across the batch
@@ -81,12 +85,15 @@ def megatron_mimo_collate_fn(
             continue
 
         modality_inputs[modality_name] = {}
+        modality_sample_indices[modality_name] = {}
 
         for key in first_non_empty.keys():
             values = []
-            for item in modality_batch_items:
+            sample_indices = []
+            for sample_idx, item in enumerate(modality_batch_items):
                 if key in item:
                     val = item[key]
+                    sample_indices.append(sample_idx)
                     if isinstance(val, torch.Tensor):
                         values.append(val)
                     else:
@@ -109,12 +116,20 @@ def megatron_mimo_collate_fn(
             elif values:
                 # Keep non-tensor values as list
                 modality_inputs[modality_name][key] = values
+            modality_sample_indices[modality_name][key] = torch.tensor(sample_indices, dtype=torch.long)
 
-    return {
+    collated = {
         "input_ids": input_ids,
         "labels": labels,
         "loss_mask": loss_mask,
         "attention_mask": attention_mask,
         "position_ids": position_ids,
         "modality_inputs": modality_inputs,
     }
+    if any(
+        indices.numel() != len(batch)
+        for field_indices in modality_sample_indices.values()
+        for indices in field_indices.values()
+    ):
+        collated[_MIMO_SAMPLE_INDICES_KEY] = modality_sample_indices
+    return collated

src/megatron/bridge/data/megatron_mimo/dp_utils.py

@@ -17,6 +17,9 @@
     from megatron.bridge.models.megatron_mimo.megatron_mimo_config import MegatronMIMOParallelismConfig
 
 
+_MIMO_SAMPLE_INDICES_KEY = "_mimo_sample_indices"
+
+
 def _batch_dim_for_tensor(key: str, value: torch.Tensor) -> int:
     """Return the batch dimension for known MegatronMIMO batch tensors."""
     # Qwen-VL MRoPE position_ids are [3, batch, seq] instead of [batch, seq].
@@ -234,11 +237,35 @@ def slice_batch_for_megatron_mimo(
         >>> local_batch = slice_batch_for_megatron_mimo(global_batch, dp_rank=1, dp_size=3)
         >>> local_batch['tokens'].shape  # torch.Size([4, 2048])
     """
-    if dp_size == 1:
+    modality_sample_indices = batch.get(_MIMO_SAMPLE_INDICES_KEY)
+    if dp_size == 1 and modality_sample_indices is None:
         return batch
 
+    global_batch_size = None
+    for key in ("input_ids", "labels", "loss_mask"):
+        value = batch.get(key)
+        if isinstance(value, torch.Tensor):
+            global_batch_size = value.size(_batch_dim_for_tensor(key, value))
+            break
+    if modality_sample_indices is not None and global_batch_size is None:
+        raise ValueError("Sparse MegatronMIMO modality inputs require a sample-batched tensor.")
+
+    sample_start = 0
+    sample_end = global_batch_size
+    if global_batch_size is not None and dp_size > 1:
+        if global_batch_size % dp_size != 0:
+            raise ValueError(
+                f"Batch size {global_batch_size} is not divisible by DP size {dp_size}. "
+                "Ensure micro_batch_size is divisible by every module's data_parallel_size."
+            )
+        local_batch_size = global_batch_size // dp_size
+        sample_start = dp_rank * local_batch_size
+        sample_end = sample_start + local_batch_size
+
     sliced = {}
     for key, value in batch.items():
+        if key == _MIMO_SAMPLE_INDICES_KEY:
+            continue
         if isinstance(value, torch.Tensor):
             batch_dim = _batch_dim_for_tensor(key, value)
             batch_size = value.size(batch_dim)
@@ -255,10 +282,35 @@ def slice_batch_for_megatron_mimo(
             index[batch_dim] = builtins.slice(start_idx, end_idx)
             sliced[key] = value[tuple(index)]
         elif isinstance(value, dict):
+            if key == "modality_inputs" and modality_sample_indices is not None:
+                local_modalities = {}
+                for modality_name, modality_values in value.items():
+                    local_values = {}
+                    field_indices = modality_sample_indices.get(modality_name, {})
+                    for field_name, field_value in modality_values.items():
+                        indices = field_indices.get(field_name)
+                        if not isinstance(indices, torch.Tensor):
+                            local_values[field_name] = field_value
+                            continue
+                        selected = (indices >= sample_start) & (indices < sample_end)
+                        if isinstance(field_value, torch.Tensor):
+                            local_values[field_name] = field_value[selected]
+                        elif isinstance(field_value, list):
+                            local_values[field_name] = [
+                                item for item, keep in zip(field_value, selected.tolist()) if keep
+                            ]
+                        else:
+                            local_values[field_name] = field_value
+                    if any(
+                        not isinstance(field_value, (torch.Tensor, list)) or len(field_value) > 0
+                        for field_value in local_values.values()
+                    ):
+                        local_modalities[modality_name] = local_values
+                sliced[key] = local_modalities
             # Patch-packed visual encoder inputs use dim 0 for different units
             # across fields (patches for hidden_states, images for grid_thw), so
             # they need joint slicing instead of normal recursive tensor slicing.
-            if _is_patch_packed_visual_dict(value):
+            elif _is_patch_packed_visual_dict(value):
                 sliced[key] = _slice_patch_packed_visual_dict(value, dp_rank, dp_size)
             else:
                 # Recurse into nested dicts (e.g. modality_inputs)

tests/unit_tests/data/megatron_mimo/test_dp_utils.py

@@ -5,6 +5,7 @@
 import torch
 import torch.distributed as dist
 
+from megatron.bridge.data.megatron_mimo.collate import megatron_mimo_collate_fn
 from megatron.bridge.data.megatron_mimo.dp_utils import (
     get_megatron_mimo_dp_info,
     get_megatron_mimo_sampling_info,
@@ -230,6 +231,42 @@ def test_recurses_into_nested_dicts(self):
         assert sliced["tokens"].shape[0] == 2
         assert sliced["modality_inputs"]["vision"]["pixel_values"].shape[0] == 2
 
+    def test_preserves_sparse_modality_sample_ownership(self):
+        """Compacted modality inputs must stay with their owning sample DP shard."""
+
+        def make_sample(sample_id, *, has_vision):
+            input_ids = torch.tensor([sample_id, sample_id])
+            modalities = {}
+            if has_vision:
+                modalities = {"vision": {"pixel_values": torch.full((1,), sample_id)}}
+            return {
+                "input_ids": input_ids,
+                "labels": input_ids.clone(),
+                "loss_mask": torch.ones(2),
+                "attention_mask": torch.ones(2),
+                "position_ids": torch.arange(2),
+                "modality_inputs": modalities,
+            }
+
+        batch = megatron_mimo_collate_fn(
+            [
+                make_sample(0, has_vision=False),
+                make_sample(1, has_vision=False),
+                make_sample(2, has_vision=True),
+                make_sample(3, has_vision=True),
+            ],
+            modality_names=["vision"],
+        )
+
+        text_only_shard = slice_batch_for_megatron_mimo(batch, dp_rank=0, dp_size=2)
+        vision_shard = slice_batch_for_megatron_mimo(batch, dp_rank=1, dp_size=2)
+
+        assert text_only_shard["modality_inputs"] == {}
+        torch.testing.assert_close(
+            vision_shard["modality_inputs"]["vision"]["pixel_values"],
+            torch.tensor([[2], [3]]),
+        )
+
     def test_preserves_non_tensor_values(self):
         batch = {
             "tokens": torch.randn(4, 10),