pmap + grad: Distributed Training Pattern

Let's put together the complete pattern for data-parallel distributed training using pmap.

import jax
import jax.numpy as jnp

# ============================================
# Model and loss (same as single-device)
# ============================================
def predict(params, x):
    w, b = params
    return x @ w + b

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

# ============================================
# Distributed training step
# ============================================
def train_step(params, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    # Synchronize gradients across devices
    grads = jax.lax.pmean(grads, axis_name='batch')
    loss = jax.lax.pmean(loss, axis_name='batch')
    # Apply update
    w, b = params
    gw, gb = grads
    lr = 0.01
    new_params = (w - lr * gw, b - lr * gb)
    return new_params, loss

# Parallelize
p_train_step = jax.pmap(train_step, axis_name='batch')

# ============================================
# Setup: replicate params, shard data
# ============================================
n_devices = jax.device_count()

# Initialize params
key = jax.random.PRNGKey(0)
k1, k2 = jax.random.split(key)
params = (
    jax.random.normal(k1, (4, 1)) * 0.1,
    jnp.zeros((1,))
)

# REPLICATE params: copy to all devices
# Shape goes from (4, 1) to (n_devices, 4, 1)
replicated_params = jax.tree.map(
    lambda x: jnp.stack([x] * n_devices),
    params
)

# SHARD data: split across devices
total_batch = 256
per_device = total_batch // n_devices
X = jax.random.normal(jax.random.PRNGKey(1), (total_batch, 4))
y = jax.random.normal(jax.random.PRNGKey(2), (total_batch, 1))

X_sharded = X.reshape(n_devices, per_device, 4)
y_sharded = y.reshape(n_devices, per_device, 1)

# ============================================
# Training loop
# ============================================
for step in range(100):
    replicated_params, losses = p_train_step(
        replicated_params, X_sharded, y_sharded
    )
    if step % 20 == 0:
        # Loss is the same on all devices; take the first
        print(f"Step {step}: Loss = {losses[0]:.6f}")