NNX Distributed Gradient Reduction Simulation

Why this matters

In data-parallel training, model replicas compute gradients from different parts of a batch. Combining those gradients before an update lets the replicas follow a shared objective. This exercise focuses on the reduction: first sum corresponding gradient arrays, then divide by the number of equally sized shards.

You will implement this with a local loop and one NNX Linear layer. It is a simulation of the gradient arithmetic, not an actual distributed training run. It has no mesh, collective communication, parameter sharding, or optimizer-state sharding.

One update from many shards

Create one model and one SGD optimizer. Every shard must evaluate the same initial kernel and bias. For each slice of the batch, calculate mean squared error and its parameter gradients, saving both without changing the model.

The loss is jnp.mean((model(x) - y) ** 2). For a shard with B examples and K output features, this divides by B * K. Equal shard sizes make the mean of these losses equal to the full-batch loss. The mean of the gradients likewise matches the full-batch gradient, up to floating-point rounding.

This equivalence depends on the model and loss. The Linear here has no batch-dependent state or randomness in its forward pass. Adding BatchNorm, dropout, uneven shard sizes, or masked examples requires reconsidering statistics, randomness, and weighting.

Sum and divide the gradient tree

An NNX gradient tree contains separate kernel and bias arrays. Sum matching leaves across shards, then divide every summed leaf by the shard count:

summed = tree_map(sum corresponding leaves, shard gradients)
average = tree_map(divide by shard count, summed)

Do not add the kernel gradient to the bias gradient. Do not apply SGD until the reduction is complete. For plain SGD, updating with the sum alone multiplies the intended step by the shard count.

This runnable example shows a two-shard reduction with named arrays:

import jax
import jax.numpy as jnp
import numpy as np

gradients = [
    {"kernel": jnp.array([[2., 4.]]), "bias": jnp.array([1., 3.])},
    {"kernel": jnp.array([[6., 8.]]), "bias": jnp.array([5., 7.])},
]
summed = jax.tree_util.tree_map(lambda *parts: sum(parts), *gradients)
average = jax.tree_util.tree_map(lambda array: array / 2, summed)
np.testing.assert_allclose(summed["kernel"], [[8., 12.]])
np.testing.assert_allclose(average["kernel"], [[4., 6.]])
np.testing.assert_allclose(average["bias"], [3., 5.])

In a real manual data-parallel program, lax.psum sums over a bound collective axis and lax.pmean computes its mean. A local Python sum does not perform that communication. Sum-then-divide describes the mathematics, not a guarantee about a communication library's implementation or bitwise reduction order.

Which loss is returned?

Save each shard's loss before the update. After applying SGD, average those saved values and return the result. Computing this average later does not turn it into a post-update loss: the values were measured with the original parameters.

Consequently, changing lr changes the updated weights but not the returned value. The output alone cannot establish that an optimizer update occurred. The task still requires exactly one update of both kernel and bias. A fresh model is constructed for every entrypoint call; updated weights are not returned or reused across calls.