Behavior of jax.pmap over GPUs with CPU heavy host_callback

[denizokt]

Hi all,

I'm on a multi-GPU device with also a lot of CPU cores. I have a function that does a bunch of computation on the GPU, but then uses host_callback to do expensive CPU operations as well.

I am currently using pmap to map the function across multiple GPUs. I can confirm that the devices are seen by Jax, but judging by the time it takes to call the functions, it seems like parallelism is not helping. As in, pmap over 8 elements takes ~8x time compared to 1 element.

Could this be due to host_callback from all devices only using one CPU somehow? Is there a way I can map e.g. 5 CPUs -> 1 GPU without overlap?

Thank you!

[YouJiacheng]

I think you need to start multiple process. (Maybe using cpu-jitted function is helpful as well.)
host_callback can modify python global variable, I think there is a GIL.
Or you can separate GPU and CPU computation in different jitted function.