how to use jax.distributed with multiple CPUs

[SingletC]

Hi, I tested the official tutorial for running distributed JAX with multiple GPUs. I am using slurm, so I assume jax.distributed.initialize takes no args.

import jax
jax.distributed.initialize()  
print("jax.devices()", jax.local_devices())

it works fine with an output of

jax.devices() [StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0)]
jax.devices() [StreamExecutorGpuDevice(id=1, process_index=1, slice_index=0)]

However, the same code seems to not work for multiple CPUs across nodes.

jax.devices() [CpuDevice(id=0)]
jax.devices() [CpuDevice(id=0)]

Does anyone knows if this is not implemented feature or if I am using it wrong

[hawkinsp]

This is not implemented for CPU at the moment. However, you may be able to use https://github.com/mpi4jax/mpi4jax as a workaround.

[sandipanpanda]

I am working on integrating JAX with Kubeflow training operator to run distributed training and fine-tuning jobs on Kubernetes in CPU env. https://www.kubeflow.org/events/gsoc-2024/#project-5-support-distributed-jax-for-training-operator

@hawkinsp @apaszke is this supported now?

[PhilipVinc]

It is, though hidden. Recently @inailuig implemented an MPI backend for jax distributed computations.

You need to build MPITrampoline and launch jax through MPI. See the brief instructions other there https://netket.readthedocs.io/en/latest/docs/parallelization.html#mpi-cpu

[hawkinsp]

We also support gloo as a distributed cpu communication framework, not just MPI.

[PhilipVinc]

True, but GLOO is a factor of 10-20 slower than MPI according to some benchmarks I had done in the past.

Depending on your workload this might be a problem.

[sandipanpanda]

Thank you @PhilipVinc!

I running into the following error while running the example:

$ python
Python 3.12.4 (main, Jun 20 2024, 21:56:28) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import os
>>> os.environ['MPITRAMPOLINE_LIB'] = '/home/user/kbflow/MPIwrapper/build/libmpiwrapper.so'
>>> import jax
>>> jax.config.update('jax_cpu_collectives_implementation', 'mpi')
>>> jax.distributed.initialize()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/user/.pyenv/versions/3.12.4/lib/python3.12/site-packages/jax/_src/distributed.py", line 199, in initialize
    global_state.initialize(coordinator_address, num_processes, process_id,
  File "/home/user/.pyenv/versions/3.12.4/lib/python3.12/site-packages/jax/_src/distributed.py", line 62, in initialize
    raise ValueError('coordinator_address should be defined.')
ValueError: coordinator_address should be defined.
>>>

Can you please help me with how it can be resolved?

[hawkinsp]

It sounds like we weren't able to autodetect your job structure. See the docs:
https://jax.readthedocs.io/en/latest/_autosummary/jax.distributed.initialize.html

You need to pass the arguments manually if we can't autodetect.

[PhilipVinc]

I think 127:0:0:1:9988 should do the job. You will also have to declare number of nodes and the local rank of every runner.

I think if you use openmpi as the backend of MPITrampoline it should work out of the box.

[sandipanpanda]

@PhilipVinc I am running into the error even with open-mpi installed

$ mpirun -np 2 python3 -m netket.tools.check_mpi
mpi4py_available             : True
mpi4jax_available            : True
available_cpus (rank 0)      : 2
n_nodes                      : 2
mpi4py | MPI version         : (3, 1)
mpi4py | MPI library_version : Open MPI v5.0.3, package: Open MPI linuxbrew@2b1a6effc901 Distribution, ident: 5.0.3, repo rev: v5.0.3, Apr 08, 2024

$ python
Python 3.12.4 (main, Jun 20 2024, 21:56:28) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import os
>>> os.environ['MPITRAMPOLINE_LIB'] = '/home/osuser/kbflow/MPIwrapper/build/libmpiwrapper.so'
>>> import jax
>>> jax.config.update('jax_cpu_collectives_implementation', 'mpi')
>>> jax.distributed.initialize()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/osuser/.pyenv/versions/3.12.4/lib/python3.12/site-packages/jax/_src/distributed.py", line 199, in initialize
    global_state.initialize(coordinator_address, num_processes, process_id,
  File "/home/osuser/.pyenv/versions/3.12.4/lib/python3.12/site-packages/jax/_src/distributed.py", line 62, in initialize
    raise ValueError('coordinator_address should be defined.')
ValueError: coordinator_address should be defined.
>>> 

[PhilipVinc]

So, by running the command

MPIwrapper/build/bin/mpiwrapperexec -n 2 env | grep SIZE
MPIwrapper/build/bin/mpiwrapperexec -n 2 env | grep RANK

I see that my env variables defining the size and rank are

PMI_SIZE
PMI_RANK

So I write the following script:

> cat test.py
import os

rank = os.environ['PMI_RANK']
size = os.environ['PMI_SIZE']
hostname = "127.0.0.1:9988"

import jax
jax.config.update('jax_cpu_collectives_implementation', 'mpi')
jax.distributed.initialize(hostname, size, rank)
import jax.numpy as jnp

print(f"r{rank}/{size} : ", jax.devices(), jax.local_devices())

and I run with

MPIwrapper/build/bin/mpiwrapperexec -n 2 python test.py