SPMD

substrax.spmd is the data-parallel training path on top of substrax.mesh: place a batch on the data axis, run one step, reduce metrics across devices.

import jax
from flax import nnx

from substrax.mesh import DeviceMeshManager
from substrax.spmd import create_data_parallel_sharding, place_batch_on_shards, spmd_train_step

mesh = DeviceMeshManager.create_device_mesh({"data": jax.device_count()})
sharding = create_data_parallel_sharding(mesh)


@nnx.jit
def train_step(model, optimizer, batch):
    return spmd_train_step(model, optimizer, loss_fn, batch)


with jax.set_mesh(mesh):
    loss = train_step(model, optimizer, place_batch_on_shards(batch, sharding))

place_batch_on_shards also takes a pytree prefix of the batch whose leaves are shardings, one per subtree, for a batch whose fields are not all split along the batch axis: rows on the data axis and a batch-level value, which has no batch axis, replicated.

from jax.sharding import NamedSharding, PartitionSpec

replicated = NamedSharding(mesh, PartitionSpec())
batch = {"rows": {"x": x, "y": y}, "batch_state": {"lam": lam}}
placed = place_batch_on_shards(batch, {"rows": sharding, "batch_state": replicated})

On several processes each passes the rows it loaded and jax assembles the global batch; a replicated leaf must be equal on every process. A leaf with fewer dimensions than its sharding splits, such as a scalar under the data-axis sharding, is refused with its path.

spmd_train_step differentiates with nnx.value_and_grad and updates the nnx.Optimizer; XLA inserts the gradient all-reduce when the parameters are sharded under jax.set_mesh. The reductions (reduce_mean, reduce_sum, reduce_max, reduce_min) and the collectives (reduce_mean_collective, reduce_sum_collective, all_gather) reduce every array leaf of the tree they are given, at any depth, so a gradient tree beside scalar metrics is reduced leaf by leaf, and return a tree of the same structure; any other leaf passes through. reduce_custom and collect_from_devices take a flat mapping from metric name to value.

substrax.spmd

Data-parallel placement, the SPMD training step and collectives.

A data-parallel gradient needs no reduction after differentiating. Under jax.jit with the batch sharded over a mesh axis, the gradient of the batch's loss is already the full gradient: the compiler inserts the all-reduce (spmd_train_step relies on this). Inside jax.shard_map each device sees its shard, and the loss is averaged over the axis before differentiating, jax.grad(lambda w: jax.lax.pmean(loss(w, shard), axis)); the gradient of a replicated parameter then comes back as the full one. Averaging the gradient afterwards is wrong by the axis size, because the gradient of a replicated parameter already arrives summed over the axis. reduce_mean_collective and reduce_sum_collective average or sum per-device metrics inside shard_map.

all_gather(metrics, axis_name='batch')

Gather metrics from all devices.

Only valid inside a pmap or shard_map context.

Parameters:

Name	Type	Description	Default
metrics	T	The metrics to gather: a mapping, or any pytree of them.	required
axis_name	str	The name of the axis to gather across.	'batch'

Returns:

Type	Description
T	metrics with every array leaf gathered across the axis.

collect_from_devices(metrics)

Collect metrics from all devices.

Call outside of a pmapped function to split per-device values from the leading device axis.

Parameters:

Name	Type	Description	Default
metrics	Mapping[str, V]	The metrics from all devices, with the first dimension corresponding to the device axis.	required

Returns:

Type	Description
dict[str, list[Array] | V]	A dictionary of metrics, with array values split into per-device lists.

reduce_custom(metrics, reduce_fn=None)

Apply custom reduction operations to metrics.

Uses standard JAX operations. Works in SPMD contexts. Each named value is reduced by its operation when it is an array; without reduce_fn this is reduce_mean.

Parameters:

Name	Type	Description	Default
metrics	Mapping[str, V]	The metrics to reduce, by name.	required
reduce_fn	Mapping[str, str | None] | None	A dictionary mapping metric names to reduction operations. Each operation should be one of {"mean", "sum", "max", "min"}. If None, defaults to "mean" for all metrics.	None

Returns:

Type	Description
Mapping[str, V | Array]	A dictionary of reduced metrics.

reduce_max(metrics)

Compute the maximum of metrics using standard JAX operations.

Works with global arrays in SPMD contexts (nnx.jit + mesh). Every array leaf with at least one dimension is reduced, at any depth of metrics; every other leaf passes through unchanged.

Parameters:

Name	Type	Description	Default
metrics	T	The metrics to reduce: a mapping, or any pytree of them.	required

Returns:

Type	Description
T	metrics with every array leaf reduced to its maximum.

reduce_mean(metrics)

Compute the mean of metrics using standard JAX operations.

