Sharding issue with devices on multiple hosts

[Findus23]

Following up to the issue reported in #16788, I now want to distribute an array over the devices and apply some function to them as a simple test.
Following the idea of https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html I created the following simple script that creates a 256x256 matrix and distributes it over the four GPUs and then applies jax.numpy.sqrt to it (as the easiest example of an element-wise function).

import jax
import numpy as np
from jax.experimental import mesh_utils
from jax.experimental.pjit import pjit
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding

jax.distributed.initialize()
jax.config.update("jax_enable_x64", True)

print(jax.devices(), jax.local_devices())
devices = mesh_utils.create_device_mesh((2, 2))

mesh = Mesh(devices, axis_names=('a', 'b'))


@pjit
def func(x):
    return jax.numpy.sqrt(x)


with mesh:
    x = jax.random.uniform(jax.random.PRNGKey(0), (2 ** 8, 2 ** 8))

    y = jax.device_put(x, NamedSharding(mesh, P('a', 'b')))
    z = func(x)
    jax.debug.visualize_array_sharding(z)
    jax.debug.visualize_array_sharding(x)

    result = np.asarray(z)
    start = np.asarray(x)
    assert np.allclose(np.sqrt(start), result)
    print(result.size)

But this fails (when run with a job-script like in #16788) with the following error:

ValueError: device_put's second argument must be a Device or a Sharding which represents addressable devices, but got NamedSharding(mesh={'a': 2, 'b': 2}, spec=PartitionSpec('a', 'b'))

Traceback

Traceback (most recent call last):
  File "/home/fs72085/lwinkler/jax-testing/distributed.py", line 32, in <module>
    y = jax.device_put(x, NamedSharding(mesh, P('a', 'b')))
  File "/gpfs/data/fs72085/lwinkler/venv-jax/lib/python3.9/site-packages/jax/_src/api.py", line 2480, in device_put
    return tree_map(
  File "/gpfs/data/fs72085/lwinkler/venv-jax/lib/python3.9/site-packages/jax/_src/tree_util.py", line 210, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/gpfs/data/fs72085/lwinkler/venv-jax/lib/python3.9/site-packages/jax/_src/tree_util.py", line 210, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/gpfs/data/fs72085/lwinkler/venv-jax/lib/python3.9/site-packages/jax/_src/api.py", line 2481, in <lambda>
    lambda y: dispatch.device_put_p.bind(
  File "/gpfs/data/fs72085/lwinkler/venv-jax/lib/python3.9/site-packages/jax/_src/core.py", line 380, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/gpfs/data/fs72085/lwinkler/venv-jax/lib/python3.9/site-packages/jax/_src/core.py", line 383, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/gpfs/data/fs72085/lwinkler/venv-jax/lib/python3.9/site-packages/jax/_src/core.py", line 815, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/gpfs/data/fs72085/lwinkler/venv-jax/lib/python3.9/site-packages/jax/_src/dispatch.py", line 660, in _device_put_impl
    raise ValueError(
ValueError: device_put's second argument must be a Device or a Sharding which represents addressable devices, but got NamedSharding(mesh={'a': 2, 'b': 2}, spec=PartitionSpec('a', 'b'))

I see that this happens because of
https://github.com/google/jax/blob/7df34779268df8639907185e3baa58432570bcf6/jax/_src/dispatch.py#L659-L662
https://github.com/google/jax/blob/7df34779268df8639907185e3baa58432570bcf6/jax/_src/sharding.py#L105-L110

but I am not sure what I am supposed to do differently.

Does anyone know how to apply the idea of distributed arrays and later shard_map to multi-host setups?

[Findus23]

I found the following function

with jax.spmd_mode('allow_all'):

from here:
https://github.com/google/jax/blob/b8019dcc9e362d222de5845162f0f9e85317827b/jax/_src/config.py#L779-L791
but it seems like it also doesn't help here.

[yashk2810]

If your setup is fully data parallel then you can use this: https://github.com/google/jax/blob/main/jax/_src/array.py#L647-L659

[Findus23]

Thank you. With this I could get a sample script to work:

It would be really great if the error message could be adapted to explain that one first needs to generate a global array.
And https://jax.readthedocs.io/en/latest/multi_process.html could also explain this and/or give a simple example of a function being applied to a global array.

import os

import jax
import numpy as np
from jax.experimental import mesh_utils
from jax.experimental.multihost_utils import process_allgather,host_local_array_to_global_array
from jax.experimental.pjit import pjit
from jax.sharding import Mesh, PartitionSpec as P

# own_jax_setup(cpu_devices=2 * 2)


jax.distributed.initialize()
jax.config.update("jax_enable_x64", True)

print(os.environ.get("CUDA_VISIBLE_DEVICES"))

print(jax.devices(), jax.local_devices())
print("process_index", jax.process_index())
devices = mesh_utils.create_device_mesh((2, 2))

mesh = Mesh(devices, axis_names=('a', 'b'))

size = 15


@pjit
def func(x):
    return jax.numpy.sqrt(x)


with mesh:
    assert jax.process_count() == 2
    local_shape = (2 ** (size - 1), 2 ** size)

    global_shape = (jax.process_count() * local_shape[0],) + local_shape[1:]
    print("global_shape", global_shape)

    if jax.process_index() == 0:
        dataX = jax.numpy.arange(2 ** (size - 1))
    else:
        dataX = jax.numpy.arange(2 ** (size - 1), 2 ** size)

    dataY = jax.numpy.arange(2 ** size)
    local_array = jax.numpy.outer(dataX, dataY)
    print("local_array shape", local_array.shape)
    assert local_array.shape == local_shape

    split = jax.numpy.split(local_array, len(mesh.local_devices), axis=0)
    print(split[0].shape)
    arrays = jax.device_put(
        split,
        mesh.local_devices
    )

    sharding = jax.sharding.NamedSharding(mesh, P(('a', 'b'), ))
    x = jax.make_array_from_single_device_arrays(global_shape, sharding, arrays)

    z = func(x)
    jax.debug.visualize_array_sharding(z)
    jax.debug.visualize_array_sharding(x)

    result = np.asarray(process_allgather(z))
    print(result.shape)
    start = np.asarray(process_allgather(x))
    assert np.allclose(np.sqrt(start), result)
    print(result.size)

[neel04]

@yashk2810 Do you have a minimal example on how we can port out jax code (that uses device_put and sharding for multi-device parallelization) to handle multi-node/host parallelization as well? Where do we need to add the communication collectives? Do we need them at all or does jax handle that automatically?