@partial(jax.jit, static_argnames=('lr', 'mesh'))
def train_step(params, X, y, lr, mesh):
"""One data-parallel step: shard_map makes the pmean collective explicit.
`X`, `y` arrive with the batch sharded across devices (P('data')) and
`params` replicated (P()) -- see `train` below; shard_map hands each device
the full parameter replica and its own batch shard. pcast marks the replica
as this device's own local copy, so the gradient below is the shard's own;
pmean then averages the shard gradients across devices."""
P = jax.sharding.PartitionSpec
def per_device(params, X, y):
def loss_fn(p):
logits = lenet(p, X[0]) # X[0]: strip the size-1 sharded axis
return optax.softmax_cross_entropy_with_integer_labels(
logits, y[0]).mean()
local = jax.lax.pcast(params, 'data', to='varying') # my own copy
grads = jax.grad(loss_fn)(local) # my shard's mean-loss gradient
grads = jax.lax.pmean(grads, 'data') # The allreduce, in one line
return jax.tree.map(lambda p, g: p - lr * g, params, grads)
step = jax.shard_map(per_device, mesh=mesh,
in_specs=(P(), P('data'), P('data')),
out_specs=P())
return step(params, X, y)Three Ways to Split
Data parallel is our subject: simplest, one sync per step, works for any model that fits. Pipeline and tensor parallel wait for the Language Models part.
Data Parallelism by Hand
Split batch → forward/backward per replica → allreduce gradients → identical update. One process; tensors moved explicitly.
Ring Allreduce
Star: hub moves (k-1)N. Ring (reduce-scatter + all-gather): \frac{2(k-1)}{k}N per device — nearly constant, bounded by 2N for any k. The identity that becomes FSDP.