mpi4jax, zero-copy MPI communication for JAX arrays

Collectives that stay inside the JAX runtime

mpi4jax enables zero-copy, multi-host communication of JAX arrays, even from jitted code and from GPU memory.

That sentence carries the whole value proposition, and each clause is doing work. Zero-copy means the data is not serialised into a Python buffer and back out, which for a large array on a GPU is the difference between a collective that costs a copy and one that costs a transfer. Multi-host means it works across nodes, not just across processes on one machine. From jitted code means the collective can be traced and compiled, so it sits inside the JAX program rather than forcing a break out of it. From GPU memory means the array never has to come back to the host to participate in the communication.

The motivation is stated plainly in the README. The JAX framework has great performance for scientific computing workloads, but its multi-host capabilities are still limited. With mpi4jax, you can scale your JAX-based simulations to entire CPU and GPU clusters, without ever leaving `jax.jit`.

That last clause is the interesting constraint, and it is what distinguishes this from the usual answer to multi-host JAX. The usual answer is to hand the array to mpi4py, let MPI do the collective in host memory, and bring the result back into the JAX world. That works, it is simple, and it costs you a device-to-host copy in each direction plus a synchronisation that stops the JAX pipeline. On a multi-node run with a large array, that overhead is usually the difference between the collective being negligible and the collective being the bottleneck.

mpi4jax's approach is to implement the MPI primitives as operations that XLA can lower, so the data stays where it is and the communication is expressed as part of the computation graph. The zero-copy property is a consequence of that design rather than an optimisation applied on top of it.

The project is MIT licensed, the primary language is Python, the documentation is at mpi4jax.readthedocs.io, and the topics name gpu, high-performance-computing, jax, jit, mpi, parallel-computing and xla. The XLA topic is the tell, since XLA is the compiler JAX uses and the layer at which these operations have to be implemented.

The example, and what the second return value is for

The usage example is short enough to take in at a glance, and it shows the shape of the API.

   from mpi4py import MPI
   import jax
   import jax.numpy as jnp
   import mpi4jax

   comm = MPI.COMM_WORLD
   rank = comm.Get_rank()

   @jax.jit
   def foo(arr):
      arr = arr + rank
      arr_sum, _ = mpi4jax.allreduce(arr, op=MPI.SUM, comm=comm)
      return arr_sum

   a = jnp.zeros((3, 3))
   result = foo(a)

The MPI communicator is the one you already have, `MPI.COMM_WORLD`, and it is the same object mpi4py uses, so the communicator setup is not duplicated.

The call itself is `mpi4jax.allreduce(arr, op=MPI.SUM, comm=comm)`, and the operation is an mpi4py enum, `MPI.SUM`. So the vocabulary of MPI is preserved rather than reinvented, which means the documentation for MPI collectives is directly usable.

The second return value is the one that needs a word of explanation. The call returns two things, and the example discards the second with a throwaway name. That second value is the request handle, the object you would use to wait on a non-blocking operation. For a blocking allreduce the handle carries the status of the completed request, and for the asynchronous variants it is what you pass to the wait call. The throwaway name is there because the blocking form has already completed by the time it returns.

The function is decorated with `@jax.jit`, and the collective is called inside it. That is the point of the library, and it is the part that a plain mpi4py call cannot do, because an mpi4py call inside a traced function would be a Python-level side effect that JAX cannot reason about. Here the collective is a primitive in the graph, so it is compiled alongside the arithmetic.

The example also prints only on rank zero, which is the usual pattern for a distributed program, and the documented output for four processes is a 3 by 3 array of sixes. The arithmetic is simple enough to check by hand: each rank adds its own rank to the zeros, and the sum across four ranks is six everywhere, so the allreduce is verified by the output rather than merely asserted.

Installing it, and the GPU path that needs care

The package is on PyPI and on conda-forge, and the plain install is one line.

   $ pip install mpi4jax                     # Pip
   $ conda install -c conda-forge mpi4jax    # conda

The interesting part is that the install has to be adapted to your JAX backend, and the README gives a separate recipe for each.

For the CPU backend, the ordinary install is enough.

   # pip install 'jax[cpu]'
   $ pip install mpi4jax

For the standard CUDA distribution, two extra steps appear: a separate install of nanobind, and the `--no-build-isolation` flag on the mpi4jax install itself.

   # pip install -U 'jax[cuda12]'
   $ pip install nanobind
   $ pip install mpi4jax --no-build-isolation

For the local CUDA distribution, the recipe sets `CUDA_ROOT` instead.

   # pip install -U 'jax[cuda12_local]'
   $ CUDA_ROOT=XXX pip install mpi4jax

The reason for these differences is in the packaging metadata rather than in the README prose. The build requirements are `setuptools>=82.0.1`, `wheel`, `mpi4py>=3.0.1`, `nanobind>=2.0.0` and `jax>=0.6.0`, and the build backend is `setuptools.build_meta`. The extension is written in C++ and bound with nanobind, and it links against both MPI and the JAX libraries. Since pip builds in an isolated environment by default, that environment has its own idea of which JAX and MPI it can see, which is why the flag is needed: with `--no-build-isolation` the build uses the JAX you already installed rather than fetching its own.

The MPI side has the same problem, and the README handles it head on. If your MPI installation is not detected correctly, it can help to install mpi4py separately, and when using a pre-installed mpi4py you must use `--no-build-isolation` when installing mpi4jax.

   # if mpi4py is already installed
   $ pip install cython
   $ pip install mpi4jax --no-build-isolation

