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 #4106 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.

Previous PR and self-review

Previous PR: #4091. A self-review was performed on the previous PR, and the resulting feedback was addressed in this PR.

[copy-pr-bot]

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

Contributors can view more details about this message here.

[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.

[terrykong]

The main win here is real: deduplicated image segments now survive the replica broadcast. The old path went through to_wire, which emits one row per logical row, so an image shared by several samples crossed the wire once per sample and every receiving rank rebuilt that many copies. broadcast_parts sends the object with its tensors removed, plus each distinct segment once, so the sharing the leader already did is still there on the other side.

Dropping the empty_packed descriptor kind is a second, quieter win worth calling out: the old branch rebuilt followers by hand as PackedTensor([None] * n_rows, ...) and silently lost _row_offsets / _segment_indices / _segment_provenance (base worker_mixin.py:241-251, linked at the base since it is gone at head). That was a real loss of state, not just extra branches.

CI. No lint or unit-test job has run on this head. The always-on gates (DCO, copyright, secrets, semantic PR title, submodule, release) passed, but gh pr checks 4107 shows no Lint check and no L0 unit job. Two things unblock it: the PR has no labels and needs a CI:* one, and someone has to post /ok to test 4fb9fbca541bcd62668651039f907acae4871050 (full SHA, re-posted after every push). Your description says you ran the data-plane and multimodal tests on 8xH100, so this is "please trigger CI", not "this is untested".

What we checked locally. Both changed test files: 16 passed, 1 skipped. The skip is the NCCL case, which needs two CUDA devices — so this is "the CPU suite passes here", not "the tests pass". Ruff is clean via uvx ruff@0.9.9, which is not the same as CI's Lint check (that also runs pyrefly and a whitelist rule).

Blast radius, since the three files differ. worker_mixin.py is shared with the DTensor v2 worker through TQWorkerMixin, so the broadcast rewrite reaches beyond the Energon path — that is the part worth the most scrutiny. local.py does not: the local store is reachable only from SFTv2, which raises unless Megatron is enabled. And multimodal_utils.py is purely additive — one widened import plus two new methods, no existing method touched — so current PackedTensor callers are unaffected.

One comment below, on a test this PR adds. Nothing here blocks the merge.

[terrykong]

/ok to test e9d40e5