Accelerate JAX computations with mpi4py

[jiayingqi]

Hello! I noticed that my JAX function cannot be accelerated by mpi4py but the numpy function works. Here is an example code:

import time
from mpi4py import MPI
import jax
import numpy as np
jax.distributed.initialize(cluster_detection_method="mpi4py")


arr1 = np.random.uniform(size=(300, 500))

@jax.jit  # Comment this line to use numpy computation
def compute(x):
    return arr1 @ x

comm = MPI.COMM_WORLD
global_size = 10000
loc_size = int(global_size/comm.size)
arr2 = np.random.uniform(size=(500, loc_size))

start = time.perf_counter()
jax.block_until_ready(compute(arr2))
end = time.perf_counter()
print(f"Time: {end-start:.3f} (s)")

Results of mpirun -n 1 python3 test.py with @jax.jit:

Time: 0.108 (s)

Results of mpirun -n 4 python3 test.py with @jax.jit:

Time: 0.114 (s)
Time: 0.119 (s)
Time: 0.120 (s)
Time: 0.119 (s)

Results of mpirun -n 1 python3 test.py without @jax.jit:

Time: 0.063 (s)

Results of mpirun -n 4 python3 test.py without @jax.jit:

Time: 0.017 (s)
Time: 0.017 (s)
Time: 0.017 (s)
Time: 0.018 (s)

So I would like to ask if it is possible to also make my JAX function benefit from mpi4py?