device_put vs device_put_sharded

[sash-a]

So if I have some data and I want to shard it which should I use? It seems the docs for device_put_sharded are not available has it been depreciated?

My big issue is that when slicing into arrays I don't seem to get the same result with both functions here's a small example

import jax
import jax.numpy as jnp

from jax.experimental import mesh_utils
from jax.sharding import Mesh
from jax.sharding import NamedSharding
from jax.sharding import PartitionSpec as P

devices = jax.devices()
n_devices = len(devices)

data = jnp.arange(n_devices)[:, jnp.newaxis].repeat(3, axis=-1)  # [[0,0,0],[1,1,1],...]

device_mesh = mesh_utils.create_device_mesh((n_devices,), devices)
mesh = Mesh(device_mesh, axis_names=('x'))
sharding = NamedSharding(mesh, P('x'))

data_dev_put = jax.device_put(data, sharding)
data_dev_put_sharded = jax.device_put_sharded([data[i] for i in range(n_devices)], devices)

print("data device put all")
jax.debug.visualize_array_sharding(data_dev_put)  # all devices
print("data device put all [0]")
jax.debug.visualize_array_sharding(data_dev_put[0])  # all devices?
print("\n\ndata device put sharded")
jax.debug.visualize_array_sharding(data_dev_put_sharded)  # all devices
print("data device put sharded [0]")
jax.debug.visualize_array_sharding(data_dev_put_sharded[0])  # just device 0

My question is why does data_dev_put[0] seem like it's on all devices? The output is a bit different from the others, only 1 colour, but clearly shows TPU 0,1,2.... Which method is recommended? Particularly interested in the case where I want to subsequently pmap the data from my previous question

Additionally it seems that device_put_sharded is a weak device placement? As when I jit this it moves it to the default device, is this intended?

[yashk2810]

You should use jax.device_put. While not officially deprecated, device_put_sharded is softly deprecated and we plan on removing it soon.

[yashk2810]

To add more on the indexing question:

That's just the semantics for jit vs pmap. That is, when you use PmapSharding, the __getitem__ portion goes via a different way giving you the device local array.

But slicing a NamedSharding array goes via the jit path and gives you a global array back. You can use .addressable_shards to get the shards of a global Array on each device.

Ideally, I would suggest to move away from pmap and use shard_map + jax.jit instead since that is more in line with the new sharding semantics in JAX.

You can read: https://jax.readthedocs.io/en/latest/sharded-computation.html, https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html and https://jax.readthedocs.io/en/latest/notebooks/shard_map.html for more information.

[sash-a]

Hi yashk2810, thanks for the response!

So if I'm understanding you correctly even though the visualisation doesn't quite show that it's on device 0, that is just the semantics of NamedSharding returning a slice of the global array and not the device local array and this is why the visualize doesn't quite show the same thing as with PmapSharding?

With your advice I've got a working example with shard_map and jit, I feel a lot more confident that the data isn't being moved unnecessarily now, because the device_put and the shard_map use the same mesh and sharding. Just want to check that this is indeed the case and shard map is clever enough to know that my array is already sharded across the mesh I've passed in. So e.g here data_tpu wouldn't be moved twice right (assuming sharding uses the same mesh and partition in_spec as the shard_map?

data_tpu = jax.device_put(data, sharding)

@partial(shard_map, mesh=mesh, in_specs=P('x'), out_specs=P('x'))
def shmap(x):
    return jax.lax.pmean(x, 'x')

y = jax.jit(shamp)(data_tpu)

My use case is that some threads are moving the data around and other threads are using it and I'd prefer certain threads to have the overhead of moving the data in order to not block others, which is why I want to do the explicit device_put before the shard_map

[yashk2810]

if data_tpu.sharding.spec == P('x'), then yes, data won't be moved.