A sharded training step

Derive the global denominator

One shard has two examples with mean gradient $1$; another has one example with gradient $4$. The global mean is $(2+4)/3=2$. Averaging the local means gives $2.5$, overweighting the small shard.

Carry contributions and counts so the final denominator matches the objective. With masked sequence losses, count valid targets rather than automatically counting devices or sequences.

The matching gradient bars are an important check on the current fixture. Add unequal valid counts to expose a weighting rule that balanced shards could hide. Numerical agreement should be tested where the implementation's assumptions are most vulnerable.

Write the objective before splitting the data

Our model predicts $Xw$ with two weights and no bias. The synthetic labels use true weights $[2,-1]$, but the initial weights are zero. The squared-error loss averages eight residuals. Its gradient is $\frac{2}{N}X^\mathsf{T}r$, where $r=Xw-y$ and $N=8$. This host formula is independent of JAX autodiff and supplies a useful reference for the first update.

At $w=0$ the targets are $[2,-1,1,4,-2,3,0,2]$. Squared targets sum to $39$, so initial loss is $39/8=4.875$. $X^\mathsf{T}y=[21,3]$, so the gradient is $[-5.25,-0.75]$. With learning rate $0.1$ the new weights are $[0.525,0.075]$. Check these values before looking at a training curve.

Give batch and parameters different placements

The feature matrix has shape $(8,2)$ and uses P("data",None); each of four devices owns two examples. Labels have shape $(8,)$ and use P("data"), so their shards align with the feature rows. Parameters have shape $(2,)$ and use P(), giving every device both weights.

The input sharding tree supplied to jit has one entry per argument. Outputs are weights, scalar loss and a two-element gradient; all three are replicated in this small example. Replication does not mean four independently trained models. The replicas represent the same global parameters and must receive the same global update.

CPU 0: X rows 0:2, y rows 0:2, full w
CPU 1: X rows 2:4, y rows 2:4, full w
CPU 2: X rows 4:6, y rows 4:6, full w
CPU 3: X rows 6:8, y rows 6:8, full w
partial contributions → global gradient → same w_next

Differentiate the global function rather than inventing a local objective

loss operates on global JAX arrays. value_and_grad evaluates the objective and obtains its parameter derivative together. jit compiles that function with explicit placement contracts; you do not call a Python training loop separately on each device. Combining contributions is necessary because each local batch alone cannot determine the global mean gradient.

This lesson uses automatic global-array parallelization, not manual shard_map or an explicit pmean. The compiler arranges required operations. Do not add an extra collective to an already global reduction merely because the input is sharded. If you later implement a manual local program, you must derive its aggregation and normalization from the same objective.

A correct first update is stronger evidence than a falling curve — on CPU and on a 4-chip TPU VM

A loss curve can decrease even with a normalization bug, especially at a small learning rate. Compare the initial loss, every gradient component and updated weights with NumPy. Then change the batch contents and order: a correct global mean should not depend on which device receives a row. Check replicated weight shards separately.

A ten-step loop is included as a transfer experiment, but this is not a complete resilient training system. It has no data pipeline, optimizer moments, checkpoint or multi-host failure handling. CPU virtual devices demonstrate placement and numerical equivalence; to verify the same sharded train_step across four physical TPU chips (v5litepod-4 or v6e-4), upload the lesson script to your TPU VM and target jax.devices()[:4] with JAX_PLATFORMS=tpu. Float32 reduction order may differ slightly across hardware, so use stated tolerances instead of bitwise comparisons.

gcloud compute tpus tpu-vm scp exercises/distributed-02.py $TPU_NAME:~/jax-tpu-lab/ --zone=$ZONE
gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE \
  --command="sed 's/jax.config.update(\"jax_platforms\", \"cpu\")/# use default TPU backend/; s/jax.devices(\"cpu\")/jax.devices()[:4]/' ~/jax-tpu-lab/distributed-02.py > ~/jax-tpu-lab/distributed-02-tpu.py && JAX_PLATFORMS=tpu ~/jax-tpu-lab/.venv/bin/python ~/jax-tpu-lab/distributed-02-tpu.py"