tree_map segmentation fault

[Sumching]

I have two machines, A and B, both equipped with 4090 GPUs, and they share a common hard drive. My uv virtual environment and code are stored on this shared drive. The code is as follows:

from jax.sharding import Mesh, PartitionSpec as P
from jax.experimental.shard_map import shard_map
import jax
import jax.numpy as jnp
from functools import partial

pos_shardings = jax.sharding.PositionalSharding(jax.devices())
k = jax.random.split(jax.random.PRNGKey(123), 2)
p = jax.jit(partial(jax.random.uniform, shape=(8,)), out_shardings=pos_shardings.replicate())(k[0])
w = jax.jit(partial(jax.random.uniform, shape=(128,)), out_shardings=pos_shardings)(k[1])

do_select = True
def scan_eval(f, w_batched, p):
    # import pdb;pdb.set_trace()
    x0 = jax.tree_map(lambda x: x[0], w_batched)
    carry_init = True, jnp.zeros_like(jax.eval_shape(f, x0, p))

    def f_(carry, x):
        is_first, y_carry = carry
        y = f(x, p)
        if do_select:
            y_reduce = jax.lax.select(is_first, y, y_carry + y)
        else:
            y_reduce = y_carry + y
        return (False, y_reduce), None

    (_, res), _1 = jax.lax.scan(f_, carry_init, w_batched, unroll=1)
    return res

def f(w, p):
    return p + w.sum()

def test(w, p):
    w_batched = w.reshape(16, -1)
    res = scan_eval(f, w_batched, p)
    return res

print('t', test(w, p).sum()) # works
print('tj', jax.jit(test)(w, p).sum())

Machine A runs fine with a single GPU, but encounters a segmentation fault on the line x0 = jax.tree_map(lambda x: x[0], w_batched) when using more than 2 GPUs. Machine B, however, works well with both single and multiple GPUs.

[Sumching]

Moreover, I use environment variables to specify my CUDA version, and CUDA is also installed on the shared drive, so both machines use the same environment variables. The JAX version is as follows:

jax 0.5.0
jax-cuda12-pjrt 0.5.0
jax-cuda12-plugin 0.5.0
jaxlib 0.5.0
jaxtyping 0.2.36

What could be the possible reasons?

[superbobry]

Hi @Sumching, I could not reproduce the crash you are observing with 2 GPUs with the nightly JAX version. Do you think you could try the nightly or the latest release (0.6.0) and let me know if the crash is still there?