NNX SGD Baseline for Sharded Training

Why this matters

Distributed training should preserve the intended optimization calculation while changing where data and model state live. A small unsharded training step is useful as a numerical baseline for checking a sharded implementation.

That baseline is what you will write here: one ordinary SGD step on one NNX Linear layer. Despite the historical name fsdp_step, this function does not implement fully sharded data parallelism (FSDP), partition parameters, or run collectives. The num_shards argument is informational and does not affect the calculation.

What fully sharded training adds

Fully sharded training distributes persistent parameters, gradients, and optimizer state across workers instead of keeping full copies of all of them on every worker. A common execution pattern gathers a parameter group before its computation and later returns to the sharded representation. Gradient reduction combines contributions across workers and can leave each worker with its gradient shard.

Exactly when parameters are gathered, kept, or released depends on the implementation and configuration. Backward computation may need another gather. Prefetching and communication overlap can keep multiple buffers live. This is more than slicing a weight matrix, and a diagram of parameter placement cannot prove the peak memory used by a training run.

With equivalent initial parameters, loss normalization and optimizer settings, a correctly synchronized sharded step should reproduce the intended baseline update within numerical tolerance. Matching a single loss value alone does not verify that update, state placement, communication, or memory use.

Separate stored parameters from peak memory

A biased Linear with input width D and output width K contains (D + 1) * K scalar parameters. If those parameters are divided evenly over S shards, each stores that count divided by S. Multiply by bytes per stored value to estimate parameter storage only.

Total and peak training memory also include gradients, activations, optimizer buffers, temporarily gathered parameters, communication buffers, and possible dtype-conversion copies. Parameter storage divided by S is not a general peak-memory formula or proof that a model will fit on a particular device.

The optimizer matters too. The constant-rate optax.sgd(lr) used here has no momentum buffer or parameter-sized optimizer state. NNX still tracks an optimizer step counter. Momentum SGD normally adds an accumulator, while Adam maintains first and second moments; their byte costs depend on shapes, dtypes and configuration.

A complete baseline example

This example uses a known scalar weight and bias. It checks the loss measured before SGD and the actual updated parameters:

import jax.numpy as jnp
import numpy as np
import optax
from flax import nnx

model = nnx.Linear(1, 1, rngs=nnx.Rngs(0))
model.kernel[...] = jnp.array([[1.]])
model.bias[...] = jnp.array([0.])
optimizer = nnx.Optimizer(model, optax.sgd(0.1), wrt=nnx.Param)
x = jnp.array([[1.], [2.]])
y = jnp.zeros((2, 1))

def loss_fn(m):
    return jnp.mean((m(x) - y) ** 2)

loss, grads = nnx.value_and_grad(loss_fn)(model)
optimizer.update(model, grads)
np.testing.assert_allclose(loss, 2.5)
np.testing.assert_allclose(model.kernel[...], [[0.5]])
np.testing.assert_allclose(model.bias[...], [-0.3])
np.testing.assert_allclose(loss_fn(model), 0.265, rtol=1e-5)

The gradients are 5 for the weight and 3 for the bias. The saved value 2.5 remains the pre-update loss even after the model changes. The new loss is a separate evaluation; this task does not return it. An arbitrary learning rate is not guaranteed to reduce the loss.