Unexpected increase in runtime when using the device option in jax.jit()

[eduardolneto]

Hi,

I am trying to integrate a JAX built python package into another non-JAX python package that uses mpi4py for parallelization.

In more detail, I am using an optimisation framework that given a number of cpu cores will run the same ammount of slightly different inputs through a series of calculations. This is done using mpi4py and thus each core will process an input in parallel. Therefore, the output results are available at their respective cpu cores.

After this, I would like to run each of these results through another application which is built in JAX. And ideally would like to do so, by running this JAX application in each of the already separated cpu cores.

Searching for the best approach to do this, I have tried to set the respetive cpu core using the device argument of jax.jit and the mpi4jax communicator rank in the following way:

result=jax.jit(jax_function,device=jax.devices('cpu')[mpi4py.rank_world])(input)

The expected timing of the single core application is ~40s. When using 2 processes I obtain the same timings for each of the processes. However, for 3 or more cpu cores, the timings increase to 600s and above. Futhermore, in these cases (3-8 cores), the application seems to be using all 8 cores of the machine, even if I explicitely chose:

mpirun -n 3 python application.py

with

os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=3'

The main idea using a small example with a diffrax application would be:

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
from jax import config
config.update("jax_enable_x64", True)
import jax
import jax.numpy as jnp
from mpi4py import MPI
import time
import diffrax
import equinox as eqx
comm_world=MPI.COMM_WORLD

class Robertson(eqx.Module):
    k1: float
    k2: float
    k3: float
    def __call__(self, t, y, args):
        f0 = -self.k1 * y[0] + self.k3 * y[1] * y[2]
        f1 = self.k1 * y[0] - self.k2 * y[1] ** 2 - self.k3 * y[1] * y[2]
        f2 = self.k2 * y[1] ** 2
        return jnp.stack([f0, f1, f2])
def main(k1, k2, k3):
    robertson = Robertson(k1, k2, k3)
    terms = diffrax.ODETerm(robertson)
    t0 = 0.0
    t1 = 1000.0
    y0 = jnp.array([200.0, 4.0, 9.0])
    dt0 = 0.0002
    solver = diffrax.Kvaerno5()
    saveat = diffrax.SaveAt(ts=jnp.array([0.0, 1e-4, 1e-3, 1e-2, 1e-1, 1e0, 1e1, 1e3]))
    stepsize_controller = diffrax.PIDController(rtol=1e-10, atol=1e-10)
    sol = diffrax.diffeqsolve(
        terms,
        solver,
        t0,
        t1,
        dt0,
        y0,
        saveat=saveat,
        stepsize_controller=stepsize_controller,
    )
    return sol

start = time.time()
sol = jax.jit(main,device=jax.devices('cpu')[comm_world.Get_rank()])(1.*comm_world.Get_rank(), 1.*comm_world.Get_rank(), 1.*comm_world.Get_rank())
end = time.time()
print(end-start, sol.ys[-1], comm_world.Get_rank())

However, in this small example this idea seems to work for any number of processes (1-8). And the timings of every process are similar to the timing of a single core run (2.0s). I also tried with a third application using adiffrax JAX Vlasov solver which takes 30s to run in single core and there I see again the behaviour I observe for the general application.

So I would assume, the duration of a single core run is important for what I am observing.

Thus, my questions are:
-If anyone is able to tell the cause behind what I am observing?
-If there is a different way to accomplish the linking between JAX and the mpi communicator in these cases, which you could point me to.

Thank you in advance!