How to use JAX pmap with multiple processes?

[leiteg]

Hey there, I would like to know if JAX supports parallelism across multiple processes. The examples of pmap in the docs focus on instances with multiple GPUs/TPUs. I have read here that it's possible to emulate many CPUs by setting an XLA flag, but would it be possible to launch two (or more) processes on different machines and have JAX recognize those as two separate devices?

[skye]

Please check out https://jax.readthedocs.io/en/latest/multi_process.html. Does that answer your question?

[leiteg]

Thanks for replying Skye! This surely helps, although I am not sure how to reproduce the examples on this example. I tried launching two instances of my Python code on my machine but when I print the devices, each process does not seem to be aware that the other exists. Am I missing something?

For example, consider this script:

import jax
import os

pid = os.getpid()
devs = jax.devices()
ldevs = jax.local_devices()

print(f"pid={pid}, devices={devs}, local_devices={ldevs}")

I tried to launch two processes with mpirun -np 2 python3 file.py. Then I get the following output:

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
pid=7318, devices=[CpuDevice(id=0)], local_devices=[CpuDevice(id=0)]
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
pid=7317, devices=[CpuDevice(id=0)], local_devices=[CpuDevice(id=0)]

The two processes do not seem do be cooperating. Do you have any idea how could I make this work? Thanks in advance!

[skye]

Unfortunately we don't support multi-process on CPU. Feel free to open a FR and then we can forward to the XLA team, but I don't think there are any near-term plans to support this.

[leiteg]

All right, thank you!

[PhilipVinc]

Since it seems you are already familiar with MPI...
Just use MPI.
Works on CPU/GPU

https://mpi4jax.readthedocs.io/en/latest/

[leiteg]

Interesting, I will look into it, thanks!