jit paralel compuation

[raresdolga]

Hi,
I am trying to execute code in parallel with jax.jit and NamedSharding as per: https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html

If I understand correctly, when you shard data on two devices the computation should happen in parallel on each of the devices, but when I look at the shape of the data, it does not seem to be split.

Reproducing code:

from jax.experimental import mesh_utils, pjit
from jax.sharding import Mesh, NamedSharding, PartitionSpec

def ab(data):
    data = data + 1
    jax.debug.print("X is {x}", x=data.shape)
    print(data.shape)
    return data

def main6():
    key = jax.random.PRNGKey(0)
    key, key1, key1b, key2 = jax.random.split(key, 4)
    
    data = jax.random.normal(key1, shape=(100000, 1000))  # Define your input data
    data2 = jax.random.normal(key1b, shape=(100000, 40))  # Define your input data
    num_devices = len(jax.local_devices())

    devices = mesh_utils.create_device_mesh((num_devices, ))
    mesh = Mesh(devices, axis_names=("B",))
    sharding_data = NamedSharding(mesh, PartitionSpec("B",))
    sharded_data = jax.device_put(data, sharding_data)
    sharded_data2 = jax.device_put(data2, sharding_data)

    print(f"Data  shape: {data.shape}")
    print(f"Shard shape: {sharding_data.shard_shape(data.shape)}") 

    f = jax.jit(ab,in_shardings=sharding_data, out_shardings=sharding_data)
    f(sharded_data)

if __name__ == '__main__':
    main6()

Output:

Data  shape: (100000, 1000)
Shard shape: (50000, 1000)
(100000, 1000)
X is (array(100000, dtype=int32), array(1000, dtype=int32))

Desired Output:

Data  shape: (100000, 1000)
Shard shape: (50000, 1000)
(50000, 1000)
X is (array(50000, dtype=int32), array(1000, dtype=int32))

Despite I see that the input data is sharded on the first dimension, with half of the elements on the first device, when the function is called all the elements seem to be passed instead of being split. If I used "shard_map" I get the desired behaviour.

[jakevdp]

Hi - thanks for the question! This question seems predicated on an incorrect understanding of what data.shape means within a jit-compiled computation over sharded data. data.shape is the logical shape of the entire array, regardless of its layout. If you want to inspect the sharding of an array at runtime within a JIT-compiled function, you can use

jax.debug.inspect_array_sharding(data, callback=print)

When I add that to your function, I see this:

PositionalSharding([[{CPU 0}]
                    [{CPU 1}]])

This indicates that the array is indeed sharded as you would expect.

[raresdolga]

Thank you for the reply! That makes sense.
Just to check if my understanding is correct:
For the code above, the function "ab" will be executed in parallel on data splits as specified by the shardings, right?

In a broader context, I am asking this because I am trying to implement DataParallel training for a neural network.

[jakevdp]

Yes, the function will execute in parallel on the sharded data.