Skip to content

Chapter 7

Overview

jax.pmap parallelizes functions across multiple devices (GPUs, TPUs). This chapter covers distributed computation.

What is pmap?

pmap = parallel map - maps function across multiple devices with automatic device mesh handling.

vmap vs pmap

import jax
import jax.numpy as jnp

def f(x):
 return x**2

# vmap
f_vmap = jax.vmap(f) # Processes batch efficiently

# pmap
f_pmap = jax.pmap(f) # Each device processes its slice

Simple Example

import jax
import jax.numpy as jnp

def f(x):
 return x**2 + 1

f_pmap = jax.pmap(f)

# Input shape
# Output shape

num_devices = jax.device_count()
x = jnp.arange(num_devices * 5).reshape((num_devices, 5))

# Each device processes its row
result = f_pmap(x)
print(result.shape) # (num_devices, 5)

Device Management

Check Available Devices

import jax
import jax.numpy as jnp

# List all devices
devices = jax.devices()
print(f"Devices: {devices}") # [gpu(id=0), gpu(id=1),...]

# Count devices
num_devices = jax.device_count()
print(f"Number of devices: {num_devices}")

# Get device count per type
gpu_count = jax.device_count("gpu")
tpu_count = jax.device_count("tpu")
cpu_count = jax.device_count("cpu")

Place Data on Devices

import jax
import jax.numpy as jnp

x = jnp.array([1., 2., 3.])

# Put on specific device
device = jax.devices()[0]
x_on_device = jax.device_put(x, device)

# Check device
print(x_on_device.device()) # Which device it's on

Basic pmap Usage

Single Function Parallelization

import jax
import jax.numpy as jnp

def f(x):
 return jnp.sum(x) # Reduce within device

f_pmap = jax.pmap(f)

# Input
x = jnp.arange(10).reshape((2, 5))

# Each device processes 5 elements
result = f_pmap(x)
print(result) # [10, 35] (sum on each device)

With Multiple Arguments

import jax
import jax.numpy as jnp

def f(x, y):
 return x @ y

f_pmap = jax.pmap(f)

# Each device
X = jnp.ones((2, 3, 4)) # 2 devices, 3x4 matrices
Y = jnp.ones((2, 4, 5)) # 2 devices, 4x5 matrices

result = f_pmap(X, Y)
print(result.shape) # (2, 3, 5)

Axis Names and Collective Operations

axis_name Parameter

import jax
import jax.numpy as jnp
import jax.lax as lax

def f(x):
 # Reduce across all devices
 total = lax.psum(x, axis_name='batch')
 return total

f_pmap = jax.pmap(f, axis_name='batch')

x = jnp.array([1., 2., 3., 4.]) # 4 devices
result = f_pmap(x)
# Each device sees sum = 10

All-Reduce Operations

import jax.lax as lax

def f(x):
 # Sum across devices
 total = lax.psum(x, axis_name='batch')

 # Average across devices
 mean = lax.pmean(x, axis_name='batch')

 # Max across devices
 max_val = lax.pmax(x, axis_name='batch')

 return total, mean, max_val

Permutation Collective (ppermute)

import jax.lax as lax

def f(x):
 # Rotate values between devices
 # Device i sends to device (i+1) % num_devices
 x_rotated = lax.ppermute(x, axis_name='batch',
 perm=[(i, (i+1) % 4) for i in range(4)])
 return x_rotated

Practical Example: Distributed Training

Data Parallelism

import jax
import jax.numpy as jnp
from jax import grad, vmap
import jax.lax as lax

def loss(params, batch):
 x, y = batch
 pred = params @ x
 return jnp.mean((pred - y)**2)

def train_step(params, batch, lr):
 grads = grad(loss)(params, batch)
 # Average gradients across devices
 grads = lax.pmean(grads, axis_name='batch')
 new_params = params - lr * grads
 return new_params, loss(params, batch)

train_step_pmap = jax.pmap(train_step, axis_name='batch')

# Initialize parameters
params = jnp.ones((2, 3)) # 2 devices, 3 params each

# Prepare batches for each device
X = jnp.ones((2, 10, 3)) # 2 devices, 10 samples, 3 features
y = jnp.ones((2, 10))
batches = (X, y)

# Run training step
params, loss_val = train_step_pmap(params, batches, 0.01)

Device Mesh (Advanced)

Multi-Axis Distribution

import jax
import jax.numpy as jnp
from jax.experimental import mesh_utils
from jax.sharding import Mesh, PartitionSpec as P
import jax.numpy as jnp

# Create device mesh
devices = jax.devices()
mesh = mesh_utils.create_device_mesh((2, 2))

print(mesh) # Shows device grid

Common Gotchas

Gotcha 1: Batch Size Must Equal Device Count

import jax

x = jnp.array([1, 2, 3]) # Size 3
num_devices = 4 # 4 devices

# Size mismatch!
# f_pmap(x) # Error

Gotcha 2: Communication Overhead

import jax
import jax.lax as lax

def f(x):
 # Expensive: communicates across devices
 total = lax.psum(x, axis_name='batch')
 return total

# Use allreduce sparingly - it's slow!

Performance Tips

Minimize Communication

# Inefficient
for step in range(100):
 grads = compute_grads()
 grads = lax.pmean(grads, 'batch') # Communicate every step

# Efficient
# Design to reduce communication patterns

Use Efficient Collective Ops

# psum, pmean, pmax are optimized
# Avoid manual loops over devices

Summary

PMAP QUICK REFERENCE:

Basic parallelization:
 f_pmap = jax.pmap(f)
 result = f_pmap(x_batched_by_devices)

With axis name:
 f_pmap = jax.pmap(f, axis_name='batch')

Collective operations:
 total = lax.psum(x, axis_name='batch')
 mean = lax.pmean(x, axis_name='batch')

Multi-axis (advanced):
 Mesh, PartitionSpec, axis_names

-

Next Steps

Checkpoint: Write distributed dot product across 4 devices

-

Last Updated: 2026-08-09