Skip to content

← Flax NNX step 84 of 100

Hard Primitives

NNX Data-Parallel Step

Jump to the task

Why this matters

Data-parallel training divides a batch among model replicas. Each replica starts with the same parameters, computes a gradient from its examples, and participates in combining the gradients. All replicas then apply the same update.

This exercise simulates that calculation on one model using a loop over equally sized batch slices. It performs one SGD update after averaging the gradients. It does not create devices, run replicas concurrently, or measure distributed throughput. More devices do not guarantee proportional speedups; communication and other costs affect real training performance.

Why equal shard sizes matter

The loss is mean squared error over every example and output value: jnp.mean((model(x) - y) ** 2). For a local shard of size B with K outputs, the denominator is B * K, not just B.

With equally sized shards, the mean of their losses is the loss over the full batch. Their gradients obey the same averaging rule, provided every shard uses the same parameters. If the shard sizes differ, weight each mean by its number of contributing values. Padding alone is not enough: padded values need a mask and the normalization must count only valid values.

An all-reduce describes communication, not a particular reduction formula: it can sum values or compute a mean. JAX's lax.pmean averages over a bound collective axis. lax.psum sums; summing gradients of local mean losses would scale this exercise's gradient by the number of equal shards unless you divide afterward.

A complete two-shard example

Start with a scalar weight of one and bias of zero. These four examples have a global mean squared error of 7.5:

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

x = jnp.array([1., 2., 3., 4.])
y = jnp.zeros(4)
params = jnp.array([1., 0.])  # weight, bias

def loss(p, inputs, targets):
    return jnp.mean((inputs * p[0] + p[1] - targets) ** 2)

local = [jax.value_and_grad(loss)(params, x[i:i+2], y[i:i+2])
         for i in (0, 2)]
mean_loss = jnp.mean(jnp.stack([value for value, _ in local]))
mean_grad = jnp.mean(jnp.stack([grad for _, grad in local]), axis=0)
updated = params - 0.1 * mean_grad
np.testing.assert_allclose(mean_loss, 7.5)
np.testing.assert_allclose(mean_grad, [15., 5.])
np.testing.assert_allclose(updated, [-0.5, -0.5])
np.testing.assert_allclose(mean_grad, jax.grad(loss)(params, x, y))

Both local gradients use the original parameters. Updating after the first slice would make the second slice use different weights, producing sequential training steps instead of this one batch update.

Model, gradient tree, and optimizer

Create one nnx.Linear(D_in, D_out, rngs=nnx.Rngs(seed)) with default initializers and an enabled bias. Construct nnx.Optimizer(model, optax.sgd(lr), wrt=nnx.Param) once.

For each shard, nnx.value_and_grad(loss_fn)(model, x_shard, y_shard) returns a scalar loss and a parameter-gradient tree. Average matching leaves with jax.tree_util.tree_map; do not flatten or mix kernel and bias gradients. Then call optimizer.update(model, avg_grads) exactly once. Plain SGD subtracts lr times the averaged gradient from both kernel and bias.

No model state changes during these Linear forward passes. Models with BatchNorm, dropout or other state require additional care; shard-local computations need not match one full-batch computation if their statistics or randomness differ.

Problem

Implement data_parallel_step(seed, x_batch, y_batch, num_devices, lr). x_batch has shape (N, D_in) and y_batch has shape (N, D_out). Assume positive dimensions, a positive num_devices dividing N, and a nonnegative learning rate.

  1. Cast seed and num_devices to Python integers. Keep lr as a scalar accepted by Optax; do not force a traced JAX scalar into Python.
  2. Build the Linear and SGD optimizer as described above.
  3. Define the mean squared error loss over all entries of each shard.
  4. Split both batches into num_devices contiguous slices along axis 0, each of length N // num_devices. Use matching slices for x and y, and compute every gradient before changing the model.
  5. Average the scalar losses and separately average the corresponding leaves of the gradient trees.
  6. Apply one optimizer update with the averaged gradient.
  7. Return the pre-update mean loss reshaped to (1,).

The reported value is measured before the update. It therefore does not depend on lr and cannot by itself prove the optimizer ran or that the updated weights are correct. Implement the update anyway; do not return a newly measured post-update loss. The model is fresh on every call, and its updated parameters are not returned.

With num_devices=1, use the full batch in one shard. With num_devices=N, each shard is one example. Both produce the same global mean loss and averaged gradient, up to floating-point rounding. The proposed shard count need not match available devices.

The diagram illustrates batch-axis partitioning on a separate mesh. Its unused model axis replicates the batch shards; it does not show gradient synchronization or the arrays created by this simulation.

Loading visualization…