using pmap on multiple GPU hosts

[igorvere]

Hello, can anybody share sample code how to start jax on multiple GPU nodes. So pmap can be used for collective operations (says sum up gradients from each node)

Thanks
Igor

[tomhennigan]

Hey @igorvere my understanding is that JAX does not currently have user friendly support for multi-host multi-gpu. There is lightly tested support in the runtime for distributed jobs that use NCCL for multi-host collective operations, but there are no Python APIs to configure this at the moment.

[Joshuaalbert]

Any idea when/if this is going to be in the roadmap?

[mattjj]

I believe multi-host multi-gpu will become a lot easier over the next 6 months. We don't have a hard-and-fast roadmap, but scaling in general is a high priority. (So far most of the action has been Google-internal and on TPUs.)

[lucasliunju]

Hi @mattjj , I would like to ask how to use multi-host tpu, such as v3-32, v3-128. Could you please provide some examples about how to connect jax and multi-host tpu, and the related operations in jax.

Thank you!

[mattjj]

Hey @lucasliunju , thanks for the question! Maybe it'd be best to start a separate thread/discussion for that, since this thread is about multi-host GPU. When you do, tag @skye, as she's the best person to ask!

[khanrc]

Hi @mattjj , can I ask the progress of the multi-node GPU training? I want to use JAX for our project, but multi-node GPU training is necessary for the project.

[PhilipVinc]

For anyone interested, MPI4Jax allows you to run computations on multiple gpu hosts.

[lucasliunju]

Hi @PhilipVinc ,

Did you try it on multi-host tpu. I would like to run my code on multi-host tpu.

Thank you!

Yong

[PhilipVinc]

We never tried. Surely it won’t support direct TPU to TPU communications, but maybe by copying to the CPU host you might get it to work

[igorvere]

Thanks! Will check that out

[igorvere]

we finally used nccl directly (thanks for unsafe pointer interface) to do distributed ops between gpu hosts

I can see there is promising pull request #8364 which perhaps can make pmap available for gpu multihost setup

Although for single host multigpu machine I see that collective operations are prohibitively slow vs direct nccl calls