Skip to content

Chapter 7: Parallelization - jax.pmap

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: Vectorization (on single device, uses vectorized instructions)
f_vmap = jax.vmap(f)  # Processes batch efficiently

# pmap: Parallelization (across multiple devices)
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: (num_devices, ...)
# Output shape: (num_devices, ...)

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: (2, 5) - 2 devices, 5 elements each
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: matrix multiply
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: 2x2 grid
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: size 3 ≠ 4 devices

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: frequent communication
for step in range(100):
    grads = compute_grads()
    grads = lax.pmean(grads, 'batch')  # Communicate every step

# ✅ Efficient: batch communication
# 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