TPU and GPU Clusters: The Anatomy of Collective Communication

Inside TPU and GPU Clusters: The Anatomy of Collective Communication

← Back to blog code]:bg-gray-200 [&_*:not(pre)>code]:text-red-600 [&_*:not(pre)>code]:px-1 [&_*:not(pre)>code]:py-0.5 [&_*:not(pre)>code]:rounded [&_*:not(pre)>code]:text-sm [&_*:not(pre)>code]:font-mono [&_ol]:text-gray-700 [&_ol]:list-decimal [&_ol]:list-inside [&_ol]:pl-6 [&_ol]:mb-6 [&_ol]:mt-3 [&_li]:mb-3 [&_ul]:list-disc [&_ul]:list-inside [&_ul]:pl-6 [&_ul]:space-y-1">In this post, I'll do a deep dive into TPU and GPU cluster topologies, and cover the core collective operations used during transformer training and inference. Why should we care? In 2026, training and serving transformers is a massively distributed systems problem. To shard models across a cluster, we rely on techniques such as data parallelism , tensor/model parallelism , FSDP , and expert parallelism . Under the hood, these techniques are built on a small set of core collective operations: Data parallel training requires gradient synchronization during backprop, this is typically implemented with all-reduce . As we'll see later in this post, all-reduce itself can be decomposed into reduce-scatter followed by all-gather . Tensor parallelism and FSDP rely heavily on all-gather and reduce-scatter in the forward and backward passes. Expert parallelism, used in MoE models, rely on the all-to-all primitive. These are just a few important examples, but they're enough to show why understanding collective communication is useful -> if you want to reason about the performance of modern transformer systems, you eventually have to reason about how data moves through the cluster. We'll start from the hardware, with the TPU and GPU cluster topologies. Understanding the physical layout of the cluster makes the collective algorithms much more grounded and easier to reason about. From there, we'll dive into the most common implementations of the core collective operations. I'll focus primarily on ring-style algorithms, since they are the natural starting point for large-message communication. For smaller payloads, latency starts to dominate, and tree-style algorithms can be a better fit (as they require only log2 steps). This post is structured into seven parts: TPU cluster topology: Superpods, Slices, DCN, PCIe, ICI Inside All-Gather: 1D/2D Rings, and Chains Reduce-Scatter (and All-Reduce): The Dual of All-Gather All-to-All: A Sharded Transpose NVIDIA GPU cluster topology: Nodes, Scalable Units, Fat Tree GPU Collectives Within the Node: Rings, Trees, and SHARP GPU Collectives Across Nodes: Hierarchical Algorithms over InfiniBand TPU cluster topology I'll start with TPUs because their topology is more uniform, and therefore arguably easier to reason about, than GPU cluster topology. The key difference between TPU and GPU clusters is the nearest-neighbor connectivity . TPU chips are connected directly to neighboring TPU chips, and each chip has either 4 or 6 nearest neighbors , depending on the TPU generation: TPU v2, v3, v5e, and v6e use a 2D torus topology, with 4 nearest neighbors. TPU v4p, v5p, TPU7x (Ironwood), and 8t use a 3D torus topology, with 6 nearest neighbors. Figure 1: TPU connectivity classes

💡Boardfly topology: Notably, Google's new inference TPU chip [1] 8i deviates from the 2D/3D torus, and instead uses boardfly, a hierarchical high-radix topology. In this post I'll ignore it.

Here is one way to build intuition for the 4-neighbor, 2D torus pattern. A 2D torus can be represented as a grid with wraparound/periodic boundaries: moving off the left edge brings you back on the right edge (and vice versa), and moving off the top edge brings you back on the bottom edge (and vice versa). Bear in mind that the TPU torus is a discrete grid; imagine it overlaid on this donut. The visualization is just for intuition: Figure 2: 2D torus intuition: a grid with wraparound boundaries

Similarly, here is the 6-neighbor, 3D torus connectivity pattern: Figure 3: 3D torus connectivity

This looks a bit messy, but the idea is simple -> each chip has neighbors along ±x, ±y, and ±z, and the edges wrap around in all three dimensions. 💡Interactive Visualization Tool: Here is a neat little visualization tool [2] I used to generate the above image, you can use it to interact with these more complex topologies.

We'll use v5e as the running example in this post, since 2D connectivity is easier to visualize. TPU chips communicate with their neighbors over ICI , short for inter-chip interconnect. The largest ICI-connected island of TPU chips is called a TPU Pod . You'll also sometimes hear people refer to this as a "superpod" ; I'll use the terms interchangeably here. For example, a v4 pod contains 16 × 16 × 16 chips, for a total of 4096 TPUs. A v5p pod contains 16 × 20 × 28 chips, or 8960 TPUs! 💡GPUs scale-up domains: Contrast this with GPUs, where the scale-up domain has traditionally been much smaller: commonly 8 GPUs in a single NVLink domain, and more recently 72 GPUs in an NVLink...

TPU and GPU Clusters: The Anatomy of Collective Communication

Related Articles

(no title)

Scientists reverse brain aging, with a nasal spray

AI has torched the market for junior programmers

Al Vigier: Canada's AI strategy shouldn't include secret Palantir bills

The Anatomy of an AI-Native Org