Practical Example: Data-Parallel Training Pattern

Let's build a complete data-parallel training example that works on any number of devices. We'll show both the pmap approach and the modern sharding approach for comparison.

Approach 1: pmap-based Data Parallel Training

import jax
import jax.numpy as jnp

# Model: simple 2-layer network
def init_params(key, in_dim, hidden_dim, out_dim):
    k1, k2 = jax.random.split(key)
    return {
        'w1': jax.random.normal(k1, (in_dim, hidden_dim)) * 0.1,
        'b1': jnp.zeros(hidden_dim),
        'w2': jax.random.normal(k2, (hidden_dim, out_dim)) * 0.1,
        'b2': jnp.zeros(out_dim),
    }

def forward(params, x):
    h = jnp.tanh(x @ params['w1'] + params['b1'])
    return h @ params['w2'] + params['b2']

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

# Data-parallel training step
def train_step(params, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    # Synchronize across devices
    grads = jax.lax.pmean(grads, axis_name='i')
    loss = jax.lax.pmean(loss, axis_name='i')
    # SGD update
    params = jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)
    return params, loss

p_train_step = jax.pmap(train_step, axis_name='i')

# Initialize
n_devices = jax.device_count()
params = init_params(jax.random.PRNGKey(0), 10, 32, 1)
replicated_params = jax.tree.map(lambda x: jnp.stack([x] * n_devices), params)

# Generate data (sharded across devices)
key = jax.random.PRNGKey(1)
X = jax.random.normal(key, (n_devices * 64, 10))
y = jax.random.normal(jax.random.PRNGKey(2), (n_devices * 64, 1))
X = X.reshape(n_devices, 64, 10)
y = y.reshape(n_devices, 64, 1)

# Train
for step in range(50):
    replicated_params, losses = p_train_step(replicated_params, X, y)
    if step % 10 == 0:
        print(f"Step {step}: Loss = {losses[0]:.6f}")

Approach 2: Modern Sharding (Recommended for New Code)

import jax
import jax.numpy as jnp

# Create mesh and shardings
mesh = jax.make_mesh((jax.device_count(),), ('dp',))
data_sharding = jax.NamedSharding(mesh, jax.P('dp'))
replicated_sharding = jax.NamedSharding(mesh, jax.P())

# Initialize params (replicated across all devices)
params = init_params(jax.random.PRNGKey(0), 10, 32, 1)
params = jax.device_put(params, replicated_sharding)

# Shard data across the data-parallel axis
X = jax.random.normal(jax.random.PRNGKey(1), (256, 10))
y = jax.random.normal(jax.random.PRNGKey(2), (256, 1))
X = jax.device_put(X, data_sharding)
y = jax.device_put(y, data_sharding)

# Training step — just use jit! The sharding handles distribution.
@jax.jit
def train_step_sharded(params, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    params = jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)
    return params, loss

# Train — XLA handles the all-reduce automatically
for step in range(50):
    params, loss = train_step_sharded(params, X, y)
    if step % 10 == 0:
        print(f"Step {step}: Loss = {loss:.6f}")

Notice how much simpler the sharding version is: no manual reshaping, no explicit pmean, no replication boilerplate. You specify how data and params are distributed, and XLA's compiler figures out the communication.

For production distributed training in the JAX ecosystem, libraries like Flax NNX and training cookbooks provide higher-level abstractions over these sharding primitives, but they all build on the same Mesh + PartitionSpec + NamedSharding foundation.