What is the SPMD partitioner in XLA?

The SPMD (Single Program, Multiple Data) partitioner is an XLA compiler pass that takes one whole-program computation graph annotated with shardings — descriptions of how each tensor is split across a mesh of devices — and rewrites it into a single per-device program that every chip runs on its own slice of the data. Where two operations want their data laid out differently, the partitioner inserts the exact cross-device collective needed to reshard between them: an all-gather, reduce-scatter, all-to-all, collective-permute, or halo exchange. The output is SPMD by construction: every device executes identical code, differing only in which shard of each tensor it holds and which slice of each collective it contributes.

What is sharding propagation (GSPMD)?

Sharding propagation is the pass that runs before partitioning. You annotate only a handful of tensors — often just the inputs and outputs — and propagation infers a sharding for every other tensor in the graph by pushing annotations along the edges. It runs both directions: forward, inferring an operation’s sharding from its operands (XLA’s ShardingPropagation calls this InferShardingFromOperands), and backward, inferring an operand’s sharding from how a consumer uses it (GetShardingFromUser). This forward-and-backward flood is the core idea of GSPMD (“General and Scalable Parallelization for ML Computation Graphs”): annotate a few tensors, let the compiler fill in the rest, then partition.

When does the SPMD partitioner insert an all-reduce?

When a matrix multiply (or any reduction) has its contracting dimension sharded across devices. If A[M,K] and B[K,N] are both split along the shared K dimension, each device multiplies only its local K-block, producing a partial sum — a full-shape result that is only part of the answer. Summing those partial sums across the devices that hold different K-blocks is exactly an all-reduce. In XLA’s dot handler this is emitted by AllReduceAlongShardingDims. If the output is itself sharded, the partitioner emits a reduce-scatter instead — the same reduction, but each device keeps only its output slice, which is cheaper. Non-contracting (row- or column-parallel) matmuls need no reduction at all: each device already computes a complete slice of the output.

What is halo exchange in SPMD partitioning?

Halo exchange is the collective the partitioner inserts for windowed operations — convolutions, pooling, reduce-window — whose output element depends on a neighbourhood of input elements. When such a tensor is sharded along a spatial dimension, computing the shard’s boundary outputs needs a few input rows that live on the neighbouring device. Each device exchanges those boundary “halo” regions with its neighbours (via collective-permutes), concatenates them onto its local shard, and then runs the op. XLA’s utility “performs halo exchange on the given dimension” and “returns nullopt if the halo is beyond the neighbor shard” — i.e. if a shard would need data from more than one hop away, the simple exchange gives up and a heavier reshard is used instead.

Why do I still need manual sharding annotations if the compiler propagates them?

Because propagation is a best-effort heuristic, not a solver, and it can reach a layout it has no efficient way to produce. When the partitioner cannot find a cheap collective to go from a tensor’s current sharding to the one an operation needs, it falls back to replicating the whole tensor and re-slicing it — XLA logs this as “Involuntary full rematerialization… which is inefficient.” Manual annotations exist to steer propagation away from those dead ends: pinning the sharding on a few key tensors (the ones that decide your parallelism strategy) keeps the inferred layouts consistent so the compiler inserts the collectives you intended rather than an accidental full replication.

The SPMD Partitioner: One Program Across a Device Mesh

You write a program as if one giant machine ran it, annotate a few tensors with how they are split across a mesh of chips, and the compiler rewrites it into N identical per-device programs — inserting every all-gather, reduce-scatter, all-to-all, and halo exchange needed to keep the maths correct. We trace how those annotations propagate through the graph, where each collective is born (and which resharding forces it), the one windowed op that needs halo exchange, and the case the partitioner cannot solve on its own — drawn, computed, and animated.

System design · AI / ML. The source ↗

A free, interactive, animated visual explainer of The SPMD Partitioner: One Program Across a Device Mesh — built to be understood, not skimmed.

Questions

What is the SPMD partitioner in XLA?: The SPMD (Single Program, Multiple Data) partitioner is an XLA compiler pass that takes one whole-program computation graph annotated with shardings — descriptions of how each tensor is split across a mesh of devices — and rewrites it into a single per-device program that every chip runs on its own slice of the data. Where two operations want their data laid out differently, the partitioner inserts the exact cross-device collective needed to reshard between them: an all-gather, reduce-scatter, all-to-all, collective-permute, or halo exchange. The output is SPMD by construction: every device executes identical code, differing only in which shard of each tensor it holds and which slice of each collective it contributes.
What is sharding propagation (GSPMD)?: Sharding propagation is the pass that runs before partitioning. You annotate only a handful of tensors — often just the inputs and outputs — and propagation infers a sharding for every other tensor in the graph by pushing annotations along the edges. It runs both directions: forward, inferring an operation’s sharding from its operands (XLA’s ShardingPropagation calls this InferShardingFromOperands), and backward, inferring an operand’s sharding from how a consumer uses it (GetShardingFromUser). This forward-and-backward flood is the core idea of GSPMD (“General and Scalable Parallelization for ML Computation Graphs”): annotate a few tensors, let the compiler fill in the rest, then partition.
When does the SPMD partitioner insert an all-reduce?: When a matrix multiply (or any reduction) has its contracting dimension sharded across devices. If A[M,K] and B[K,N] are both split along the shared K dimension, each device multiplies only its local K-block, producing a partial sum — a full-shape result that is only part of the answer. Summing those partial sums across the devices that hold different K-blocks is exactly an all-reduce. In XLA’s dot handler this is emitted by AllReduceAlongShardingDims. If the output is itself sharded, the partitioner emits a reduce-scatter instead — the same reduction, but each device keeps only its output slice, which is cheaper. Non-contracting (row- or column-parallel) matmuls need no reduction at all: each device already computes a complete slice of the output.
What is halo exchange in SPMD partitioning?: Halo exchange is the collective the partitioner inserts for windowed operations — convolutions, pooling, reduce-window — whose output element depends on a neighbourhood of input elements. When such a tensor is sharded along a spatial dimension, computing the shard’s boundary outputs needs a few input rows that live on the neighbouring device. Each device exchanges those boundary “halo” regions with its neighbours (via collective-permutes), concatenates them onto its local shard, and then runs the op. XLA’s utility “performs halo exchange on the given dimension” and “returns nullopt if the halo is beyond the neighbor shard” — i.e. if a shard would need data from more than one hop away, the simple exchange gives up and a heavier reshard is used instead.
Why do I still need manual sharding annotations if the compiler propagates them?: Because propagation is a best-effort heuristic, not a solver, and it can reach a layout it has no efficient way to produce. When the partitioner cannot find a cheap collective to go from a tensor’s current sharding to the one an operation needs, it falls back to replicating the whole tensor and re-slicing it — XLA logs this as “Involuntary full rematerialization… which is inefficient.” Manual annotations exist to steer propagation away from those dead ends: pinning the sharding on a few key tensors (the ones that decide your parallelism strategy) keeps the inferred layouts consistent so the compiler inserts the collectives you intended rather than an accidental full replication.

The SPMD Partitioner: One Program Across a Device Mesh

Questions

Related explainers