Frédéric Bastien

@hawkinsp Can you take a look?

Note, an example is being added here: https://github.com/google/jax/pull/20179

Replaced by https://github.com/google/jax/pull/20179

Replaced by https://github.com/google/jax/pull/20179

There is 2 ways of doing multi-GPU on one node. One process that handle all the GPUS, or each GPU have a different process. When doing multi-node, we end up...

@Findus23 do you have a patch to JAX to have it handle this correctly? If so, where it is? It would be good to update JAX to handle this.

@mehdiataei

Yes, it is now fixed. I dumped the buffer allocation file and both at the same and start with: ``` BufferAssignment: allocation 0: size 4000000, parameter 0, shape |f32[100,100,100]| at...

I think it would be better to merge this PR when all is working. Ping us in that case. I know on helios, you can request interactive sessions. Could you...

@tsirif is everything working? Do you have more tests to make sure it work well?