Combine two arrays on different devices

[mehdiataei]

I want to concatenate arrays 'a' and 'b' located on devices 0 and 1, respectively, along a specific axis without transferring the data back to a single device (as it would significantly hinder the performance). I am using the new jax.Array and am unsure if device_put_sharded is a viable solution for this task.

Any thoughts?

Thanks

[mehdiataei]

Another possible solution would be so somehow only access addressable_data in a sharded array and modify it in-place. But I don't think this is possible in JAX currently. You also cannot access addressable_data in a JITed function.

[jakevdp]

If you want to concatenate along the leading axis, effectively creating a sharded array, I suspect the best approach would be using jax.device_put_sharded; something like this:

import jax
import numpy as np

devices = jax.devices()
x = jax.device_put(np.zeros(10), device=devices[0])
y = jax.device_put(np.ones(10), device=devices[1])

xy = jax.device_put_sharded([x, y], devices=devices[:2])

print(xy)
# [[0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
#  [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]]

print(xy.devices())
# [CpuDevice(id=1), CpuDevice(id=0)]

[yashk2810]

I am still unclear on why you don't want the transfer to device 1 happen because it is not going to be expensive at all.

But anyways, you can do this:

In [21]: mesh = jax.sharding.Mesh(np.array(jax.devices()).reshape(2,), axis_names=['x',])

In [22]: s = jax.sharding.NamedSharding(mesh, P('x'))

In [23]: @pjit
    ...: def f():
    ...:     x = jnp.zeros(8)
    ...:     y = jax.device_put(jnp.ones(8), jax.devices()[1])
    ...:     out = jnp.concatenate([x, y])
    ...:     return jax.lax.with_sharding_constraint(out, s)
    ...:

In [24]: f().devices()
Out[24]: [CpuDevice(id=0), CpuDevice(id=1)]

In [25]: f().sharding
Out[25]: OpShardingSharding({devices=[2]0,1})

Here the SPMD partitioner will propagate the sharding backwards and materialize x and y on the right devices to begin with so as to not incur any device transfers.

But again, this seems like overkill to me when you can just keep 1 array uncommitted since transferring it to device 1 is not going to be that expensive anyways.

[mehdiataei]

Thanks for the example, Yash. I had some issues with a similar approach today, but I'll investigate further tomorrow and update you on the results.

As for the overkill, when the Array size is +40GB on each device, transferring is either infeasible, or extremely slow 😮‍💨

[yashk2810]

As for the overkill, when the Array size is +40GB on each device, transferring is either infeasible, or extremely slow 😮‍💨

Ahh, I see. This was the missing part :)

[mehdiataei]

Hey @yashk2810 and @jakevdp:

This is the solution in case you're wondering:

jax.make_array_from_single_device_arrays(shape, sharding, list_of_arrays_on_single_devices)

This is really useful in constructing jax Arrays when the array sizes are very large on each GPU.

[yashk2810]

Ohh cool! If that works for you, then awesome!

I guess I still didn't understand your question or else i would have recommended that :)