A question on manual device transfer and pmap

[sash-a]

Hi there,

If I have manually sharded some data and then try to pmap that data will the pmap move this data again? And if so how could I avoid this extra move and still perform a pmean?

For example:

from functools import partial
import jax
import jax.numpy as jnp


devices = jax.devices()
n_devices = len(devices)
data = [jnp.ones(3) + i for i in range(n_devices)]  # [[0,0,0], [1,1,1],...]
data_tpu = jax.device_put_sharded(data, devices)


@jax.jit
def jitted(x):
    return jnp.mean(x)


@partial(jax.pmap, axis_name="devices") 
def pmapped(x):
    return jax.lax.pmean(x, "devices")


print("data jitted")
jax.debug.visualize_array_sharding(jitted(data_tpu))  # only one 1 device and is just a scalar
print("data pmapped")
jax.debug.visualize_array_sharding(pmapped(data_tpu))  # still on all devices and retains the original shape

[sash-a]

I see that this was mostly an issue with jax.device_put_sharded which does a weak device placement, so the jit moves all the data to device[0]. However, if I place with a jax.device_put and a simple sharding across all devices then jnp.mean does a per device mean and not an all reduce mean, is this possible to do without a pmap or should I not worry about the potentially extra device transfer and use pmap with already sharded data?

[yashk2810]

Seems like I have answered this here? #24279