JAX on CPU cluster

[Honza9723]

Dear All,

sorry for beginner question. I would like to ask for some guidance regarding CPU parallelization in JAX.

Background: I am working on so-called planner problem from macroeconomics. It is inf-horizon, discrete time stochastic control problem. I am using neural networks for approximating optimal control functions. As a loss function, I am using negative of average discounted reward across 15k periods and 2.5k initial points. Because of that, each evaluation of loss function is painfully slow, even mini-batch gradients (with 150 points) takes a hell lot of time, hence I want to parallelize this thing.

My networks are relatively small (simple FFW networks with 32 neurons in hidden units), so GPU/TPU wouldn't help (I think). While my loss function include long for loops (accumulating discounted reward over long horizon), I need to evaluate this operation (and its gradient) on large number of points, so I hoped that I can use pmap, or ideally vmap/pmap to distribute this task on multiple CPUs on my university's cluster.

So, my question is, should I expect JAX to work out of the box and register those CPUs just after installation, or would It require more work? Is there some tutorial for this?

Any guidance would be very welcomed!

Best,
Honza

[PhilipVinc]

As far as I know, jax does not support this use-case out of the box.
However, have a look at mpi4jax, which allows you to do exactly that through MPI primitives.

Disclaimer: I'm the author of the library.

[Honza9723]

@PhilipVinc Thank you for your response & link to your awesome library! I would like to ask one more thing. So, "vanilla" pmap works just in shared memory (single node), and mpi4jax is needed to parallelize in distributed memory?

Best,
Honza

[dionhaefner]

So, "vanilla" pmap works just in shared memory (single node), and mpi4jax is needed to parallelize in distributed memory?

It's complicated - it depends on the platform.

pmap is most powerful on TPU, where you can do multi-host communication. On GPU, multi-host support is on the horizon, but not yet mature. CPU currently feels like an afterthought, and I haven't heard of any plans to support it (which doesn't mean that there are none).

There is some more context by the JAX devs here: #4129

mpi4jax fully supports multi-host communications on CPU and GPU, but it obviously depends on MPI as an additional dependency (which your university cluster should supply).

[jekbradbury]

Even though your networks are small, running them on large numbers of data points can create the parallelism needed to take advantage of GPUs or TPUs. So I would definitely try running your code on e.g. a GPU colab (making sure to use vmap for running on the data points in parallel).

[Honza9723]

@jekbradbury Thank you! I will definitely try that, vmap looks like a great tool. Is this possible speed up only a "brute-force" thing, in the sense that GPU/TPU has a capacity to run multiple of those computations in parallel, or is there some XLA magic in play?

[lkhphuc]

I have posted similar question, but I think the difference in my case the networks are of different size. Would love to know if you have any tips for this: #7236