TPU and GPU Clusters: The Anatomy of Collective Communication

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 domain with NVIDIA GB200 NVL72. We'll analyze GPUs in much more depth a bit later.

For TPU pods whose chips have 6 neighbors,the smallest full 3D torus is a 4×4×4 cube. If you request a smaller topology, for example 2×2×2, you lose the wraparound links, so the slice stops being a torus and becomes a mesh.

💡Terminology:

A TPU slice is a collection of chips within a single TPU Pod connected through ICIs (an "ICI-island"). For v5e those can be: 2×2, 2×4, 8×8, etc. The largest such slice is the (super)pod itself.

This can roughly double the time of ring-style collectives along that axis, as we'll soon see. So this is something to be aware of - if your application is communication-heavy, you may want to request slice shapes that preserve the torus structure, for example 4×4×8.

v5e, on the other hand, has a 16×16 2D torus pod. The wraparound exists at size 16, so if you take a smaller slice, such as 8×16, you lose the wraparound along the shorter axis (again, collective ops along that axis pay roughly a 2× penalty).

To scale beyond a single pod, TPUs use DCN, short for data center networking. DCN has much lower throughput than ICI, so you have to be careful: if too much communication crosses pod boundaries, it can easily become the bottleneck during training.

Putting all of that together, let's visualize the topology of a 16×16 v5e pod. Spend some time analyzing this (zoom in if needed):

Figure 4: Topology of a v5e TPU 16x16 superpod

We can connect multiple pods into an even larger compute cluster through a shared DCN fabric (maybe this is what we should call a superpod, heh):

Figure 5: Multiple pods connected through a DCN backbone

A few more details are worth knowing.

As can be seen in the figure above, within the pod the 0th row connects back to the 15th row, and the 0th column connects back to the 15th column. This is what makes the topology a torus/donut instead of a regular mesh. The torus geometry constrains path lengths, because data often has to flow through intermediate TPUs to get from the source chip to the target chip.

For example, if TPU (15, 15) wants to send data to TPU (2, 15), the shortest path goes through the wraparound link:

(15, 15) → (0, 15) →(1, 15) → (2, 15)

instead of going the long way through (14, 15), (13, 15), and so on.

💡Extra info:

TPUs also support “twisted torus” configurations, which alter the wraparound connectivity to reduce hop counts for communication patterns such as All-to-All. This is an implementation detail that improves efficiency, but we don't need it for the purposes of this post.

Each TPU chip is also connected to a "dedicated" CPU host through PCIe. In the case of v5e, one host is connected to a 2×4 block of TPU chips, giving us 8 PCIe connections between the host and those 8 chips.

Importantly, notice that to reach the DCN backbone, data has to flow through PCIe first, which means DCN communication is even slower than PCIe.

Concretely, cross-pod data flows from the source TPU chip's HBM, over PCIe to the source host, then egresses over the DCN fabric, ingresses into the target host, and finally goes back over PCIe into the target TPU chip's HBM.

This gives us a natural bandwidth hierarchy. The closer we are to the compute die on the TPU chip, the faster the data movement is; the farther we move out into the cluster, the slower it gets.

Let's put the relevant v5e cluster bandwidths on one picture:

Figure 6: Bandwidth hierarchy in a v5e TPU cluster

Now that we have the topology and bandwidth hierarchy in mind, let's look at a few concrete examples of how data actually flows through a TPU slice.

💡Book recommendation:

A few of the examples in this blog, as well as the broader motivation for this post, were inspired by the excellent Scaling Book [3], which I highly recommend.

Suppose we request a 4×4 v5e slice from GCP. Since both axes are smaller than 16, we don't get any wraparound links. So this slice is not a torus, it is a regular 2D grid, and some paths between chips are longer than they would be with wraparound links.

Let's pose the following question: how long does it take to move a (2048, 2048) bf16 matrix from TPU chip (3, 3) to TPU chip (0, 0)?

Figure 7: Moving an 8 MiB bf16 matrix across a 4×4 v5e mesh using two ICI paths

Note that we ignored link latency in the calculation above. In practice, ICI links have roughly 1 μs of latency per hop.

In our example, each path is 6 hops long. Since the two paths run in parallel, this adds roughly 6 μs of latency to the transfer, which is negligible here, but may not be negligible for smaller messages.

This is why it's important to understand whether we are in a latency-bound or throughput-bound regime.

A simple way to estimate this is to ask: assuming we saturate the unidirectional ICI bandwidth of 45 GB/s, how much data can flow through one link in 1 μs?

The answer is:

45 GB/s × 1 μs = 45 KB

So if our message chunks are around this size, latency matters a lot. Ignoring a 1 μs per-hop latency could make the estimate off by a large factor, so the simple bandwidth-only approximation is no longer valid.

Let's do one more example, this time using PCIe, ICI, and HBM → VMEM links as well.

💡what is VMEM?

VMEM is fast on-chip SRAM, roughly equivalent to programmer-managed shared memory on GPUs. It feeds directly into the matmul units; the systolic array, in the case of a TPU chip.

The details of the TPU chip are not important for understanding the main topic of this post, so we won't spend more time on VMEM here.

Assume we have a (128 * 1024, 128 * 1024) bf16 matrix sharded over a 4 × 4 TPU slice. Each chip therefore owns a (32 * 1024, 32 * 1024) submatrix (128/4=32).

Also assume these submatrices have been offloaded to host DRAM.

How long does it take to move all of this data to TPU (0, 0) and do a matmul with a (128 * 1024, 128) bf16 matrix?

Figure 8: Gathering a sharded matrix onto TPU(0,0) for matmul

With the TPU topology and bandwidth hierarchy in mind, we're ready to jump into collective operations.

Let's start with All-Gather.

Inside All-Gather: 1D/2D Rings, and Chains

Let's start with a motivating example.

Figure 9: All-Gather motivation: gathering shards of A so each chip can run the matmul locally

So how do we efficiently implement All-Gather?

A common approach is to use a ring algorithm. Before looking at the algorithm itself, let's first understand where this "ring" comes from:

Figure 10: 1D bidirectional rings naturally appear along both axes of a 16×16 v5e torus

For easier visualization, let's use a smaller ring of 8 TPU chips. In real v5e pods, the wraparound happens at size 16, but for the next few diagrams we'll pretend that an 8-chip row also wraps around:

Figure 11: Simplified representation of an 8-chip ring

Here is All-Gather over a bidirectional 1D ring. It's worth spending a bit of time on this diagram, zoom in and follow one shard as it moves around the ring:

Figure 12: All-Gather over a bidirectional 1D ring

As an exercise, what would the algorithm look like if the ICI links were not full-duplex?

In that case, we could run All-Gather over a unidirectional ring. To keep the diagram small, let's use a hypothetical ring of size N = 4, so the algorithm completes in only N - 1 = 3 steps:

Figure 13: All-Gather over a unidirectional 1D ring

More realistically, if we take a TPU slice without wr

[truncated for AI cost control]