liveBook · Manning

SPMD

As explained in Deep Learning with JAX

A parallel programming model in which the same program runs simultaneously on multiple devices while each device processes different data.

Single-program multiple-data (SPMD) runs the same program on multiple devices while each device processes a different chunk of data. In JAX, pmap() expresses this model: it compiles a function with XLA, replicates it across devices, and executes the replicas in parallel. section 7.0 section 7.1.2

SPMD resembles vmap() because both map a function over array axes, but they parallelize at different levels. vmap() adds a batch dimension to primitive operations, whereas pmap() distributes replicas across devices. Functional purity makes these transformations composable, so a pipeline can combine compilation, vectorization, parallelization, and differentiation. section 7.1.2 section 1.2.2

Use SPMD when the computation naturally divides across devices:

  • Split input data into one leading slice per device; pmap() expects that mapped axis to match the available device count. section 7.1.2
  • Keep the same model or computation on every device when using data parallelism. section 7.3
  • Communicate and aggregate results such as gradients when replicas must update identical model weights. section 7.3
  • Expect sublinear scaling because communication between devices adds overhead. section 7.0

Mapping a dot product across eight devices

For ten million pairs of three-dimensional vectors, reshape each input from (10000000, 3) to (8, 1250000, 3). Applying pmap(dot) produces one result per device with shape (8,), while pmap(vmap(dot)) processes every local batch and produces (8, 1250000). Flattening that result gives (10000000,), and comparing it with the vmap() result returns Array(True). Applying pmap() directly to the original arrays instead requests ten million logical devices and raises a device-count ValueError.

Listing 7.5 Applying vmap and pmap to our arrays
dot_batched = jax.jit(jax.vmap(dot))                            #1
 
x_vmap = dot_batched(v1s, v2s)                                  #2
 
x_vmap.shape                                                    #3
>>> (10000000,)                                                 #3
 
dot_parallel = jax.pmap(dot)                                    #4
 
x_pmap = dot_parallel(v1s,v2s)                                  #5
 
>>> ...
>>> ValueError: compiling computation that requires 10000000 logical devices, 
but only 8 XLA devices are available (num_replicas=10000000)    #6

Inspect the runnable comparison showing why pmap() requires an input mapping axis aligned with the available devices. section 7.1.2

How Deep Learning with JAX covers it

How deep it goes
Explains pmap() compilation, replication, input layout, nesting with vmap(), and communication overhead, but defers tensor sharding.
Examples
Demonstrates vector dot products across eight devices and data-parallel MNIST image-classification training.
What it assumes
Assumes familiarity with mapped functions and arrays while teaching device-aligned sharding, functional purity, and data-parallel replication.
What you can do after
Apply it via JAX pmap() and vmap() to distribute batched computations across devices and verify equivalent results.
How it connects
Connects pmap() with vmap(), XLA, functional purity, data parallelism, model parallelism, and collective operations.
Cover of Deep Learning with JAX
liveBookmanning.com
Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.