ProDeploy — Learn the fabric that keeps AI clusters alive

TPU Pod 3D torus — communication patterns

The amber node is the active TPU. Toggle to see how communication differs between convolutional and transformer workloads.

Convolution: each TPU communicates only with its 6 torus neighbours. Topology matches workload perfectly.

Google's TPU Pod is the most sophisticated deployment of torus topology for machine learning workloads. It shows both the ceiling of what torus can achieve and why transformer workloads push past that ceiling.

A TPU v3 Pod connects 2,048 TPUs in a 3D torus using a dedicated inter-chip interconnect (ICI) running at 656 GB/s per chip in each direction. The torus has dimensions of roughly 16x16x8 depending on pod size. Links are short (chips are physically adjacent in 3D); the folded torus ensures no cable is long. Latency is extremely low.

The workload optimisation is explicit: TPUs were designed for matrix multiplication in convolutional neural networks. Convolution has strong spatial locality in the tensor parallelism pattern Google uses. The convolution computation is partitioned across the torus so that each TPU primarily communicates with its immediate neighbours in the 3D mesh. This matches the torus perfectly.

Google's mesh tensor parallelism strategy -- XLA SPMD partitioning -- explicitly arranges the model tensor dimensions onto torus axes so that the communication during each layer is local. A convolutional layer with spatial height and width can be partitioned along those spatial axes, which map directly to torus axes.

The transformer problem: transformer self-attention is not spatially local. The attention operation requires every token to attend to every other token. In tensor parallelism across a transformer, the AllGather and ReduceScatter operations that synchronise activations across devices are inherently all-to-all. There is no partition of the computation that keeps communication local.

Google has invested heavily in optimising TPUs for transformer workloads -- their later TPU generations use wider per-chip bandwidth, better collective operation primitives, and more sophisticated SPMD partitioning strategies. But the fundamental tension between torus topology and all-to-all communication patterns remains. NVIDIA's choice of fat-tree for DGX deployments reflects a different set of workload assumptions: the world of AI training was dominated by transformer-based models, where all-to-all AllReduce is the bottleneck.

Both choices are defensible. They reflect honest disagreement about which bottleneck matters more: the bisection bandwidth that fat-tree provides vs. the lower latency and hardware simplicity that torus provides.