Sixteen JAX devices on a CPU, and one gradient bug worth finding

My favorite use for a pretend cluster is a real correctness failure. Here, replacing an average with a sum makes the gradient 16 times larger, even though every shard computes a perfectly sensible local loss. 🐛

Rishiraj Acharya's new Jaxotron walkthrough shows how to expose logical CPU devices to JAX and explore distributed layouts without renting accelerators. I took a smaller route: one linear model, one batch and a deliberately wrong reduction. The whole check fits in a single script. [1]

Original explanatory diagram, with values from the executed JAX example. Logical CPU devices share one host; the drawing does not represent separate GPUs.

Start with a reference you can calculate

The model has two weights. Its prediction is x @ w, and the loss is the mean squared error over the batch. For a matrix of inputs X, labels y and N examples, its gradient is:

gradient = 2 * X.T @ (X @ w - y) / N

I use 32 synthetic examples, split into 16 equal shards. Each device gets two examples and the same initial weights. The data are deterministic; no random seed or downloaded model is involved.

With equally sized shards, the average of their mean-loss gradients equals the mean-loss gradient over the full batch. That follows directly from the denominator: averaging sixteen groups of two restores the division by thirty-two. 💡

Summing those local means keeps the division by two. The result is sixteen times the intended gradient. An optimizer may still reduce the loss, especially if someone compensates with a smaller learning rate. A falling curve is therefore a weak test for this particular mistake.

Run the small failure before the big training job 🧪

Use Python 3.12 and a fresh process. The script sets the CPU device flag before importing JAX; setting it after the runtime initializes is too late for this setup.

I pinned JAX 0.4.35 to keep the demonstrated pmap behavior reproducible. This is an intentional example version, not a claim that it is the newest release. With the uv package runner installed, execute:

uv run --python 3.12 --with jax==0.4.35 --with numpy==2.1.3 python mesh_check.py

Save this as mesh_check.py:

import os
os.environ['JAX_PLATFORMS'] = 'cpu'
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=16'
import jax
import jax.numpy as jnp
import numpy as np
from functools import partial
x = jnp.arange(64, dtype=jnp.float32).reshape(32, 2) / 32
y = x @ jnp.array([2., -1.]) + .25
w = jnp.zeros(2, dtype=jnp.float32)
def loss(w, x, y):
    return jnp.mean((x @ w - y) ** 2)
reference = jax.grad(loss)(w, x, y)
@partial(jax.pmap, axis_name='data', in_axes=(None, 0, 0))
def distributed(w, x, y):
    local = jax.grad(loss)(w, x, y)
    return jax.lax.pmean(local, 'data'), jax.lax.psum(local, 'data')
good, bad = distributed(w, x.reshape(16, 2, 2), y.reshape(16, 2))
np.testing.assert_allclose(good, np.broadcast_to(reference, good.shape), rtol=1e-6)
np.testing.assert_allclose(bad, np.broadcast_to(16 * reference, bad.shape), rtol=1e-6)
print('JAX:', jax.__version__)
print('logical CPU devices:', jax.local_device_count())
print('reference gradient:', np.asarray(reference).tolist())
print('largest pmean error:', float(jnp.max(jnp.abs(good - reference))))
print('psum / reference:', np.asarray(bad[0] / reference).tolist())

pmap maps the function across devices. in_axes=(None, 0, 0) broadcasts the weights and slices the leading dimension of the input and label arrays. The named data axis tells pmean and psum which group to reduce. JAX documents these as a cross-device mean and sum respectively. [2][3]

My output:

JAX: 0.4.35
logical CPU devices: 16
reference gradient: [-2.966796875, -3.041015625]
largest pmean error: 0.0
psum / reference: [16.0, 16.0]

✅ The average matches the single-device gradient exactly in this run. Both gradient coordinates are nonzero, so the displayed ratio is meaningful. The assertions check every replica, rather than just the convenient first row. I also checked the reference with exact rational arithmetic independently of JAX.

The mean of means has a condition

⚠️ Equal shard sizes are doing real work in that derivation. If one shard contains more valid examples than another, an unweighted average of local means gives each shard equal influence, rather than each example.

For unequal valid counts, reduce the sum of per-example gradients and divide by the total valid count. With padding, the same mask must affect the numerator and count. A padded row that contributes zero loss but still increases the denominator changes the objective.

That distinction is more useful to debug on a tiny deterministic batch than inside a long training run. Keep the reference, compare actual gradients, and only then add the next parallel dimension.

What the CPU trick buys

This runs real JAX autodiff and collectives on logical CPU devices. It does not emulate a GPU network, accelerator memory limits or communication latency. The zero error above is specific to this data and arithmetic; floating-point reductions need not match bit for bit in general.

I did not reproduce Jaxotron's full training, checkpoint restoration or all its parallel layouts. This is a focused prerequisite check: if splitting the batch changes the objective, a larger device mesh will only distribute the mistake more efficiently. 🚦

References

[1] Rishiraj Acharya, How to Simulate a 16-Device Cluster on Your Laptop.

[2] JAX, lax.pmean.

[3] JAX, lax.psum.

Original diagram and synthetic numerical example created for this article.

AI disclosure

AI was used to draft the article.