parallelization in jax when the data are more than the number of devices/cores

[yyang97]

Hi jax lovers and developers,

I am trying parallel computing in jax. However, in my case, I have 100 data but only 8 cores/devices. It looks like jax.pmap() does not support this case.

I just want to know what is the easiest way when I have more data than the devices in the parallellization in jax.

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
import jax


key = jax.random.PRNGKey(0)
N = 100
x = jax.random.normal(key, (N,))


jax.pmap(jax.numpy.square)(x)

Of course it reports the error:

compiling computation that requires 100 logical devices, but only 8 XLA devices are available (num_replicas=100).

Is there any way for me to parallelize the calculation when I have 100 data but only 8 devices in jax? E.g., I want that if a core/device finishes a job and then it automatically starts the next job.

[jakevdp]

You can do something like this with shard_map, which is the more general successor of pmap. Note however that it requires operating on an axis size which is an integer multiple of the number of devices. For example:

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'

import jax
from jax.experimental import mesh_utils
from jax.experimental.shard_map import shard_map
from jax.sharding import Mesh, PartitionSpec as P

mesh = Mesh(mesh_utils.create_device_mesh(8,), axis_names=('i',))

key = jax.random.PRNGKey(0)
N = 800
x = jax.random.normal(key, (N,))

y = shard_map(jax.numpy.square, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)

[ForceBru]

requires operating on an axis size which is an integer multiple of the number of devices

Is it possible to use sharding with any random axis size that's not necessarily an integer multiple of the number of devices?

I have this:

In [60]: sharding = PositionalSharding(jax.devices()); sharding
Out[60]: PositionalSharding([{CPU 0} {CPU 1} {CPU 2} {CPU 3}], shape=(4,))
In [61]: x = jax.random.normal(jax.random.PRNGKey(0), (1024 * 4 + 1,))
In [62]: x.shape[0] / 4 # make sure it's NOT divisible by the number of devices
Out[62]: 1024.25
In [63]: jax.device_put(x, sharding)
...
ValueError: One of device_put args was given the sharding of PositionalSharding([{CPU 0} {CPU 1} {CPU 2} {CPU 3}], shape=(4,)), which implies that the global size of its dimension 0 should be divisible by 4, but it is equal to 4097 (full shape: (4097,))

The error makes total sense, but does it mean it's impossible to shard this array at all? I could put it on a single CPU core, but I'd like to use all 4.

---

A naïve question: is it possible to somehow automatically shard everything such that operations on large arrays automatically run in parallel (at least on the CPU) when possible? TBH, I find this API really complicated. Compared to JAX, Julia's LoopVectorization runs everything in parallel automatically, I don't even need to think about sharding or CPU cores being treated as different devices. I write @turbo and LoopVectorization does the magic. Can JAX do this?