Cython is added to that recipe because a pre-installed mpi4py is likely a source build, and the isolated build environment would otherwise not have what it needs to compile against it.

The last piece of the installation story is the compiler. The custom build extension in `setup.py` looks up the MPI compiler wrapper, taking `MPI4JAX_BUILD_MPICC` from the environment and falling back to whatever `mpicc` is on the path, and if neither is found it prints a warning. That environment variable is the one to set on a cluster with a non-default MPI, where `mpicc` on the path may belong to a different MPI than the one your launcher uses.

A compiled extension, and how it is put together

The repository layout shows a hybrid package, and the split explains the installation complexity.

There is a `mpi4jax/` package directory and a `setup.py` that builds a C++ extension. The setup script names the extension module explicitly, with `CPP_SUBMODULE_NAME = "mpi4jax._src.xla_bridge"` and the path `mpi4jax/_src/xla_bridge`. The name is the whole design in one identifier: an XLA bridge, a layer that translates between XLA's custom call mechanism and MPI.

The `setup.py` also shows the dependency detection logic. It tries to import nanobind and sets a flag, tries to import mpi4py and sets another, and the runtime dependency list is built from those: `BASE_DEPENDENCIES = ["mpi4py>=3.0.1", "numpy", f"jax>={JAX_MINIMUM_VERSION}"]` with `JAX_MINIMUM_VERSION = "0.6.0"`.

The versioneer file at the root, plus the `versioneer.py` import and the `MANIFEST.in`, show the standard pattern for deriving a package version from the git tags rather than writing it in a file. The `CITATION.cff` at the root is the machine-readable citation, and it matches the reference the README asks for.

The test setup is in `pyproject.toml`, with the test paths set to `tests`, coverage configured with branch and parallel enabled and the source set to `mpi4jax`, and the version file omitted from coverage. The formatting is Black with a line length of 88 and target versions from py38 through py311, and the dev dependencies pin `black==23.9` and `flake8==3.9.2` exactly, along with pytest, pytest-cov, coverage with the toml extra, pre-commit and tqdm.

The example files give one more hint about the intended use. There is a `examples/shallow_water.py`, which is a classic test case for a communication library, since a shallow water solver is a stencil computation with halo exchange between neighbouring ranks, which is exactly the kind of workload where communication cost and correctness both show up immediately.

The project also has a JOSS paper, which is the marker of a piece of research software that went through peer review, and the README asks you to cite it if you use the library.

Differentiation, and what the library is not

The README mentions one capability that sits apart from the rest, and it is the one that makes this more than a wrapper.

In the spirit of differentiable programming, mpi4jax also supports differentiating through some MPI operations.

The qualifier some is doing the work there, and it is the honest part. A collective that is a sum over ranks has a well-defined derivative with respect to each input, so reverse-mode differentiation through it is meaningful. A collective whose result depends on the collective state of all ranks, such as a reduction that picks an argument based on a comparison, does not have a local derivative at all, and there is no sensible gradient. So the supported set is a subset, and the documentation for the API page lists which operations are supported.

That capability changes what the library is for. A JAX program in which the physics is differentiable is a program in which you can compute gradients of a loss with respect to parameters. If the communication step is inside that loss, and the communication is not differentiable, then the gradient is wrong, and the optimisation silently converges to the wrong thing. That is a specific and serious failure mode, and it is why the differentiability support is the feature that distinguishes mpi4jax from a plain mpi4py call inside a jitted function, which would break the gradient chain entirely.

The other limit is the scope of the operations. The README says allreduce is just one example of the MPI primitives you can use, and points to the API page for the full list. The set is the standard set of MPI collectives and point-to-point operations, but it is a set of primitives, not a distributed programming environment. There is no task scheduler, no data distribution layer, no fault tolerance, and no support for the things a full MPI application gives you. mpi4jax is the communication layer, and the rest of the distributed application is still yours to write.

The practical alternative is mpi4py alone, and the difference is exactly the copy. With mpi4py you convert a JAX array to a NumPy array, or to a buffer, send it, and convert the result back. On a single node that is a device-to-host copy per collective, which for a large array is expensive and for a small array is dominated by the latency of the host round trip. On multiple nodes it is worse, because the data has to travel over the network from a host buffer rather than directly from the device.

So the choice is between a library that keeps your arrays on the device and gives you differentiated collectives, and a library that is pure Python, works with anything, and costs a copy per collective. For a small number of large arrays the first is clearly right. For a workflow with many small collectives, or a codebase that does not use JAX at all, mpi4py is the better answer.

Editorial conclusion

mpi4jax suits a team that already writes JAX for scientific computing and needs to scale a simulation across multiple hosts, especially with GPUs, without reimplementing the collective operations outside the JAX runtime. It does not suit a project that wants a pure-Python install with no compiler, since the package builds a C++ extension against both MPI and JAX. Verify first that the JAX version you run is 0.6.0 or newer, which is the stated minimum, that your MPI installation is detected by the build, and that the `MPI4JAX_BUILD_MPICC` environment variable points at the right compiler wrapper if you are on a cluster with a non-default MPI.

Frequently asked questions

Official sources

1. License: MIT
2. mpi4jax/mpi4jax on GitHub
3. Project website
4. README
5. Releases