How did you solve the problem of jax.pmap hanging?

[SmearingMap]

I have an embarrassingly parallel problem and I seem to be running into the problem you brought up a couple of years ago in jax-ml/jax#5065, which doesn't seem to have been resolved fully in Jax.

I was wondering if you found some useful heuristic workarounds to prevent this and if you wouldn't mind sharing them. It seems you must have, since there seems to be plenty of use of pmap in your code. Did you have to settle for the (generally inadvisable) method of wrapping pmaps in jits, or did you find something better?

[Joshuaalbert]

Hi @halgorthim you can use https://github.com/Joshuaalbert/jaxns/blob/master/jaxns/internals/maps.py#L101.

It behaves just like pmap with a few extra parameters. Use it like this:

from jaxns.internals.maps import chunked_pmap

def embarassingly_parallel_func(*args, **kwargs):
  pass

parallel_func = chunked_pmap(embarassingly_parallel_func, chunksize=...)

results = parallel_func(*args, **kwargs)

#if args[0] is a pytree then you need to specify the batch_size parameter too.
parallel_func = chunked_pmap(embarassingly_parallel_func, chunksize=..., batch_size=...)

chunck_size is slightly misnamed, and means how many parallel workers to use. Much be less that number of devices. batch_size is the leading dimension, which is inferred if the first arg is an array.

Behind the scenes it will pad and reshape arguments and then stick the items in a number of queues that each get executed in parallel with pmap.

Hope this helps!