Works with global arrays in SPMD contexts (nnx.jit + mesh). Every array leaf with at least one dimension is reduced, at any depth of metrics; every other leaf passes through unchanged.

Parameters:

Name	Type	Description	Default
metrics	T	The metrics to reduce: a mapping, or any pytree of them.	required

Returns:

Type	Description
T	metrics with every array leaf mean-reduced.

reduce_mean_collective(metrics, axis_name='batch')

Compute the mean of metrics using collective operations.

Only valid inside a pmap or shard_map context.

Parameters:

Name	Type	Description	Default
metrics	T	The metrics to reduce: a mapping, or any pytree of them.	required
axis_name	str	The name of the axis to reduce across.	'batch'

Returns:

Type	Description
T	metrics with every array leaf averaged across the axis.

reduce_min(metrics)

Compute the minimum of metrics using standard JAX operations.

Works with global arrays in SPMD contexts (nnx.jit + mesh). Every array leaf with at least one dimension is reduced, at any depth of metrics; every other leaf passes through unchanged.

Parameters:

Name	Type	Description	Default
metrics	T	The metrics to reduce: a mapping, or any pytree of them.	required

Returns:

Type	Description
T	metrics with every array leaf reduced to its minimum.

reduce_sum(metrics)

Compute the sum of metrics using standard JAX operations.

Works with global arrays in SPMD contexts (nnx.jit + mesh). Every array leaf with at least one dimension is reduced, at any depth of metrics; every other leaf passes through unchanged.

Parameters:

Name	Type	Description	Default
metrics	T	The metrics to reduce: a mapping, or any pytree of them.	required

Returns:

Type	Description
T	metrics with every array leaf summed.

reduce_sum_collective(metrics, axis_name='batch')

Compute the sum of metrics using collective operations.

Only valid inside a pmap or shard_map context.

Parameters:

Name	Type	Description	Default
metrics	T	The metrics to reduce: a mapping, or any pytree of them.	required
axis_name	str	The name of the axis to reduce across.	'batch'

Returns:

Type	Description
T	metrics with every array leaf summed.

create_data_parallel_sharding(mesh, data_axis='data')

Create the sharding that splits a batch's leading axis over a mesh axis.

Parameters:

Name	Type	Description	Default
mesh	Mesh	The device mesh to use for sharding.	required
data_axis	str	The name of the mesh axis to use for data parallelism.	'data'

Returns:

Type	Description
NamedSharding	A NamedSharding over mesh with spec PartitionSpec(data_axis).

place_batch_on_shards(batch, sharding)

Turn this process's batch into global arrays, each subtree on its sharding.

Array leaves, whether jax.Array, a NumPy array as a host data loader yields it, or a NumPy scalar, go to jax.make_array_from_process_local_data in one call. On a single process that is one batched jax.device_put. On several processes each process passes the slice of the global batch it loaded, and jax stitches the slices into one global array; a leaf on a replicated sharding must be equal on every process. Any other leaf, such as a string, is returned as it is. Call it on the host batches a loader yields, outside jax.jit: inside a traced function the result is not placed on sharding.

Parameters:

Name	Type	Description	Default
batch	PyTree	The batch this process loaded.	required
sharding	Sharding | PyTree	One sharding for every array leaf, or a pytree prefix of batch whose leaves are shardings, each applying to its subtree: rows sharded on the batch axis and a batch-level field without one replicated, for example.	required

Returns:

Type	Description
PyTree	The batch with every array leaf replaced by the global array on its sharding.

Raises:

Type	Description
ValueError	If a leaf has fewer dimensions than its sharding partitions (a scalar under a sharding that splits the batch axis); the message names the leaf's path.

place_nnx_state_on_shards(state, mesh, filter_sharding)

Shard a Flax NNX state tree using current NNX sharding helpers.

spmd_train_step(model, optimizer, loss_fn, batch)

Execute a data-parallel training step using SPMD.

Uses nnx.value_and_grad for automatic differentiation. The XLA compiler handles gradient AllReduce automatically when model parameters are sharded across devices via jax.set_mesh or explicit NamedSharding.

This function should be called inside an @nnx.jit decorated function with a mesh context active.

Parameters:

Name	Type	Description	Default
model	M	The NNX model to train.	required
optimizer	Optimizer[M]	The NNX optimizer wrapping the model.	required
loss_fn	Callable[[M, PyTree], Array]	Function (model, batch) -> loss scalar.	required
batch	PyTree	The training batch (should be pre-sharded).	required

Returns:

Type	Description
Array	The loss value.

Example

mesh = DeviceMeshManager.create_device_mesh({"data": 4}) rules = data_parallel_rules()

@nnx.jit def train_step(model, optimizer, batch): return spmd_train_step(model, optimizer, my_loss_fn, batch)

with jax.set_mesh(mesh): loss = train_step(model, optimizer, batch)