TopolearnModels and batches that exceed one accelerator need sharding; JAX makes data and model parallelism declarative.
JAX scales from one device to thousands with the same program. The model is sharding: you describe how arrays are split across a grid of devices and jit compiles a program that runs on all of them, inserting communication as needed.
The pieces:
- A mesh names the device grid:
mesh = jax.make_mesh((4, 2), ("data", "model")). - A PartitionSpec says which array axis maps to which mesh axis:
P("data", None)splits rows overdataand replicates columns. NamedSharding(mesh, spec)combines them;jax.device_put(x, sharding)lays an array out accordingly.x.shardingandjax.debug.visualize_array_sharding(x)show the layout.
Then just call a jitted function. XLA's partitioner propagates shardings from inputs to outputs and inserts all-reduces and all-gathers automatically. Data parallelism is "shard the batch, replicate the params"; model parallelism shards weight matrices over the model axis.
You can constrain intermediates with jax.lax.with_sharding_constraint, or drop to manual per-device code with jax.shard_map, where the function sees local shards and you write collectives such as jax.lax.psum yourself.
Locally you can simulate several CPU devices with jax.config.update("jax_num_cpu_devices", 8) before any computation, which is great for testing sharding logic without hardware.
Example
import jax
jax.config.update("jax_num_cpu_devices", 8) # simulate 8 devices on CPU
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P
mesh = jax.make_mesh((4, 2), ("data", "model"))
x = jnp.ones((16, 8))
w = jnp.ones((8, 4))
x = jax.device_put(x, NamedSharding(mesh, P("data", None))) # shard batch
w = jax.device_put(w, NamedSharding(mesh, P(None, "model"))) # shard columns
@jax.jit
def layer(x, w):
return jnp.tanh(x @ w)
y = layer(x, w)
print(y.shape) # (16, 4)
print(y.sharding.spec) # P('data', 'model')
print(len(y.addressable_shards), y.addressable_shards[0].data.shape) # 8 (4, 2)Remember
- A mesh names a device grid; PartitionSpec maps array axes to mesh axes
jitpropagates shardings and inserts collectives automatically- Use
shard_mapfor explicit per-device code with manual collectives
Watch out for
- Setting jax_num_cpu_devices after JAX has already initialised its backends, which has no effect.