The Training Run

Nothing in this chapter is new, and that's deliberate. Eleven chapters built the pieces; a training run is what happens when you wire them together in the order a real loop needs them. Params and optimizer state ride as one registered pytree, the pattern chapters 7 and 9 built. The step itself is a jitted function computing gradients, the machinery chapters 3 and 4 gave you. Iteration is a scan or a plain Python loop over steps, chapter 6's territory. Randomness folds in per step by index, the discipline chapter 8 insisted on. If the mesh is real, the arrays carry sharding from chapter 10. And every number the loop reports, loss, throughput, time per step, gets measured the way chapter 11 taught: warmup, block_until_ready, provenance stated.

This chapter walks one complete run: small, real, and fully specified, an MLP fitting a synthetic regression target end to end. The mastery work at the end is where you scale it, with real data, a bigger model, a real mesh. The loop's shape doesn't change when you do.

A training run is composition, not invention.

One jitted function does the actual work of a step: compute the loss and its gradient together with value_and_grad, hand the gradients to the optimizer's update, apply the resulting updates to the parameters, and return the new state alongside whatever metrics you want to keep. On an accelerator, donate_argnums lets that function recycle the old parameters' buffers instead of allocating fresh ones for the new values, since the old buffer is dead the instant the new one exists.

Data stays outside the jitted world in spirit, even though it's passed as an argument: it lives as an iterator of ordinary host or NumPy arrays, handed in fresh each step rather than carried inside the traced function. Metrics leave the same way params do, as returned values from the jitted function, or through io_callback when you need them to stream out mid-run instead of waiting for the step to finish, the callback family chapter 9 introduced.

import jax
import jax.numpy as jnp
import optax

def init(key):
    k1, k2 = jax.random.split(key)
    return {
        "w1": jax.random.normal(k1, (8, 32)) * 0.1, "b1": jnp.zeros(32),
        "w2": jax.random.normal(k2, (32, 1)) * 0.1, "b2": jnp.zeros(1),
    }

def predict(p, x):
    h = jax.nn.relu(x @ p["w1"] + p["b1"])
    return h @ p["w2"] + p["b2"]

def loss_fn(p, batch):
    x, y = batch
    return jnp.mean((predict(p, x) - y) ** 2)

opt = optax.adam(1e-2)

@jax.jit
def train_step(p, opt_state, batch):
    loss, grads = jax.value_and_grad(loss_fn)(p, batch)
    updates, opt_state = opt.update(grads, opt_state, p)
    return optax.apply_updates(p, updates), opt_state, loss

key = jax.random.key(0)
params = init(key)
opt_state = opt.init(params)
x = jax.random.normal(jax.random.key(1), (256, 8))
y = jnp.sum(x, axis=1, keepdims=True)

for step in range(200):
    params, opt_state, loss = train_step(params, opt_state, (x, y))
print(f"final loss {loss:.5f}")   # final loss 0.00335 (verified, jax 0.4.38 CPU)

A checkpoint that only saves params isn't a checkpoint of your training run. It's a checkpoint of a different run that happens to start from the same weights. Adam without its moment estimates is a different optimizer wearing the same name, and resuming from params alone while re-initializing everything else changes the trajectory even if nobody notices for a while. The unit of checkpointing is the whole state tree: params, optimizer state, step, and key together.

Orbax is the library that saves and restores that tree, as pytrees, sharding-aware, so a checkpoint written by a sharded run restores correctly onto a different mesh shape. The key half of the tree resumes cleanly because of a choice chapter 8 already made: fold_in(key, step) derives each step's randomness from the step number directly, so restoring at step k and continuing produces exactly the same bits from k onward that an uninterrupted run would have produced.

The bar for resume working isn't a plausible-looking loss curve. It's bitwise: kill the run, restore the checkpoint, continue, and the loss at every subsequent step matches the run that was never interrupted, digit for digit. Prove that. Don't assume it.

Multi-host JAX isn't a different programming model. It's the same program run once per host, after each host calls jax.distributed.initialize() to find the others. Once initialized, a jax.Array can be sharded globally across every device on every host, not just the devices attached to one process, and jit compiles a single SPMD program that runs across the whole fleet.

Each host feeds only its own local shard of the batch into the step, the same way each device did inside a single host in chapter 10, and the collectives that keep gradients and activations synchronized cross host boundaries exactly as they crossed device boundaries there. Nothing about the mesh or the PartitionSpec vocabulary changes; the ring just gets bigger.