Inside TPU and GPU Clusters: The Anatomy of Collective Communication

TL;DR

A long-form pedagogical walkthrough of the four collective operations that underpin transformer training and inference at scale — All-Gather, Reduce-Scatter, All-Reduce, and All-to-All — grounded in the actual hardware topology of TPU pods and NVIDIA DGX H100 SuperPods. The post first builds a bandwidth hierarchy (HBM → VMEM → ICI → PCIe → DCN on TPUs; HBM → NVLink/NVSwitch → InfiniBand leaf → spine on GPUs), then derives ring-based collective algorithms over 1D/2D torus and fat-tree topologies, contrasts them with tree/recursive-doubling variants, and covers NVSwitch SHARP in-network reduction and multicast. The through-line: collective performance is a topology problem, and reasoning about distributed transformers requires reasoning about which physical links data traverses.

Key claims

TPU chips are directly connected to 4 (v2/v3/v5e/v6e — 2D torus) or 6 (v4p/v5p/TPU7x/8t — 3D torus) nearest neighbors over ICI, forming a wraparound torus at pod scale; the newer 8i inference TPU deviates to a hierarchical high-radix “boardfly” topology [§“TPU cluster topology”].
Requesting a TPU slice smaller than the pod’s wraparound size loses the ring structure along that axis and roughly doubles ring-collective time — e.g., an 8×16 slice on a 16×16 v5e pod pays ~2× on the shorter axis [§“TPU cluster topology”].
The TPU bandwidth hierarchy is monotonic in physical distance: HBM → VMEM (fastest on-chip) → ICI (nearest neighbor) → PCIe (chip-to-host) → DCN (cross-pod, slowest, requires PCIe-DCN-PCIe hop) [§“TPU cluster topology”, bandwidth summary figure].
Latency-bound vs throughput-bound regime is delimited by the per-link “45 GB/s × 1 µs = 45 KB” rule of thumb on v5e ICI: below ~45 KB per hop, per-hop latency dominates and simple bandwidth-only models mispredict transfer time by a large factor [§“TPU cluster topology”].
All-Gather over a bidirectional 1D ring shards the payload and rotates chunks in both directions simultaneously; the dual operation (Reduce-Scatter) uses the same schedule but reduces at each hop instead of copying, so throughput-bound time complexity is identical [§“Inside All-Gather”, §“Reduce-Scatter”].
All-Reduce decomposes as Reduce-Scatter followed by All-Gather, so without hardware support it costs ~2× a single primitive; 2D collectives that exercise both torus axes achieve ~2× speedup over a 1D collective by using all four neighboring ICI links in parallel [§“Reduce-Scatter”].
When a TPU slice loses its ring (no wraparound), collectives fall back to a 1D path/chain algorithm — no algorithmic change beyond removing the wraparound step [§“Inside All-Gather”, §“Reduce-Scatter”].
All-to-All is a sharded transpose: each chip starts grouped by source and ends grouped by destination; in the balanced-routing MoE case (each chip has equal data for every expert), it maps cleanly onto ring/path algorithms, mirroring All-Gather’s structure [§“All-to-All”].
The DGX H100 SuperPod reference architecture is a full fat tree with 8-GPU NVLink/NVSwitch nodes → 32-node Scalable Units connected via InfiniBand leaf switches → SUs connected via spine switches; full bisection bandwidth means any equal split of N nodes can communicate cross-partition at N × 400 GB/s [§“NVIDIA GPU cluster topology”].
Inside a DGX H100 node, NVSwitch provides effective all-to-all connectivity, so the “ring” used by collectives is a logical ordering over the switch fabric, not a physical topology — unlike TPUs where the ring maps directly onto physical ICI links [§“GPU Collectives Within the Node”].
NVSwitch SHARP performs in-network reduction (the switch itself sums partial values instead of routing them back for GPU-side reduction) and multicast for All-Gather; theoretical All-Reduce speedup approaches 2× (ideal ~1.75× on 8-GPU nodes), but in practice NCCL delivers only ~1.3× (~30% improvement) — corroborated independently by Axel Feldmann on 32× H100 IB with NCCL 2.29.7 on 4 GB All-Reduce, and by Pranjal Shankhdhar; NVLink SHARP multicast wastes 1/8 of NVLink bandwidth on the source GPU because multicast still includes it [§“GPU Collectives Within the Node”, References [5], [6]].
Inside an NVSwitch node, balanced All-to-All can be ~2× faster than a bidirectional 1D torus ring at equal bandwidth because every GPU can send directly to every destination without traversing intermediate ring hops [§“GPU Collectives Within the Node”].
Tree/recursive-doubling collectives reduce communication rounds from N−1 to log₂(N) — better for small messages where latency dominates — but achieve lower effective bandwidth on large tensors because they are less pipeline-friendly than rings; NCCL selects among ring, tree, and hybrid algorithms per message size and topology [§“GPU Collectives Within the Node”].
Hierarchical cross-node All-Gather / Reduce-Scatter are well-approximated by T_total ≈ max(D/BW_gpu, D/BW_node) in the pipelined limit, and by T_total ≈ D/BW_node (400 GB/s per H100 node) as a first-order model over IB [§“GPU Collectives Across Nodes”].
In Megatron-style hybrid parallelism, the All-Reduce is not performed across tensor-parallel ranks (they own different shards with nothing element-wise to reduce) — instead, for each fixed TP shard, one All-Reduces across the corresponding data-parallel replicas [§“GPU Collectives Across Nodes”].
Empirically, All-Reduce on H100 nodes needs very large messages (multi-GB) to approach peak bandwidth, and effective bandwidth drops noticeably below ~100 MB; TPUs by contrast approach peak collective bandwidth at ~10 MB message sizes [§“GPU Collectives Within the Node”, Ref [5]].

