[RFE] Add support for distributed CPU-backend mode

[PhilipVinc]

Unless I am mistaken, it is only possible to use the distributed backend (initialised with jax.distributed.initialize ) with the GPU and TPU backends.

However, I believe that Tensorflow, thus XLA should also support the CPU backend.
Would it be possible to support it in Jax as well so that it will be possible to use it with pjit & co?

[hawkinsp]

The main thing that JAX is missing to make this work is an implementation of the various collective operations in XLA that works across processes.

Two possibilities are mpi, in which case the third-party package mpi4jax may be of interest.

Another possibility might be to plug in something likegloo into XLA to implement its collectives: https://github.com/facebookincubator/gloo
This would probably not be hard to do. Currently the collectives implemented in XLA/CPU are naive reference implementations.

[PhilipVinc]

(I am the author of mpi4jax so I know that one pretty well 😄 )

Thanks for the explanation, I see the issue now.

The reason I'm asking that is mainly because 'pjit' and distributing along different axes along a mesh with pjit are very interesting to me and I would like to play with it in our packages, but MPI limits us to only to something like data-parallelism.
Unfortunately I have many CPU-based users and that's why I need CPU support.

The other reason is that I'd like to support multiple-GPUs/CPUs per node (which is exactly what pjit/GlobalArray does) but that would mean supporting mpi primitives within pmap which... I'm not sure how to do. Right now I have to force users to launch 1 jax process per GPU but that's particularly annoying in some HPC setups.

[hawkinsp]

Yes, in essence you are asking to replace the collectives emitted by XLA internally with calls to mpi or something similar. This is a bit different to mpi4jax where you added a separate set of collective ops unknown to XLA on the side.

Right now, XLA:CPU emits calls to functions in a small helper library that implement collectives:
https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/compiler/xla/service/cpu/cpu_runtime.h;drc=6eeb889576593a803bce51871b11fb2b27f8f2b3;l=174

Someone would need to teach XLA:CPU how to either call different runtime library functions, or to change those runtime library functions to call MPI (etc.) That would most likely need a bit of refactoring so multiple collective implementations can be plugged in.

[PhilipVinc]

Yes, in essence you are asking to replace the collectives emitted by XLA internally with calls to mpi or something similar.
This is a bit different to mpi4jax where you added a separate set of collective ops unknown to XLA on the side.

Hmm, yes, that would be amazing, but I can imagine that would be quite a bit of work and I'm unsure if google's interested in that. Though surely academic groups working with HPC would be interested and benefit into it.

Someone would need to teach XLA:CPU how to either call different runtime library functions, or to change those runtime library functions to call MPI (etc.) That would most likely need a bit of refactoring so multiple collective implementations can be plugged in

Getting my hands dirty in XLA itself is a bit beyond the amount of time I have available now, unfortunately. Is there anything I can do to help you In the process/convince you that this is an useful path?

[/begin off topic]
However, something that might make me temporarily, slightly happier would also be a way to insert MPI custom calls into pmap-ped functions. Right now we define C functions that respect the XLA calling convention and then specify how to encode those on cpu and gpu but I have no idea how to support pmap in this context. If you have any pointers I'd be happy to take them.
[/end off topic]

[PhilipVinc]

@hawkinsp did anything change on your end about this recently, or you're still not really planning on supporting this-?

[alelovato]

@hawkinsp, are there any updates on this end? I would like to start using JAX on our CPU-based cluster.

[jon-chuang]

Hello, I may be interested in taking this on by implementing gloo-based collective ops, replacing the naive implementations.

How does this sound @PhilipVinc @hawkinsp ?

[jon-chuang]

I will start with all-reduce and broadcast, in that order.

[hawkinsp]

@jon-chuang well that sounds excellent, if you wanted to contribute that! I would agree: start with all-reduce, which is by itself enough for data-parallel training.

[jon-chuang]

Here is the MVP target:

1. Can perform psum on separate processes running on local with XLA CPU runtime; e2e test (specifically, multiprocess_cpu_test.py, similar to gpu test).

[jon-chuang]

@PhilipVinc could you advise on the degree to which we should be able to perform a collective operation across both CPU and GPU (e.g. GPU+CPU offloading).

In this case:

1. CPU<->GPU: MPI (incl. cross-process, same device)
2. CPU<->CPU: MPI
3. GPU<->GPU: NCCL

The way we can implement it is to e.g. for all-reduce:

1. do a local all-reduce first (CPU<->GPU + GPU<->GPU(same host)),
2. then use either GPU<->GPU or CPU<->CPU (cross-host).

I think that GPU<->GPU should have better performance via NCCL?

See also: alpa-projects/alpa#694

Note that I don't think that even torch.distributed allows for hybrid cluster?

EDIT: gloo implements local reduction in CPU memory - see e.g. cuda_allreduce_ring, cuda_allreduce_ring_chunked. The latter leverages NCCL for same-host multi-GPU reduce/scatter.

[PhilipVinc]

@jon-chuang thank you for looking into this. It's something that would greatly benefit many people including me...

As per your question, If I understand your question correctly, you want to know what reduction operations must be implemented?
In XLA, as of today, there exist only CPU-based reductions (so CPU to CPU) or GPU-based reductions (so GPU to GPU).
That's because an XLA compiled executable can only run on one platform.

So you should not worry for CPU-GPU reductions and can always assume that the devices executing your distributed operation are homogeneous.
At least, that assumption has worked very well for mpi4jax so far.

There might be plans to allow for hybrid computations (@hawkinsp will know for sure) but I'd leave that out of scope for a first implementation.

--

I think the only operation you need to implement is CPU-CPU reductions, possibly using MPI. GPU-GPU reductions are already implemented (using NCCL I think).

[jon-chuang]

Actually, I got a hint that the new PJRT runtime can handle a mixed CPU<->GPU workload. Could you confirm @hawkinsp ?