Collective Operations: Communicating Between Devices

In data-parallel training, each device computes gradients on its own data slice. But you need to synchronize these gradients before updating parameters. This is done using collective operations — communication primitives that operate across all devices.

import jax
import jax.numpy as jnp

# psum: sum values across all devices (all-reduce sum)
@jax.pmap
def sum_across_devices(x):
    # x is this device's local value
    total = jax.lax.psum(x, axis_name='devices')
    return total

# Note: pmap needs to know the axis name for collectives
sum_across_devices = jax.pmap(
    lambda x: jax.lax.psum(x, axis_name='i'),
    axis_name='i'
)

# pmean: average across devices
mean_fn = jax.pmap(
    lambda x: jax.lax.pmean(x, axis_name='i'),
    axis_name='i'
)

# pmax: maximum across devices
max_fn = jax.pmap(
    lambda x: jax.lax.pmax(x, axis_name='i'),
    axis_name='i'
)

The All-Reduce Pattern for Gradient Synchronization

This is the core pattern for data-parallel training: each device computes gradients on its data, then all devices average the gradients so they all have the same update.

import jax
import jax.numpy as jnp

def loss_fn(params, x, y):
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y) ** 2)

def train_step(params, x, y, lr):
    """Training step that runs on each device."""
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)

    # Average gradients across all devices
    grads = jax.lax.pmean(grads, axis_name='devices')
    loss = jax.lax.pmean(loss, axis_name='devices')

    # Update (all devices now have the same gradients → same params)
    new_params = params - lr * grads
    return new_params, loss

# Parallelize with pmap
parallel_train_step = jax.pmap(train_step, axis_name='devices')