Method

Not a research paper — a design-principles walkthrough structured around seven sections: TPU cluster topology (superpods, slices, DCN/PCIe/ICI); ring/chain algorithms for All-Gather; the dual construction of Reduce-Scatter and All-Reduce; All-to-All as a sharded transpose; NVIDIA GPU cluster topology (nodes, Scalable Units, fat tree); intra-node collectives (rings, trees, SHARP); and hierarchical inter-node collectives over InfiniBand. Each section proceeds visualization-first — every algorithm is animated as chunk movement across a small worked example — and derives throughput-bound cost models by counting bytes traversing the slowest link in the pipeline. The author uses v5e (16×16 2D torus) as the running TPU example and the DGX H100 SuperPod reference architecture as the GPU example.

Results

Illustrative numbers rather than benchmarks; the useful ones:

TPU v5e ICI: unidirectional 45 GB/s per link. v4 pod: 16×16×16 = 4096 TPUs. v5p pod: 16×20×28 = 8960 TPUs. Latency: ~1 µs per ICI hop; 45 KB break-even between latency and bandwidth regimes [§“TPU cluster topology”].
DGX H100 node: 8 GPUs × 50 GB/s unidirectional IB links = 400 GB/s injection bandwidth per node; NVLink 4 GPUs bisection = 4 × 450 GB/s = 1.8 TB/s unidirectional, 3.6 TB/s bidirectional; SU bisection (16-vs-16 nodes) = 6.4 TB/s uni / 12.8 TB/s bi; full 128-node cluster bisection (64-vs-64) = 25.6 TB/s uni / 51.2 TB/s bi [§“NVIDIA GPU cluster topology”].
SHARP All-Reduce ideal speedup: ~2× in the large-N limit, ~1.75× on 8-GPU nodes; observed on NCCL 2.29.7 on H100 IB: ~1.3× on 4 GB All-Reduce; H100 collectives saturate near ~1 GB message size [§“GPU Collectives Within the Node”, References [5], [6]].
NVSwitch balanced All-to-All ~2× faster than bidirectional 1D ring at equal bandwidth [§“GPU Collectives Within the Node”].

Why it’s interesting

This is the missing pedagogical prerequisite for essentially every training-infra paper on the wiki — the concept page Distributed training parallelism tracks what the parallelism axes are (FSDP, TP, PP, EP, SP) and how they compose non-uniformly at trillion-parameter scale, but it presumes you understand the collective primitives underneath. This post fills that gap and grounds the primitives in physical topology (torus vs fat tree, ICI vs NVLink vs IB), which changes how you reason about failures. The SHARP theoretical-2×-vs-practical-1.3× gap is the specific instance of what Loads and Loads of Fluffy Kittens: Compute-Communication Kernels with Multi-GPU ThunderKittens tackles at the kernel layer — Hazy Research’s ParallelKittens gets GEMM+RS overlap below 1% overhead by picking the right transfer mechanism (TMA vs copy engine vs register/multimem) per message granularity, which is the actionable answer to why NCCL leaves ~40% on the table. Complements Demystifying NVSHMEM: A System-Level Analysis on Symmetric Memory and Device-Initiated Operations in GPU Communication (the NVSHMEM communication runtime underneath) and Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon (the algebraic-structure-exploitation axis, orthogonal to the transfer-mechanism axis this post exposes).