How do I shard jax.ops.segment_sum correctly?

[benmoseley]

I am trying to parallelise jax.ops.segment_sum across multiple GPU devices using sharding:

from functools import partial
import jax
import jax.numpy as jnp
from jax.experimental import mesh_utils
from jax.sharding import PositionalSharding
import numpy as np

@partial(jax.jit, static_argnums=(2,))
def ss(y, x, m):
    s = jax.ops.segment_sum(y, x, indices_are_sorted=False, num_segments=m)
    return s

m,n = int(1024), int(2*1024)
x = jnp.array(np.random.randint(0,m,n))
y = jnp.arange(n, dtype=int)

mesh = mesh_utils.create_device_mesh((2,))
sharding = PositionalSharding(mesh)
x = jax.device_put(x, sharding)
y = jax.device_put(y, sharding)
print(x.shape, x.sharding)
print(y.shape, y.sharding)

s = ss(y, x, m)
print(x.shape, y.shape, (s.shape))

but this gives the following error:

(2048,) PositionalSharding([{GPU 0} {GPU 1}])
(2048,) PositionalSharding([{GPU 0} {GPU 1}])

XlaRuntimeError: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.all_reduce' failed: external/xla/xla/service/gpu/nccl_all_reduce_thunk.cc:244: NCCL operation ncclGroupEnd() failed: unhandled cuda error.: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

Where am I going wrong?

System info:

python=3.10.10
jax=0.4.8
jaxlib=0.4.7+cuda11.cudnn86

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.125.06   Driver Version: 525.125.06   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:A1:00.0 Off |                  N/A |
|  0%   28C    P8    27W / 370W |  18945MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  On   | 00000000:C1:00.0 Off |                  N/A |
|  0%   26C    P8    38W / 370W |  18935MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

[benmoseley]

Update: This works when upgrading to jax=0.4.20
Output:

(2048,) PositionalSharding([{GPU 0} {GPU 1}])
(2048,) PositionalSharding([{GPU 0} {GPU 1}])
(2048,) (2048,) (1024,)