SPMD

The JAX path taught the surface of this from above, at /jax/sharding: a Mesh, a PartitionSpec, a NamedSharding, and a compiler that quietly inserted whatever collectives your program's math required. This chapter is the mechanism underneath that surface. SpmdPartitioner (xla/service/spmd/spmd_partitioner.h:274) is an HloModulePass, and its job is exactly what it sounds like: take a module written for one logical, global-shape array and rewrite it into the module every device actually runs, one physically local shape per participant.

Every all-gather, reduce-scatter, and all-reduce you never wrote by hand in the JAX path came from this pass deciding, at a specific point in the dataflow, that a sharding boundary demanded data movement. Nothing about your Python asked for a collective explicitly. The collective is what SpmdPartitioner inserts when two adjacent operations disagree about how a value should be split.

Rewriting the module is the second half of this pass's job; propagation is the first, and it runs to completion before any rewriting starts. Sharding annotations you gave explicitly flow forward and backward through the graph until nothing changes anymore, a fixpoint, and every instruction you did not annotate simply inherits whatever its neighbors settled on.

It is easy to assume every instruction needs its own explicit annotation, and that assumption makes the mental model far more work than it needs to be. Most of a real module carries no sharding you wrote at all; propagation supplies it. The cases worth watching are the disagreements, where two neighbors would prefer conflicting shardings for the same value, and those get resolved by priority rules that favor the more informative or more constrained side.

GSPMD, the design this whole mechanism implements, gets its name from the acronym: generalized single program, multiple data. Every device that participates ends up running the identical compiled binary; what differs from device to device is never the code, only the shard of data that binary happens to be looking at.

One program, partitioned, is still one program.

That single sentence is the whole contract, and it is why the partitioner has to be this careful: a collective it inserts on one device's copy of the program has to be matched by the exact same collective, at the exact same logical point, on every other device's copy. Get that wrong and you are not looking at a slow program, you are looking at chapter 8's failure mode.

A newer dialect, sdy, short for Shardy, is changing where the sharding annotation actually lives: in MLIR, attached to StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →, before the module ever becomes HLO at all. That is an earlier altitude than the mechanism this chapter describes, not a different mechanism. Whether the annotation started life in MLIR or was inherited by an HLO-level pass, the destination is the same partitioned module SpmdPartitioner has always produced.