Exposing the XLA communicator through the FFI interface

[therooler]

A while back I posted a question about getting access to the XLA communicator from JAX' FFI: jax-ml/jax#35969

I never got a response there so I figured I would try my luck here.

The short version is:

I'm interested in calling cuSOLVERMp, NVIDIA's distributed linear algebra library from JAX. I have coded up a cuSOLVERMg interface here https://github.com/flatironinstitute/jaxmg, but I used shared memory to set up barriers between threads and processes. Ideally I would use the underlying XLA communicator to block process and control my program across threads/nodes.

Is it currently possible to access the XLA communicator somehow on the C++ side with JAX FFI?

I would appreciate any insight!

[PatriosTheGreat]

Hi Roeland,
You can use XLA communicator through FFI, though it would require to link statically with XLA since so far all primitives are defined there.

A small example you can find in collective_ops_ffi_test here: https://github.com/openxla/xla/blob/main/xla/backends/gpu/tests/collective_ops_ffi_test.cc#L313

This test runs a simple custom call implementation of all-reduce over 2 GPUs and uses XLA communicator.
The GpuDeviceCommunicator also can be passed to the kernel as an argument (see: https://github.com/openxla/xla/blob/main/xla/backends/gpu/tests/collective_ops_ffi_test.cc#L373). Having it you can use NCCL LSA API from kernel itself to perform device-side rendezvous: https://github.com/openxla/xla/blob/main/xla/backends/gpu/tests/collective_ops_ffi_kernels.cu.cc#L53

[therooler]

Hi Levon,

Thank you for the reply. This is exactly what I was looking for! I will look into replacing my current shared memory hack with this for JAXMg. Thank you for this.

[therooler]

Hi @PatriosTheGreat ,

I just wanted to let you know that we were able to get cuSOLVERMp hooked up to JAX through FFI using the XLA communicator, resulting in a stable code that we are now developing in parallel with JAX and cuSOLVERMp.

We ended up statically linking to XLA, using the docker image JAX uses for compilation to ensure compatibility.

As an example of what we can do: With this new design we were able to do a calculation over 64 H200s across 8 nodes, performing a cholesky solve in JAX of a matrix of size 1.5 million x 1.5 million in ~10 minutes.