diff --git a/src/megatron/bridge/data/megatron_mimo/collate.py b/src/megatron/bridge/data/megatron_mimo/collate.py index 482f352a8a..2221f6e41b 100644 --- a/src/megatron/bridge/data/megatron_mimo/collate.py +++ b/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,8 +116,9 @@ 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, @@ -118,3 +126,10 @@ def megatron_mimo_collate_fn( "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 diff --git a/src/megatron/bridge/data/megatron_mimo/dp_utils.py b/src/megatron/bridge/data/megatron_mimo/dp_utils.py index 2519c426a1..e6ed0f5b43 100644 --- a/src/megatron/bridge/data/megatron_mimo/dp_utils.py +++ b/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) diff --git a/tests/unit_tests/data/megatron_mimo/test_dp_utils.py b/tests/unit_tests/data/megatron_mimo/test_dp_utils.py index 01d1f69014..f2d0bce568 100644 --- a/tests/unit_tests/data/megatron_mimo/test_dp_utils.py +++ b/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),