perf(data-plane): avoid local batch copies

[rohitrango]

What does this PR do ?

Avoids unnecessary copies when fetching full or subset batches from the local data plane. It also preserves PackedTensor physical-segment sharing across replica broadcasts when that sharing reaches the worker intact, including the local data-plane path. The TQ storage codec remains unchanged.

Packed payloads are coalesced and moved to the collective device one field at a time, which bounds GPU staging memory for large multimodal batches.

Adds coverage that verifies full local fetches share tensor storage and replica broadcasts preserve PackedTensor row mappings, sharing metadata, preprocessing configuration, and physical tensors. A two-rank NCCL test covers CPU-to-GPU staging, CPU restoration, and int16 transport.

This PR is stacked on #4082 via rohit/packedtensor_padding.

Issues

None.

Usage

No user-facing API changes.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? The replica-broadcast scope is documented above.

Additional Information

Validated on Slurm job 18295320 (pool0-01445, 8xH100):

uv run --frozen --group test pytest --confcutdir=tests/unit/data_plane -q tests/unit/data_plane/test_leader_broadcast.py tests/unit/data_plane/test_local_sft.py
uv run --frozen --group test pytest --confcutdir=tests/unit/data -m "not mcore" -q tests/unit/data/test_multimodal_dict.py

Results: 17 data-plane tests passed, including the two-rank NCCL test; 59 PackedTensor tests passed. Ruff and Pyrefly passed for the changed implementation files.

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[rohitrango]

/ok to test 33d6b8e

[rohitrango]

/ok to test b7663d1

[rohitrango]

/ok to test 696bd95

[rohitrango]

Review of the three commits on this branch (33d6b8e6 → b7663d13 → 696bd95c)

No functional correctness bug found. I specifically cleared: leader/follower symmetry of the if tensor.numel(): guard, the dtype/device string round-trips, torch.Size(shape).numel() vs the numel *= dimension loop (they agree, including 0-d and 0-size shapes), narrow().view() contiguity, and the all-None segment case. The inline comments cover one perf regression, one over-broad comment, one maintainability footgun, and one test gap.

Two things these commits quietly fix — worth adding to the commit messages

• 696bd95c removes a follower-only crash. The old empty_packed branch built PackedTensor([None] * n_rows, ...), and for a zero-logical-row field (minted by PackedTensor.empty_rows_like(v, 0) at batched_data_dict.py:840, giving tensors=[]) to_wire() returns (None, []) — so the descriptor carried n_rows=0 and every follower hit AssertionError: Input tensors to PackedTensor must be a non-empty list while the leader sailed through. The new path bypasses __init__ and reproduces len() == 0 correctly.
• It also removes a leader/follower layout divergence. from_wire always sets _row_offsets / _segment_indices, so for a legacy-layout value the leader had _row_offsets is None and len(.tensors) == 3 while followers got _row_offsets=[0,1,2,2] and len(.tensors) == 2. len(), as_tensor() and logical_segment_counts_by_row() all agreed — which is all the existing test checks — but .tensors[i] indexing diverged. All ranks now carry the leader's exact layout.

One candidate I chased and dropped

33d6b8e6's identity short-circuit means the local adapter hands out the same PackedTensor object the partition holds, and both PackedTensor.to() and as_tensor(device=...) mutate self.tensors in place — so training looked like it could silently move the stored partition to GPU.

The chain breaks at the consumer. SFTv2 is the only local-data-plane user and is Megatron-only (sft_v2.py:400), and the Megatron worker never device-moves the full fetched batch — the only .to("cuda") is on microbatches (megatron/data.py:144), which always come from PackedTensor.slice and are therefore fresh wrappers. Not a bug today.

The residual is worth a line of comment somewhere: the invariant "never device-move the fetched batch in place" is now load-bearing, undocumented, and untested. Wiring the local plane to a DTensor worker (which does data.to("cuda") on the full batch) would make it bite.

---

Caveat: reviewed by reading only — I could not execute the test suite in this environment.

[rohitrango]

/ok to test a609bff