Need to run the same computation on several accelerators at once? pmap() parallelizes a function across local devices using single-program multiple-data (SPMD) execution. It compiles the function with XLA, replicates it, and runs each replica on its own device, so a separate jit() transformation is unnecessary. section 7.0 section 7.1.2
Use pmap() when the work can be divided into independent chunks:
- Prepare identical participating devices and reshape each input so its leading mapping axis contains one chunk per device. section 7.1 section 7.1.1
- Keep the mapped axis no larger than jax.local_device_count(); otherwise compilation fails because JAX cannot provide enough local replicas. section 7.1.2 section 7.5
- Control mapped inputs with in_axes, output placement with out_axes, and cross-device communication with axis_name and collective operations. section 7.2 section 7.5
- Expect scaling to be sublinear when communication overhead offsets the benefit of adding devices. section 7.0
Parallel dot products
A dot product handles one pair of vectors. Split ten million pairs across eight devices, then add vmap() inside pmap() to process each device's local batch: v1sp = v1s.reshape((8, 1_250_000, 3)) v2sp = v2s.reshape((8, 1_250_000, 3)) dot_parallel = jax.pmap(jax.vmap(dot)) result = dot_parallel(v1sp, v2sp) The result has shape (8, 1_250_000). Reshape it to (10_000_000,) to remove the device and local-batch axes; its values match the vmap() result. section 7.1
pmap() also supports data-parallel neural-network training. Each device receives a data shard and a local parameter copy, computes gradients, and uses a collective sum such as lax.psum() so every device applies the same update. Across hosts, each process sees local devices while collective operations communicate across global devices. section 7.3.2 section 7.4
Listing 7.2 Emulating a multidevice system on a CPU
import os
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8' #1
import jax #2
jax.devices("cpu")
WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set
TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0), #3
CpuDevice(id=1), #3
CpuDevice(id=2), #3
CpuDevice(id=3), #3
CpuDevice(id=4),
CpuDevice(id=5),
CpuDevice(id=6),
CpuDevice(id=7)]
#1 Sets the env variable with the desired number of devices before importing JAX
#2 Imports JAX
#3 Eight CPU devices are available.
Emulate multiple CPU devices locally to test pmap() semantics before using a multi-GPU or TPU system. section 7.1.1
