Need to train the same neural network faster across several devices? Data parallelism splits each batch into chunks, gives every device an identical model copy, aggregates gradients, and keeps the copies synchronized. section 7.3
In JAX, pmap distributes the update across a device axis. The update computes value_and_grad, combines gradients with lax.psum, and relies on pmap’s internal JIT compilation; grad, vmap, and pmap can work together. section 7.3.2
A practical implementation needs to account for:
- Reshape inputs to expose a device dimension and a per-device batch dimension, such as (8, 32, 784) for eight devices. section 7.3.2
- Handle a final batch that cannot divide evenly across available devices; otherwise reshaping can fail. section 7.3.2
- Decide whether evaluation should also be parallelized; the demonstrated accuracy calculation runs on one device. section 7.3.2
Listing 7.27 Loss and update functions
from functools import partial
INIT_LR = 1.0
DECAY_RATE = 0.95
DECAY_STEPS = 5
NUM_EPOCHS = 20
def loss(params, images, targets): #1
"""Categorical cross entropy loss function."""
logits = batched_predict(params, images)
log_preds = logits - logsumexp(logits)
return -jnp.mean(targets*log_preds)
@partial(jax.pmap,
axis_name='devices',
in_axes=(None, 0, 0, None),
out_axes=(None,0)) #2
def update(params, x, y, epoch_number):
loss_value, grads = value_and_grad(loss)(params, x, y)
grads = [(jax.lax.psum(dw, 'devices'), #3
jax.lax.psum(db, 'devices')) #3
for dw, db in grads] #3
lr = INIT_LR * DECAY_RATE ** (epoch_number / DECAY_STEPS)
return [(w - lr * dw, b - lr * db)
for (w, b), (dw, db) in zip(params, grads)], loss_value
#1 The same loss function as before
#2 Uses pmap to parallelize function and map over the data only; parameters are replicated and not mapped over
#3 Accumulates gradients on each of the devices
Inspect the loss and update functions, including pmap axes and lax.psum gradient aggregation. section 7.3.2
