Batch per-slice device_put in transfer_arrays_to_host to reduce D2H dispatch overhead

[anveshjayakrishna]

Summary

In replica_slices.transfer_arrays_to_host, the pinned-host transfer path issues one jax.device_put per replica slice (replica_slices.py#L451). Each device_put carries a fixed per-call dispatch cost independent of slice size, so when saving many arrays the dispatch count can dominate the transfer step.

Proposal

jax.device_put accepts a list of arrays plus a matching list of shardings and issues them in a single dispatch. Batching all pinned-host slices into one device_put pays the dispatch cost once. Non-pinned slices keep the existing copy_to_host_async() path. This is behavior-preserving: the same slices/data are transferred with the same replica-parallel semantics, and peak host memory is unchanged (all buffers were already held until the final await). e.g.:

# Current: one device_put per slice
for rslice in rslices:
    data = rslice.data()
    if use_pinned_host_transfer(data.device):
        data = jax.device_put(
            data,
            jax.sharding.SingleDeviceSharding(data.device, memory_kind='pinned_host'),
        )
    else:
        data.copy_to_host_async()

# Proposed: batch all pinned slices into a single device_put
pinned = [s.data() for s in rslices if use_pinned_host_transfer(s.data().device)]
on_host = jax.device_put(
    pinned,
    [jax.sharding.SingleDeviceSharding(d.device, memory_kind='pinned_host') for d in pinned],
)

Question for maintainers

Is a single batched device_put across all pinned slices acceptable, or was the per-slice form chosen for a specific reason (ordering, memory, or backend constraints)?

[cpgaffney1]

Do you have a micro-benchmark demonstrating a performance improvement from the proposed alternative? Happy to merge if so, just make sure to preserve the correct logic under use_pinned_host_transfer vs not use_pinned_host_transfer branches.

[anveshjayakrishna]

Yes. On a 16-node run (~600 arrays, ~5,000 slices), batching collapsed per-slice device_put dispatch ~98% (1285 → 30 calls), byte-identical output. I'll add a standalone micro-benchmark to the PR for the isolated wall-time.

Both branches preserved: only pinned-host slices are gathered into the single batched device_put; non-pinned slices stay on copy_to_host_async() as before.

Working through CLA clearance — PR shortly.