how to use pmap in multi-host cpus #10745
-
|
I am able to run the multi-host cpu using jax. This ONE is not the solution? |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 1 reply
-
|
Hi Jiahao It's not an answer, but since I could not manage to use pmap at the time for that use case, we wrote mpi4jax , so you just use standard jit and run the same code everywhere. It requires you to have a working MPI installation but I'm pretty sure at UC you have that... Nowdays there's some experimental support through Global Sharded Devices Arrays but it's not greatly documented. I'll let other comment on how to do that... |
Beta Was this translation helpful? Give feedback.
-
|
thanks @PhilipVinc, appreciated your answers! |
Beta Was this translation helpful? Give feedback.
Hi Jiahao
It's not an answer, but since I could not manage to use pmap at the time for that use case, we wrote mpi4jax , so you just use standard jit and run the same code everywhere. It requires you to have a working MPI installation but I'm pretty sure at UC you have that...
Nowdays there's some experimental support through Global Sharded Devices Arrays but it's not greatly documented. I'll let other comment on how to do that...