Parallelization runs computation concurrently across multiple devices. In JAX, it is useful when independent portions of work can be processed simultaneously, including large-scale neural-network training and scientific simulations. section 7.0
The most direct JAX mechanism is pmap(), or parallel map. It applies the same program to data chunks on multiple devices using single-program multiple-data (SPMD) execution. Adding devices increases throughput, although communication overhead usually makes scaling sublinear. section 7.0
Control parallel execution with:
- in_axes, which selects the input dimension mapped across devices. section 7.5
- out_axes, which controls the layout of mapped outputs. section 7.5
- axis_name and optionally axis_index_groups, which enable collective communication across all devices or device groups. section 7.5
- Array shapes whose mapped dimension does not exceed jax.local_device_count(); inputs must be reshaped when necessary. section 7.5
Because JAX transformations are composable, pmap() can be nested or combined with vmap() and grad(). It also supports data-parallel neural-network training. In multihost configurations, each process handles local devices and data while collective operations communicate across global devices. section 7.5 section 1.1.2
Tensor sharding provides a more implicit alternative: shard arrays first, then leave computation code unchanged while the compiler partitions work across devices. Earlier mechanisms such as xmap() and pjit() are retained mainly for historical or legacy-code context. section 7.0 section 8.1
Parallel dot-product calculations
A distributed-array example defines a dot(v1, v2) function with jnp.vdot, creates large vectors, and shards the arrays before computation. The dot-product code remains unchanged while the calculation proceeds in parallel; array-sharding visualization reveals tensor placement. section 8.1
Listing 8.1 Parallelizing dot product calculations using distributed arrays
from jax.experimental import mesh_utils #1
from jax.sharding import PositionalSharding #1
from jax import random
def dot(v1, v2): #2
return jnp.vdot(v1, v2)
rng_key = random.PRNGKey(42)
vs = random.normal(rng_key, shape=(8_000,10_000)) #3
v1s = vs[:4_000,:] #3
v2s = vs[4_000:,:] #3
v1s.shape, v2s.shape
>>> (4000, 10000), (4000, 10000))
jax.debug.visualize_array_sharding(v1s) #4
>>> ┌────────────────────────────────────────────────────────────┐
>>> │ │
>>> │ │
>>> │ │
>>> │ │
>>> │ TPU 0 │
>>> │ │
>>> │ │
>>> │ │
>>> │ │
>>> └────────────────────────────────────────────────────────────┘
sharding = PositionalSharding(mesh_utils.create_device_mesh((8,1))) #5
sharding
>>> PositionalSharding([[{TPU 0}] #6
>>> [{TPU 1}] #6
>>> [{TPU 2}] #6
>>> [{TPU 3}] #6
>>> [{TPU 6}] #6
>>> [{TPU 7}] #6
>>> [{TPU 4}] #6
>>> [{TPU 5}]]) #6
v1sp = jax.device_put(v1s, sharding) #7
v2sp = jax.device_put(v2s, sharding) #7
type(v1sp)
>>> jaxlib.xla_extension.ArrayImpl #8
jax.debug.visualize_array_sharding(v1sp) #9
>>> ┌────────────────────────────────────────────────────────────┐
>>> │ TPU 0 │
>>> ├────────────────────────────────────────────────────────────┤
>>> │ TPU 1 │
>>> ├────────────────────────────────────────────────────────────┤
>>> │ TPU 2 │
>>> ├────────────────────────────────────────────────────────────┤
>>> │ TPU 3 │
>>> ├────────────────────────────────────────────────────────────┤
>>> │ TPU 6 │
>>> ├────────────────────────────────────────────────────────────┤
>>> │ TPU 7 │
>>> ├────────────────────────────────────────────────────────────┤
>>> │ TPU 4 │
>>> ├────────────────────────────────────────────────────────────┤
>>> │ TPU 5 │
>>> └────────────────────────────────────────────────────────────┘
d = jax.vmap(dot)(v1sp, v2sp) #10
d.shape
>>> (4000,) #11
jax.debug.visualize_array_sharding(d) #12
>>> ┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐
>>> │ TPU 0 │ TPU 1 │ TPU 2 │ TPU 3 │ TPU 6 │ TPU 7 │ TPU 4 │ TPU 5 │
>>> └───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘
#1 Imports modules for sharding
#2 Our familiar function for calculating a dot product between two vectors
#3 This time, we generate vectors wider than before.
#4 Visualizes sharding; the whole tensor resides on a single device
#5 Creates 2D mesh to shard the rank-2 tensor’s first axis across eight TPUs
#6 create_device_mesh returns the most performant ordering of devices given the shape.
#7 Shards tensors across their first axis
#8 Checks the type
#9 Checks that the tensor is now sharded across multiple devices
#10 Performs the computation
#11 The shape is as expected.
#12 The result is also distributed across many devices.
Inspect the distributed-array listing to see array creation, slicing, dot-product definition, and sharding visualization. section 8.1
