What's the difference between xmap, shmap, pmap (pjit?) and which should I use?

[Artur-Galstyan]

Hello everyone, I'm a bit confused about the many possibilities about using JAX on multiple devices.

Can someone please explain the differences between xmap, shmap and pjit and which I should use in the future? It says in the docs that xmap is going to be deprecated? I thought xmap is the more general case of shmap?

Thanks everyone :)

[jakevdp]

Hi - this is a good question, and things are admittedly a bit confusing now because they're in flux and still not well documented.

The TL;DR is you should use shard_map for explicit parallelism, and jit for automatic parallelism; xmap and pmap are deprecated, and pjit is now part of jit.

Some details:

• pmap is the oldest and least flexible parallelizing transformation. It is severely limited: for example, you can only map over axes with the same shape as the number of devices, and it can't be easily nested. It's mostly replaced by shard_map, though it may live on as a convenient wrapper of shard_map.
• xmap is a slightly-less-old attempt to generalize vmap and pmap, but it has mostly been subsumed by shard_map and will be deprecated in the next JAX release (see Add a deprecation warning to the xmap function and the corresponding tutorial #20182). Most existing uses can be replaced with some combination of shard_map and vmap.
• pjit as an API is gone, but its core features are now part of standard jit: when you use jax.jit, you are actually using what was originally called pjit. It is the preferred method for automatic parallelism.
• shard_map has been experimental, but its API has stabilized and it will soon graduate out of jax.experimental in the sign that it is a recommended API (see [shard-map] move shmap out of experimental, add docstring #20124). It is the preferred method for explicit parallelization going forward.

The documentation for implicit parallelism via jit and explicit parallelism via shard_map is still somewhat lacking, but we're actively working on it. The best discussion currently of parallelism via jit is probably this: https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html. The best discussion of shard_map currently is this: https://jax.readthedocs.io/en/latest/notebooks/shard_map.html. We'll soon have a more unified tutorial that introduces both and shows how to use them effectively.

[johannahaffner]

Looking forward to the unified tutorial!

Are you going to benchmark explicit parallelism with shard_map against automatic parallelism with jit? And maybe include guidance on when to choose one over the other?

I’m quite new to JAX and it’s unclear to me what I should prefer for which use-case (for performance/efficiency).

[jakevdp]

Parallelism via shard_map and via jit lower to the same kind of operation; the difference is that in the former the user is in control, while in the latter the compiler is in control. If you have reason to think you can do better than the compiler, then use shard_map. If not, then use jit.

[johannahaffner]

Thank you for the clear response! That makes a lot of sense. So I pick shard_map when I see cores or devices sitting idle, and otherwise trust the compiler.

[arnoegw]

Thank you, @jakevdp , for the superb clarification!

In the meantime, jax.smap has arrived. Is that the "convenient wrapper of shard_map" as which pmap lives on, just with the added benefit of allowing to nest?

[jakevdp]

Hi - pmap has not yet been completely merged with shard_map (the pmap_shmap_merge configuration optionally enables this). smap is basically a simplified wrapper of shard_map that applies to a single axis.