Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models

Model-parallel layers must gather weights before multiplying and scatter results after — Google's compiler pass breaks each collective and its dependent matrix multiply into matched slices that flow around a ring, so TPUs compute on one shard while the next is in flight, pushing 500-billion-parameter training to 72% of peak FLOPS.

Why this problem is worth solving

Scaling laws made big models a strategy: more parameters, better sample efficiency. But a trillion-parameter model does not fit in any accelerator’s memory, so production systems use

Partitioning manufactures communication. In the standard sharded transformer layer, weights live split across chips; before each

The interesting constraint: this must work automatically. Google’s stack compiles JAX and TensorFlow models through XLA for TPUs; an optimization that requires hand-rewriting each model would never reach production workloads. (Contrast CoCoNet — covered elsewhere on this site — which attacks the same barrier on GPUs with a user-facing DSL and autotuner; this paper’s answer is an invisible compiler pass.)

What people did before, and where it broke

Overlap independent work. Data-parallel training overlaps layer N’s gradient AllReduce with layer N−1’s backward computation — legal because they touch different tensors. Compiler literature extracts independent loops via fission, tiling, and pipelining. None of this applies when the communication’s output is the computation’s input.

Distributed matmul algorithms. Cannon-style 2D algorithms interleave shard rotation with partial products — the germ of the right idea, but specialized: they cover operands partitioned along contracting dimensions, not the zoo of partitioning strategies real models use (the paper treats Cannon’s algorithm as a special case of its scheme).

Offload engines. Proposals like ACE add hardware near the network to run collectives. That helps independent communication but cannot manufacture overlap where dependence forbids it.

The baseline that mattered. Google’s own production stack (the GSPMD partitioner in XLA) generated well-tuned synchronous collectives: gather, then multiply, then scatter — each individually optimal, jointly serialized.

The idea

The core observation: a collective and its dependent Einsum are only monolithically dependent — piecewise, they are not. The AllGather’s output is a concatenation of N shards; the Einsum over that concatenation is algebraically a combination of N partial Einsums, one per shard. Pick an order for both sequences such that partial computation i uses only shards that have already arrived, while transfer i ships the shard needed later. Same-index pairs share no data — so they run in parallel.

AllGather → Einsum becomes: each chip starts multiplying with the shard it already owns while asynchronously forwarding that shard to its ring neighbor. Next iteration, it computes with the shard that just arrived while forwarding it onward. After N iterations every chip has used all N shards; a DynamicUpdateSlice has been stitching each partial result into place. The full gathered tensor never needed to exist before compute began.

Einsum → ReduceScatter inverts the flow: what circulates is not an operand but the accumulating result shard. Each chip computes a partial product, adds it into the accumulator it just received from its neighbor, and passes the growing sum along. After N hops, each shard has visited every chip, collected every contribution, and landed on its home partition — the scatter happened during the multiplies.

Getting this to production quality is the paper’s second half:

• Asynchrony with a budget. Transfers become CollectivePermuteStart/Done pairs; a bottom-up scheduler pushes Starts early and Dones late while respecting the hardware’s limited synchronization flags. (It beat a top-down alternative by ~5%.)
• Bidirectional rings. TPU inter-chip links carry full bandwidth both ways; splitting each shard stream into clockwise and counterclockwise halves cuts transfer trips — worth the most on heavily-partitioned dimensions.
• Fusion diplomacy. XLA’s fusion heuristics, left alone, could fuse an addition into the wrong Einsum and create a false dependence on a pending transfer, serializing the loop. The pass rewrites concatenations into fusable Pad/Max forms and steers the fuser to keep independent work independent.
• Knowing when to decline. The decomposed ring uses only part of the torus’s bandwidth (one mesh dimension), so if compute is too small to hide the longer transfer, overlap loses. A simple peak-rate cost model gates the rewrite per operation, and picks which of two candidate collectives to overlap when an Einsum offers both.

timeline: [AllGather —— idle compute] → [Einsum —— idle links]

Why it works

• Concatenation and summation are linear → an Einsum over gathered shards equals combined partial Einsums → decomposition is exact, not approximate, so it can ship as a default-on compiler pass.
• Same-index slices share no data → the dependence that forbade overlap exists only at the monolithic granularity → at shard granularity, compute and transfer are free to run concurrently.
• Transfers become Start/Done pairs scheduled far apart → the interconnect works during multiplies → when compute time covers transfer time, communication costs approach zero observed time.
• Rings map onto the TPU torus’s dedicated inter-chip links, both directions → bidirectional streams halve trips → even heavily-partitioned dimensions keep transfers under compute.
• A cost model gates each rewrite → operations that would lose (thin compute, small shards) keep native collectives → the pass never has to be globally right, only locally honest — which is what makes fleet-wide automatic deployment safe.

What it costs

• Ring bandwidth discount. The decomposed sequence drives one mesh dimension of the torus rather than the full collective topology — the reason the cost gate exists at all.
• Prologue/epilogue overhead. Bidirectional transfer needs a pre-shift before the loop and a realigning addition after it, charged as unoverlappable in the estimate.
• Scheduler and hardware coupling. In-flight asynchronous transfers are capped by the system’s synchronization flags, and the whole design leans on TPU ICI links and XLA internals — porting to GPU clusters (the paper suggests NVLink-class fabrics) is plausible but undemonstrated here.
• Not all communication is covered. AlltoAll (mixture-of-experts routing) and AllGathers without a dependent Einsum stay synchronous — visible in T5’s ~10% AlltoAll runtime share and GLaM’s ~40% utilization ceiling.
• Compute must be big enough to hide behind. Narrow architectures (BigSSL_10B, GLaM_1T) get modest benefit — overlap hides latency; it cannot create arithmetic.

The evaluation

The technique ships in XLA and is measured on production-class workloads: six models from 10 billion to 1.03 trillion parameters — GPT-style (1T), the Meena conversational model (507B), Google’s MLPerf giant BERT (258B), T5 (513B), the GLaM mixture-of-experts model (1T), and the BigSSL speech model (10B) — running on TPU v4 pods of 128 to 2,048 chips, via JAX and TensorFlow frontends. The baseline is the same partitioning with ordinary synchronous collectives — Google’s own tuned production configuration.