Development strategies for dealing with backend changes in XLA

[therooler]

Hi,

I'm one of the developers of JAXMg, which is a JAX library that allows users to call cuSOLVERMp routines for distributed linear algebra over multiple GPUs.

JAXMg currently gets access to the underlying NCCL communicator through JAX/XLA’s internals. To ensure compatibility, we compile in the docker image that JAX provides, which pins a specific XLA version.

The problem with this approach is that every new JAX release requires us to recompile and release a new version of JAXMg, since changes in the XLA backend can cause subtle API changes that cause all sorts of problems.

This means the dependencies in our pyproject.toml look like this:

dependencies = [
    "jax==0.11.1; sys_platform == 'linux'",
]

cuda12 = [
    "jax[cuda12]==0.11.1",
    "nvidia-cusolvermp-cu12==0.9.1.9318.post1",
]

So we hard pin to a specific version of JAX per release (although the latest one happens to be backwards compatible with 0.11.1 as well).

JAX	JAXMg
0.10.1	1.0.0, 1.1.0, 1.1.1
0.11.0	1.2.0
0.11.1	1.3.0
0.11.1, 0.11.2	1.4.0 (to be released)

Do you have any advice for how to deal with this? Are we stuck with this type of versioning due to the nature of this project?

[jakevdp]

It sounds like your project is tightly coupled with internal details of XLA, in which case I don't think it's possible to get around that hard version pinning.

There may be one alternative: JAX internally solves this kind of skew issue by bumping some internal counters each time an incompatible change is made – look for uses of ifrt_version (in python code) and JAX_IFRT_VERSION_NUMBER in the C code. If the problematic changes align with bumps in those version numbers, you may be able to depend on those in some cases.

I should emphasize that these version trackers are not meant for public consumption, and there are no stability guarantees around using them – but since you're already depending on internal details of XLA, it could be useful to you.

[therooler]

Hi Jake,

Thanks for the response, this is very useful! We will check the ifrt changes going forward and see if they coincide with the issues we're seeing with XLA changes. Maybe the JAX_IFRT_VERSION_NUMBER will allow us to do some plumbing to stretch what versions we support.

[hawkinsp]

It's also possible we could support some sort of API here. We've previously had requests from users to provide their own NCCL communicator to XLA, this in a sense is the reverse.

[hawkinsp]

Can you elaborate on:

JAXMg currently gets access to the underlying NCCL communicator through JAX/XLA’s internals.

What do you do at the moment?

[therooler]

I see. We really needed the communicator on the C++ side to orchestrate the data layout. cuSOLVERMp requires the data to be in 2D-block cyclic form, and remapping the matrix shards was not really possible with JAX' API (at least not in a memory efficient way). I asked about the best way to do this a while back here, which is what we ended up implementing openxla/xla#42689.

---

On the Python side we use JAX' FFI interface, and register the targets (src/jaxmg/_setup.py):

"cusolvermp_potrs",
        {
            "prepare": "XlaCusolverMpPotrsPrepareFFI",
            "execute": "XlaCusolverMpPotrsFFI",
        },

On the C++ side, we use the XlaCusolverMpPotrsPrepareFFI call to get the CollectiveParams and CollectiveClique objects from XLA required for communication (src/cuda/xla_ffi/handlers.cc):

// Registers the POTRS prepare target that asks XLA to construct the P2P
// communicator clique before runtime.
XLA_FFI_DEFINE_HANDLER_SYMBOL(
    XlaCusolverMpPotrsPrepareFFI, XlaCusolverMpPotrsPrepare,
    ffi::Ffi::BindPrepare()
        .Ctx<ffi::CollectiveParams>()
        .Ctx<ffi::CollectiveCliqueRequests>());

At execution, XlaCusolverMpPotrsFFI gets passed the cached collective_params and collective_cliques,

// Registers the runtime POTRS target that receives padded JAX shards and
// returns the solved right-hand side plus a status vector.
XLA_FFI_DEFINE_HANDLER_SYMBOL(
    XlaCusolverMpPotrsFFI, XlaCusolverMpPotrsDispatch,
    ffi::Ffi::Bind()
        .Ctx<ffi::Stream>()
        .Ctx<ffi::PlatformStream<cudaStream_t>>()
        ... // attributes and data
        .Ctx<ffi::CollectiveParams>() // collective_params
        .Ctx<ffi::CollectiveCliques>()); //collective_cliques

and we then obtain the communicator (src/cuda/cusolvermp_routines/cusolvermp_potrs.cc):

  absl::StatusOr<GpuCommunicator*> gpu_comm = collective_cliques->GetComm(
      *clique_key, collective_params->global_device_id); // Obtain gpu communicator from collective_cliques
  ... 
  void* platform_handle = (*gpu_comm)->platform_comm().handle; // Get raw handle
  ...
  ncclComm_t nccl_comm = reinterpret_cast<ncclComm_t>(platform_handle); // Cast

This communicator is then used for synchronization and data communication across all the devices.