[data] fix: preserve sparse MIMO modality DP ownership

[yaoyu-33]

Problem

MegatronMIMO explicitly supports batches where a modality is absent from some
samples. The collator compacts modality inputs by dropping absent entries, but
module-local DP slicing previously split that compacted tensor independently
from the full sample batch.

With four samples whose vision presence is [false, false, true, true] and
vision/language DP 2, DP shard 0 owns the first two text-only samples but
received the vision tensor belonging to sample 2. Equal-DP MCore routing carries
no sample identity to repair this, so training fails with a modality-token versus
embedding-count mismatch. Sparse counts not divisible by DP failed earlier in
the slicer.

Fix

The MIMO collator now records private per-field sample indices only when a field
is sparse. Module-local slicing uses those indices to keep tensor and list values
with their owning sample shard, then removes the metadata before the model call.
All-present batches retain the existing fast path, and DP 1 preserves the
model-facing batch while consuming the private metadata.

Validation

A deterministic source-level reproducer using the real collator, DP slicer, and
pinned MCore was run unchanged before and after the production fix:

PYTHONPATH=3rdparty/Megatron-LM uv run python /tmp/reproduce_sparse_mimo_dp.py
before: exit 1; shard 0 unexpectedly contained pixel_values tensor([[2]])
after:  exit 0

Adjacent contract checks covered all-present DP 2 behavior, a single present
modality, aligned list metadata, and sparse DP 1 cleanup; all passed. The focused
regression is included in test_dp_utils.py.

uv run pre-commit run --all-files
Passed: end-of-file, trailing whitespace, ruff, ruff-format

git diff --check
Passed

Scope: this validates deterministic CPU data ownership and does not claim a full
suite, GPU, distributed end-to-end, convergence, or performance run. The normal
local pytest environment was unavailable because a locked dependency has no
wheel for this workstation platform; CI should execute the included focused unit
test in the project container.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test c84347a

[claude]

Review — sparse MegatronMIMO modality DP slicing

The change tracks per-sample ownership of modality inputs at collate time (_mimo_sample_indices) and uses it in slice_batch_for_megatron_mimo so compacted modality tensors follow their owning sample into the correct DP shard, instead of being naively split by dim 0. This fixes a real correctness bug where, e.g., only samples 2 and 3 carry vision inputs but a 2-way split would hand rank 0 (samples 0-1) the vision data. Logic looks correct and the new test exercises the core case well.

Findings

• Medium — patch-packed encoder dicts bypassed in sparse mode (dp_utils.py:285): when _mimo_sample_indices is set, the new modality_inputs branch fully replaces the _is_patch_packed_visual_dict / recursion branches. A nested patch-packed {hidden_states, grid_thw} dict would hit the final else and be passed through unsliced to every shard. Only matters if a patch-packed modality and a sparse modality can share a batch — please confirm reachability and add coverage if so. (Inline comment posted.)

• Low — collate sample-index correctness depends on per-key membership (collate.py:90-119): sample_indices is rebuilt per key from 'if key in item'. This holds only if every field of a given sample-modality appears together. Fields present on a different subset of samples than the modality's other fields would desync the tensor and its indices. Current datasets appear to emit whole modality dicts per sample, so this holds — worth a brief comment noting the invariant.

Notes

• The 'dp_size == 1 and modality_sample_indices is None' early-return preserves the fast path; when indices are present it proceeds even at dp_size==1, which is intended (start=0, end=global_batch_size selects all).
• _MIMO_SAMPLE_INDICES_KEY is duplicated as a literal in both collate.py and dp_utils.py. Not a bug, but a shared constant would prevent future drift.
• The key is only emitted when some field is sparse (numel != len(batch)), so the dense fast path is unaffected.

Suggested test cases

• tests/unit_tests/data/megatron_mimo/test_dp_utils.py::TestSliceBatch::test_preserves_sparse_modality_sample_ownership (added — the core sparse-ownership case)
• Sparse ownership at dp_size=1 (indices present but no split): all owning samples retained
• Sparse modality co-existing with a patch-packed visual dict in the same batch (verifies patch-packed joint slicing still applies under the new branch)
• Sparse modality with a non-tensor (list) field, verifying the list-filter path
• A modality owned by samples spanning a shard boundary (e.g. samples 1 and 2 with dp_size=2) to confirm the mask splits ownership across shards
• Collate emits _mimo_sample_indices only when a field is sparse, and omits it when every sample carries the modality

No perf tests impacted.

[claude]

When _mimo_sample_indices is present, this branch handles all of modality_inputs and the _is_patch_packed_visual_dict(value) / recursion branches below are never reached for it. A patch-packed encoder dict ({hidden_states, grid_thw}) nested under a modality would then fall into the final else: local_values[field_name] = field_value (it is a dict, not a Tensor/list) and be passed through unsliced — every DP shard would receive the full encoder input.

This only occurs when a patch-packed modality co-exists in a batch that also has a sparse modality (which sets the key). If that combination is reachable in real VLM SFT data, the joint patch-packed slicing should still apply here. Could you confirm whether these two paths can co-occur, and add coverage if so?