39 Distributed Generative AI

GPT-4 required approximately 25,000 NVIDIA A100 GPUs running continuously for months before a single token could be generated. Stable Diffusion XL consumed 256 A100s in its training run. Gemini Ultra, Google's largest model, reportedly required a cluster of TPU v4 pods containing tens of thousands of chips. These are not engineering curiosities. They are the load-bearing facts beneath every generative AI system you have ever used.

No single machine can hold a trillion-parameter model. The weights of a 1.8-trillion-parameter model in 16-bit floating point occupy approximately 3.6 terabytes of memory - more than 45 times the 80,GB of memory available on a single A100 GPU, the most capable card commercially available at the time these models were trained. Distribution is not a performance optimisation. It is a prerequisite for existence.

Yet the story is richer than mere size. Even if a single chip could hold the weights, training on a trillion tokens of data at a throughput that completes in this decade requires a parallelism that no sequential pipeline can achieve. And even if compute were free, the data movement between GPU memory, CPU memory, NVLink interconnects, and InfiniBand networks would throttle every training step. The three walls - memory, compute, and communication - form a trilemma that defines the entire discipline of distributed generative AI.

This chapter maps that trilemma from first principles. We begin with the arithmetic of scale that forces distribution, move through the hardware that makes it possible, and then develop the algorithmic techniques - data parallelism, tensor parallelism, pipeline parallelism, and their modern hybrids - that close the gap between what physics allows and what engineers deliver. Understanding distributed systems is the difference between using generative AI and building it.

Why Distributed Computing for Generative AI?

The Scale Explosion

The numbers are striking enough that it is worth stating them plainly before drowning them in formalism.

In 2018, BERT-Large had 340 million parameters. In 2019, GPT-2 reached 1.5 billion. In 2020, GPT-3 crossed 175 billion. By 2022, the Chinchilla paper revealed that optimal training required not just larger models but proportionally more data. In 2023, GPT-4 is estimated to contain approximately 1.8 trillion parameters in a mixture-of-experts architecture. The growth rate over five years - from 340 million to 1.8 trillion - represents a factor of roughly 5,000.

The training data has grown in parallel. GPT-2 was trained on 40 gigabytes of web text. GPT-3 consumed roughly 570 gigabytes of filtered data (the Common Crawl corpus after deduplication and quality filtering). LLaMA-3 was trained on 15 trillion tokens, which at roughly 4 bytes per token represents 60 terabytes of raw text.

The compute required tracks with both. If a model with $N$ parameters is trained on $D$ tokens, the approximate number of floating-point operations required is: (Compute Estimate)$$\begin{array}{ccc}& C\approx 6ND& \text{<span class="tml-eqn"></span>}\end{array}$$ where the factor of 6 accounts for the forward pass (2 FLOPs per multiply-add per parameter), backward pass (approximately 2$\times$ the forward), and the gradient update pass. For GPT-3 ($N=175\times {10}^9$ parameters, $D=300\times {10}^9$ tokens), this gives: (GPT3 Compute)$$\begin{array}{ccc}& C_{\text{GPT-3}}\approx 6\times 175\times {10}^9\times 300\times {10}^9\approx 3.14\times {10}^{23}\text{ FLOPs.}& \text{<span class="tml-eqn"></span>}\end{array}$$ An A100 GPU achieves approximately $3.12\times {10}^{14}$ BF16 FLOPs per second (312 teraFLOPS). Training GPT-3 on a single A100 would therefore require: (Single GPU TIME)$$\begin{array}{ccc}& t_{\text{single}}=\frac{3.14\times {10}^{23}}{3.12\times {10}^{14}}\approx 1.0\times {10}^9\text{ seconds}\approx 31.7\text{ years.}& \text{<span class="tml-eqn"></span>}\end{array}$$ GPT-3 was actually trained in roughly one to two months on thousands of A100s. Distribution is not an optimisation - it is the difference between a model that exists and one that does not.

Why Students Should Care

There is a tempting misunderstanding in the practitioner community: that distributed training is the province of infrastructure teams at large technology companies, and that researchers and ML engineers can remain blissfully unaware of how their models actually run. This view is increasingly untenable.

The reason is that algorithmic choices - architecture decisions, optimiser configurations, precision formats, gradient accumulation schemes - are not separable from the distributed substrate on which they execute. A model designed without regard for communication overhead will underperform its theoretical potential by orders of magnitude. A training run scheduled without understanding memory layouts will fail with out-of-memory errors that no amount of hyperparameter tuning will fix. The researcher who understands tensor parallelism can design attention mechanisms that partition efficiently across heads. The one who does not will be bottlenecked by the infrastructure team's defaults.

More practically: the modern toolkit for distributed training - PyTorch FSDP, DeepSpeed ZeRO, Megatron-LM, JAX pmap and pjit - has reached a level of sophistication where using it effectively requires understanding its abstractions. “It ran” is not the same as “it ran efficiently”. A poorly configured FSDP run may achieve 5% hardware utilisation on a cluster costing thousands of dollars per hour.

Data Centre Anatomy

Before formalising the mathematics of distributed training, it is worth grounding the discussion in physical reality. A machine learning data centre is organised in a strict hierarchy, and the bandwidth available at each level of that hierarchy determines which parallelism strategies are viable.

At the top level, a data centre contains rows of racks. Each rack is a 42-unit (42U) or 52U cabinet housing multiple compute nodes. A node is a server: a motherboard with one or more CPUs, DRAM, local storage, and - for AI workloads - a set of GPUs. A modern DGX H100 node, for example, packs 8 H100 GPUs into a single system with a combined GPU memory of 640 gigabytes.

Within a node, GPUs communicate via NVLink, NVIDIA's proprietary high-bandwidth interconnect. In the H100 generation, NVLink provides 900 GB/s of aggregate bidirectional bandwidth between GPUs in the same node - more than 7$\times$ the bandwidth of PCIe Gen 5.

Between nodes, the dominant interconnect is InfiniBand (or, in Google's clusters, the custom ICI interconnect, and in some cloud environments, RoCE over Ethernet). Current-generation NDR InfiniBand provides 400 Gbps (50 GB/s) per port, and nodes are typically connected to a high-radix switch fabric providing non-blocking or near-non-blocking connectivity across hundreds to thousands of nodes.

The ratio of intra-node to inter-node bandwidth is approximately 18:1 (900 GB/s NVLink vs. 50 GB/s InfiniBand per node). This ratio is the single most important fact in distributed AI: it means that algorithms that can keep their most expensive communication within a node will dramatically outperform those that must communicate across nodes.

Formal Definition of Distributed Training

We now lay the mathematical foundations.

The key constraint in Definition Definition 38 is the memory constraint: each worker holds only a fragment of the full computation. This motivates the two fundamental axes of distribution: data parallelism (each worker holds a full copy of the model but processes a different shard of data) and model parallelism (each worker holds a fragment of the model). Real systems use both simultaneously, leading to 3D and 4D parallelism strategies discussed in later sections.

The Three Walls

The central challenge of distributed generative AI can be decomposed into three interacting bottlenecks, each of which is a “wall” that limits achievable throughput.

The Memory Wall.

A model with $N$ parameters stored in BF16 (2 bytes each) requires $2N$ bytes just for the weights. During training, gradient storage adds another $2N$ bytes. The Adam optimiser (the standard choice) adds two momentum terms per parameter, requiring an additional $8N$ bytes in FP32. The total model state memory footprint is: (Memory Footprint)$$\begin{array}{ccc}& M_{\text{state}}=2N+2N+8N=12N\text{ bytes.}& \text{<span class="tml-eqn"></span>}\end{array}$$ For a 7-billion-parameter model, this is $12\times 7\times {10}^9=84\times {10}^9$ bytes $=84$,GB, exceeding the 80,GB memory of a single A100. For GPT-4 at 1.8 trillion parameters: $M_{\text{state}}=12\times 1.8\times {10}^{12}\approx 21.6$,TB - requiring at least 270 A100s for model state alone, before activations.

The Compute Wall.

Even when memory is not the binding constraint, raw throughput limits training speed. The model FLOP utilisation (MFU) measures what fraction of theoretical peak FLOP/s a training job achieves in practice: (MFU)$$\begin{array}{ccc}& \text{MFU}=\frac{\text{observed throughput (FLOPs/s)}}{\text{peak hardware throughput (FLOPs/s)}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ State-of-the-art distributed training typically achieves MFU of 35–55% for large models. The gap between 100% and achieved MFU represents time spent on communication, memory movement, load imbalance, and framework overhead - all of which are targets for the techniques in this chapter.

The Communication Wall.

The most subtle and often most binding wall is communication. In a data-parallel setting with $K$ workers, each worker computes gradients over its local batch and then must aggregate these gradients across all workers before the next parameter update. This all-reduce operation transfers $2N$ gradient values per step (send $+$ receive), and with BF16 gradients requires: (Allreduce Volume)$$\begin{array}{ccc}& V_{\text{allreduce}}=2\times 2N=4N\text{ bytes per step.}& \text{<span class="tml-eqn"></span>}\end{array}$$ For GPT-3 ($N=175\times {10}^9$), this is 700,GB per step. Over a 400,Gbps (50,GB/s) InfiniBand link, the minimum communication time for a single step is: (COMM TIME)$$\begin{array}{ccc}& t_{\text{comm}}\ge \frac{700\text{GB}}{50\text{GB/s}}=14\text{seconds per step.}& \text{<span class="tml-eqn"></span>}\end{array}$$ Since compute for a step with batch size 2048 and sequence length 2048 is on the order of 1–2 seconds per GPU, the communication wall is the dominant bottleneck in naive data-parallel training.

Historical Development

Connection to Generative Models

Generative models have three properties that make distribution especially necessary:

Parameter density. Language models and diffusion models are dense - almost every parameter participates in every forward pass. This contrasts with sparse models (e.g., sparse autoencoders) where most weights are zero or masked. Density means every byte of the model state is active memory, so the memory wall is hit at smaller model sizes than for sparse architectures.

Sequence length and attention. Transformer-based generative models must hold the entire attention context - the key-value cache - in memory for each sequence. At sequence length $L$ with $H$ attention heads each of dimension $d_h$, and $\ell$ layers, the KV cache requires: (KV Cache Memory)$$\begin{array}{ccc}& M_{\text{KV}}=2\times L\times H\times d_h\times \ell \times 2\text{ bytes}& \text{<span class="tml-eqn"></span>}\end{array}$$ (the factor of 2 for key and value, and 2 bytes for BF16). For a 70B model with $L=\mathrm{128,000}$, $H=64$, $d_h=128$, $\ell =80$: $M_{\text{KV}}\approx 3.3\text{TB}$ for a single sequence - already requiring distribution at inference time, not just training.

Iterative denoising in diffusion. Diffusion models generate samples through hundreds of sequential denoising steps. While the model itself may be smaller than a large LLM (Stable Diffusion XL has 2.6 billion parameters), high-resolution generation requires storing intermediate activations at full spatial resolution, and text-to-video models operating at $1024\times 1024$ resolution over 120 frames produce activation tensors measured in tens of gigabytes per forward pass.

Exercises

GPU Architecture and the Accelerator Landscape

The choice of hardware for distributed generative AI is not a commodity decision. Different accelerators differ by an order of magnitude in the bandwidth and compute they provide for specific operations, and these differences determine which parallelism strategies are effective, which precision formats are viable, and ultimately how fast a model can be trained or served. This section develops the architectural understanding needed to make those choices.

GPU vs. CPU: The SIMT Paradigm

A central processing unit (CPU) is designed for latency: it executes a single instruction stream as fast as possible, with out-of-order execution, branch prediction, large caches, and speculative prefetching to minimise the time any single instruction waits. A modern server CPU (e.g., an AMD EPYC 9654) has 96 cores, each running at up to 3.7,GHz with AVX-512 vector units.

A GPU is designed for throughput: it sacrifices the latency-hiding machinery of the CPU in favour of massive parallelism over thousands of simple cores. An NVIDIA H100 GPU has 16,896 CUDA cores organised into 132 Streaming Multiprocessors (SMs), each SM running up to 2048 concurrent threads. The total concurrency is over 270,000 threads.

This architecture is called SIMT: Single Instruction, Multiple Threads. Groups of 32 threads (a warp) execute the same instruction simultaneously. If threads in a warp diverge (take different branches), the hardware serialises the divergent paths, causing warp divergence and reducing throughput. Efficient GPU kernels minimise branch divergence and maximise the fraction of cycles where all warp threads execute the same instruction path.

For matrix multiplication - the dominant operation in transformer training - this architecture is ideal. A matrix multiply $𝐀\in {\mathrm{ℝ}}^{m\times k}$ times $𝐖\in {\mathrm{ℝ}}^{k\times n}$ decomposes into $mn$ independent dot products, each of which can be assigned to a different thread. The regularity of the computation means near-zero warp divergence and near-peak hardware utilisation, which is why matrix multiply is the one operation for which GPUs achieve hardware utilisation close to 90% of theoretical peak.

The GPU Memory Hierarchy

Understanding GPU memory is inseparable from understanding GPU performance. The memory hierarchy has three relevant levels for AI workloads:

High Bandwidth Memory (HBM).

The main GPU memory, implemented as stacked DRAM dies connected to the GPU die via a wide memory interface (typically 5120 bits). The A100 uses HBM2e at 2,TB/s bandwidth and 80,GB capacity. The H100 uses HBM3 at 3.35,TB/s and 80,GB capacity (or HBM3e at 3.35,TB/s and 141,GB in the H200 variant).

Shared Memory / L1 Cache (SRAM).

Each SM has 228,KB of combined L1 cache and shared memory on the H100. SRAM operates at approximately 19–30,TB/s effective bandwidth, roughly 10$\times$ faster than HBM. Kernels that can stage their working data in shared memory - a technique called tiling - achieve dramatically higher throughput. The FlashAttention algorithm, which achieves linear memory complexity for attention, works precisely by tiling the attention computation to fit in SRAM.

Register File.

Each CUDA thread has access to a private register file. The H100 has 65,536 32-bit registers per SM, shared among resident threads. Register spilling - when a kernel uses more registers than available and must store the overflow in slower shared memory - is a common source of unexpected performance degradation in hand-written kernels.

NVIDIA GPU Evolution: V100 to B200

The cadence of NVIDIA GPU releases has been the hardware heartbeat of the AI training era. Each generation brought architectural improvements in matrix computation, memory bandwidth, and interconnect that enabled the next step in model scale.

V100 (Volta, 2017). The V100 introduced Tensor Cores - specialised matrix multiply units that execute a $4\times 4$ matrix multiply per clock per Tensor Core. Peak throughput: 125 TFLOPS in FP16. Memory: 32,GB HBM2 at 900,GB/s. NVLink 2.0: 300,GB/s bidirectional. The V100 was the GPU of choice for GPT-2 and early BERT pre-training runs.

A100 (Ampere, 2020). Introduced BF16 training (the dominant format for LLM training), third-generation Tensor Cores with sparsity support, and NVLink 3.0. Peak throughput: 312 TFLOPS in BF16. Memory: 80,GB HBM2e at 2,TB/s. NVLink 3.0: 600,GB/s bidirectional. The A100 powered GPT-3, PaLM, and the first generation of large-scale diffusion training.

H100 (Hopper, 2022). The Hopper architecture introduced the Transformer Engine, which dynamically selects FP8/BF16 precision per layer, effectively doubling usable compute for transformer workloads. Peak throughput: 989 TFLOPS BF16 (or $\sim$1979 TFLOPS with FP8). Memory: 80,GB HBM3 at 3.35,TB/s. NVLink 4.0: 900,GB/s bidirectional. DGX H100: 8$\times$H100 with full-mesh NVSwitch for 900,GB/s between any pair. The H100 is the primary training chip for GPT-4, Llama-3, and Gemini.

B200 (Blackwell, 2024). The Blackwell B200 introduced a 2-die GPU design (connected by a 10,TB/s NVLink-C2C interconnect), FP4 support for inference, and dramatically expanded memory. Peak throughput: up to 4500 TFLOPS FP8. Memory: 192,GB HBM3e at 8,TB/s per GPU. NVLink 5.0: 1800,GB/s bidirectional. The GB200 NVL72 system connects 72 B200 GPUs in a rack-scale unit with 130,TB/s total NVLink bandwidth, effectively making 72 GPUs behave as a single logical device for many workloads.

NVLink and NVSwitch: Intra-Node Fabric

NVLink is a high-bandwidth, point-to-point interconnect between GPU dies developed by NVIDIA. Unlike PCIe (which routes through the CPU or a PCIe switch), NVLink provides a direct GPU-to-GPU path with much higher bandwidth and lower latency.

In the H100 DGX node, 8 GPUs are connected via NVSwitch, a full-bisection crossbar that provides 900,GB/s between every pair of GPUs simultaneously. The full-bisection property means there is no contention: GPU 0 can simultaneously send to GPU 7 at full bandwidth while GPU 1 sends to GPU 4, GPU 2 sends to GPU 5, and so on. This makes the 8-GPU DGX node behave as a single 640,GB virtual GPU for all-reduce and all-to-all operations.

The architectural implication is profound: tensor parallelism (which requires all-reduce communication at every transformer layer boundary) is practical within a node because the communication cost is amortised over the 900,GB/s NVLink fabric. Extending tensor parallelism across nodes, where the available bandwidth is 50,GB/s per InfiniBand port, is 18$\times$ more expensive in communication time.

Definition: Arithmetic Intensity and the Roofline Model

For the H100 in BF16, $\pi =989\times {10}^{12}$ FLOPs/s and $\beta =3.35\times {10}^{12}$ bytes/s, giving: (H100 Ridge)$$\begin{array}{ccc}& I_{\text{ridge}}^{\text{H100}}=\frac{989\times {10}^{12}}{3.35\times {10}^{12}}\approx 295\text{ FLOPs/byte.}& \text{<span class="tml-eqn"></span>}\end{array}$$ Any kernel with arithmetic intensity below 295 FLOPs/byte is memory-bandwidth-limited on the H100, regardless of how many CUDA cores are available.

The roofline model provides a useful first-pass classification of every operation in a generative model:

The implications of the roofline model for generative AI design are direct. Embedding lookups, layernorm, and element-wise activations all have arithmetic intensity well below the ridge point and are therefore memory-bandwidth-limited. Large matrix multiplications - the core of MLP layers and QKV projection - have high arithmetic intensity and are compute-limited. Naive attention has low intensity (quadratic memory traffic, quadratic compute, but the memory traffic dominates for small batch sizes), while FlashAttention raises intensity by tiling into SRAM.

TPU Architecture: Systolic Arrays and Pods

Google's Tensor Processing Units (TPUs) take a different architectural approach to accelerating matrix computation. Where GPUs use many-core SIMT parallelism, TPUs use systolic arrays - a hardware design that pipes data through a 2D grid of multiply-accumulate units in a wave-like (systolic) pattern, eliminating memory traffic for intermediate results.

A systolic array of size $M\times M$ can compute a $(M\times K)\times (K\times M)$ matrix multiplication by feeding the left matrix row-by-row down the columns and the right matrix column-by-column across the rows, with each cell accumulating partial products. The key property is that each element is read from memory once and reused $M$ times as it travels across or down the array, achieving arithmetic intensity of $O(M)$ - which for $M=128$ (the TPU v4 matrix unit size) gives $\sim$128 FLOPs/byte for matrix multiply.

TPU v4 and v5. The TPU v4 chip achieves 275 TFLOPS in BF16 with 32,GB of HBM. TPU v4 pods connect 4096 chips via a 3D torus ICI (Inter-Chip Interconnect) with 1.2,Tb/s bandwidth between adjacent chips. The ICI topology means that all-reduce across an entire pod is fast: ring-based all-reduce on a 3D torus achieves much lower latency than on a fat-tree InfiniBand network for large cluster sizes. TPU v5p scales to 8960 chips per pod.

Accelerator Comparison

Several observations follow from Table tab:distgen:accelerator-comparison:

• The B200 achieves a $\sim$4.6$\times$ compute advantage over the H100 in BF16, driven primarily by larger Tensor Core arrays and FP8 compute units operating in BF16-equivalent mode.

• The TPU v5p has lower BF16 throughput than the H100 but its 4.8,Tb/s ICI bandwidth (across the pod) enables all-reduce operations that would bottleneck on InfiniBand clusters.

• Intel's Gaudi 3 offers competitive compute and high on-chip interconnect bandwidth, representing a viable alternative for workloads where NVIDIA supply constraints are binding.

• Memory capacity, not raw compute, is often the binding constraint for serving very large models; the B200's 192,GB advantage over the H100's 80,GB is therefore significant for inference workloads.

Example: Training Time for a 7B Model

Exercises

Network Topologies for Distributed Training

Every distributed training system is a graph. Nodes are processing units - GPUs, TPUs, or other accelerators - and edges are the communication links that carry gradients, activations, and parameters between them. The shape of that graph, its topology, determines almost everything that matters for performance: how long all-reduce takes, how much bandwidth is available between any pair of nodes, how many hops a message must traverse before it arrives.

A model parallelism scheme that works beautifully on a ring of eight GPUs may perform catastrophically on a fat-tree network with a hundred nodes, because the communication pattern it generates is fundamentally mismatched to the graph it runs on. Topology awareness is therefore not an implementation detail to be delegated to a network engineer. It is a first-class design principle that shapes every algorithmic choice in distributed deep learning.

This section develops the formal theory of network topologies, surveys the five families that matter most in modern clusters, and establishes the quantitative tools - bisection bandwidth, diameter, cost - that allow principled topology selection. We end with the NVIDIA DGX architecture as a concrete case study that combines multiple topology levels within a single physical system.

Why Topology Matters: The Geometry of Communication

A Motivating Thought Experiment

Imagine 16 GPUs arranged in a line: GPU 0 connected to GPU 1, GPU 1 to GPU 2, and so on up to GPU 15. To broadcast a value from GPU 0 to every other GPU, the message must traverse up to 15 hops. Now rearrange the same 16 GPUs in a $4\times 4$ grid with wraparound edges. The maximum distance between any two GPUs drops to 4. The physical hardware has not changed. The bandwidth of each link has not changed. Only the wiring pattern has changed - yet communication time is cut by more than half.

This is the fundamental insight: topology controls the effective bandwidth and latency available to any distributed algorithm, not just the raw link speeds. A 100 Gbps link is useless if messages must traverse 15 of them in sequence to reach their destination.

Topology-Algorithm Co-Design

Modern distributed training frameworks select collective communication algorithms - all-reduce, broadcast, all-gather, reduce-scatter - based on the topology they detect at runtime. NCCL (NVIDIA Collective Communications Library) uses topology auto-detection to choose between ring, tree, and direct all-reduce patterns. The right algorithm for a ring is not the right algorithm for a fat-tree; running the wrong one can cost $2\times$–$5\times$ in communication time.

The formal tools that enable this co-design begin with a precise definition of the communication graph.

Formal Definitions

The Five Topology Families

Ring Topology

The ring is the simplest topology: each node is connected to exactly two neighbors, forming a cycle. Formally, in a $P$-node ring, the edge set is $E=\{(i,(i+1)modP):i=0,\dots ,P-1\}$.

The ring has diameter $D=\lfloor P/2\rfloor$ and bisection bandwidth $2b$ where $b$ is the per-link bandwidth (the worst cut severs exactly two links). Despite its small bisection bandwidth, the ring supports an elegant all-reduce algorithm - the ring all-reduce - that is bandwidth-optimal in the sense that every link carries exactly $2(P-1)/P$ bytes per byte of model data. We prove this in Proposition 2.

The ring is the topology of choice for NVLink connections within a single DGX node and for the baseline all-reduce in many deep-learning frameworks.

Torus Topology (2D and 3D)

A $k$-dimensional torus arranges $P=n^k$ nodes in a $k$-dimensional grid with wraparound edges in every dimension. In a 2D torus with $\sqrt{P}\times \sqrt{P}$ nodes:

• each node $(i,j)$ is connected to $(i\pm 1mod\sqrt{P},j)$ and $(i,j\pm 1mod\sqrt{P})$;

• diameter $D=2\cdot \lfloor \sqrt{P}/2\rfloor$;

• bisection bandwidth $2\sqrt{P}\cdot b$ (two rows of links are cut).

The 3D torus - used in the IBM Blue Gene series and the Google TPU Pod interconnect - further reduces diameter and increases bisection bandwidth at the cost of higher wiring complexity. Each node has 6 links (two per dimension), the diameter is $3\lfloor P^{1/3}/2\rfloor$, and the bisection bandwidth is $2P^{2/3}\cdot b$.

Fat-Tree Topology

The fat-tree, introduced by Charles Leiserson in 1985 and popularized in the Thinking Machines CM-5 and modern InfiniBand clusters, achieves full bisection bandwidth: the total bandwidth across any bisection equals the total bandwidth at the leaves.

A $k$-ary fat-tree has three levels of switches.

• $k^2/4$ edge (leaf) switches, each with $k/2$ downlinks to compute nodes and $k/2$ uplinks.

• $k^2/4$ aggregation switches.

• $k^2/4$ core switches.

Total nodes: $k^3/4$. Diameter: $2{\log }_{k/2}P+2$ hops. Bisection bandwidth: equal to total endpoint bandwidth (ideal).

In practice, fat-trees are often deployed with oversubscription ratios of 2:1 or 4:1 to reduce cost: uplinks carry less bandwidth than downlinks. A $k$-ary fat-tree with $r$:1 oversubscription has bisection bandwidth $\mathrm{BB}=P\cdot b/(2r)$ where $b$ is the per-endpoint bandwidth.

Fat-trees are standard in InfiniBand HPC clusters and are used by cloud providers including AWS (EFA networks) and Azure (InfiniBand for A100/H100 nodes).

Dragonfly Topology

The dragonfly topology, introduced at SC'08, targets very large systems ($P>\mathrm{10,000}$ nodes) where the cost of full-bisection fat-trees becomes prohibitive.

A dragonfly organizes nodes into groups (router chassis), each connected in an all-to-all pattern within the group. Groups are then connected by a smaller number of global links. This two-level hierarchy achieves:

• diameter 3 (host$\to$local router$\to$global link$\to$ destination router$\to$destination host);

• bandwidth cost roughly half that of a fat-tree at the same scale, by reducing the number of global high-cost links;

• the need for adaptive routing to avoid global link congestion when many pairs communicate simultaneously.

The Cray Slingshot interconnect used in the Frontier and Aurora exascale supercomputers implements a dragonfly+ topology.

Hypercube Topology

The $d$-dimensional hypercube connects $P=2^d$ nodes so that two nodes are adjacent if and only if their binary addresses differ in exactly one bit. Properties:

• diameter $D=d={\log }_{2}P$ - the shortest diameter possible for any topology connecting $P$ nodes with the same node degree;

• node degree $d={\log }_{2}P$ (increases with $P$, limiting scalability);

• bisection bandwidth $P/2\cdot b$ (the minimum cut severs $P/2$ links).

Hypercubes were prominent in early parallel computers (Intel iPSC, nCUBE) and remain theoretically important because many collective algorithms - butterfly all-reduce, recursive halving-doubling - map naturally onto hypercube communication patterns and can be emulated on fat-trees.

TikZ Figure: Five Topologies

Ring All-Reduce: Cost Analysis

The ring all-reduce algorithm proceeds in two phases. Let $P$ be the number of nodes, each holding a gradient vector of $n$ bytes.

• Reduce-scatter phase: $P-1$ rounds. In round $r$, each node $i$ sends its chunk of size $n/P$ to node $(i+1)modP$ and receives a chunk from node $(i-1)modP$, accumulating the partial reduction. Each link carries $n/P$ bytes per round.

• All-gather phase: $P-1$ rounds. Each node $i$ forwards its fully-reduced chunk to node $(i+1)modP$. Again, each link carries $n/P$ bytes per round.

Topology Comparison

tab:distgen:topology-comparison summarizes the key properties of the five topology families.

NVIDIA DGX Architecture: A Multi-Level Topology

The NVIDIA DGX A100 system exemplifies how modern AI hardware implements a hierarchy of topologies within a single product.

Within the Node: NVSwitch All-to-All

A single DGX A100 contains 8 NVIDIA A100 GPUs interconnected by 6 NVSwitch chips. Unlike a ring (where each GPU connects to 2 neighbors) or a tree, NVSwitch implements a non-blocking all-to-all switch fabric: every GPU can simultaneously communicate with every other GPU at full bandwidth. This is topologically equivalent to a complete graph $K_8$ on the GPU level.

• NVLink 3.0 bandwidth: 600 GB/s bidirectional per GPU.

• NVSwitch total bisection bandwidth: $8\times 300$ GB/s = 2.4 TB/s.

• All-reduce of $n$ bytes across 8 GPUs: $T\approx \alpha +n{\beta }_{\text{NVLink}}$ (single-step, no multi-hop bottleneck).

Between Nodes: InfiniBand Fat-Tree

Multiple DGX nodes are connected via InfiniBand HDR (200 Gb/s per port) in a fat-tree or dragonfly topology. Each DGX A100 has 8 InfiniBand ports, one per GPU, providing aggregate inter-node bandwidth of $8\times 25$ GB/s = 200 GB/s.

The bandwidth mismatch is stark:

• Intra-node (NVLink): 600 GB/s per GPU.

• Inter-node (InfiniBand): 25 GB/s per GPU.

This $24\times$ mismatch means that all-reduce operations that cross node boundaries are 24 times slower per byte than operations that stay within a node. This is why tensor parallelism (which requires all-to-all within a node) is placed on NVLink-connected GPUs, while pipeline and data parallelism (which require less frequent, larger messages) are placed across InfiniBand-connected nodes.

The DGX SuperPOD

A DGX SuperPOD connects 32 DGX A100 nodes ($32\times 8=256$ GPUs) via a fat-tree InfiniBand fabric. The full system delivers $256\times 312$ TFLOPS = 79.9 PFLOPS of FP16 compute with $256\times 25$ GB/s = 6.4 TB/s of inter-node bandwidth.

Example: All-Reduce Through a Fat-Tree

Key Idea: Topology as Constraint

Exercises

The Alpha-Beta Communication Model

The Fundamental Question

How long does it take to send a message from one GPU to another?

This question sounds simple, but it underlies every performance prediction, algorithm design decision, and system tuning effort in distributed deep learning. The answer depends on two completely separate physical phenomena that happen to coexist in every communication link:

• Latency: the time that elapses before even the first bit of your message arrives at the destination. This is governed by the speed of light, switch forwarding delays, software stack overhead, and network protocol negotiation. It does not depend on how much data you are sending.

• Bandwidth: the rate at which data flows once the link is active. This is governed by the physical signalling speed of the link and is (approximately) fixed by the hardware.

These two quantities are independent. A 400 Gbps InfiniBand link has extraordinary bandwidth but still incurs $\sim 1\mu \mathrm{s}$ of hardware latency plus $\sim 2\mu \mathrm{s}$ of software (MPI/NCCL) overhead. Conversely, a 1 Gbps Ethernet link may have the same latency but 400$\times$ less bandwidth. An algorithm that sends millions of tiny messages may be latency-bound even on a 400 Gbps fabric; an algorithm that sends a few enormous messages may be bandwidth-bound even on a high-latency WAN.

The Train Analogy

Before formalizing the model, consider a vivid physical analogy that makes the two terms immediately intuitive.

Imagine you need to move a large crowd of passengers from city A to city B by train. The journey has two completely distinct phases:

1. Waiting for the first coach to arrive at the platform. No matter how many passengers you have, no matter how long the train is, you must wait for the engine and the first coach to pull into the station. This waiting time is fixed: it depends on the distance between cities, the speed of the engine, and the time the driver takes to respond to the departure signal. It does not change whether you have 10 passengers or 10,000 passengers. This is latency: $\alpha$.

2. Passengers disembarking once the train has arrived. Once the first coach is at the platform, passengers begin streaming off. If the doors are wide and the platform is long, passengers disembark quickly (high bandwidth, small $\beta$). If the doors are narrow or the coaches are long, disembarkation takes longer (low bandwidth, large $\beta$). And crucially: the more passengers there are, the longer this takes. Time $\propto$ (number of passengers) $\times \beta$. This is the bandwidth contribution: $N\beta$.

Total journey time = waiting for the first coach ($\alpha$) + time for all passengers to disembark ($N\beta$) = $\alpha +N\beta$.

Once the first coach arrives, the rest follow quickly - coach after coach at a steady rate determined by track capacity (bandwidth). That is why reducing latency matters more than you might think: if you are sending many small messages, the $\alpha$ term dominates and no amount of bandwidth improvement will help. Conversely, for a single enormous transfer, $\alpha$ is a one-time fixed cost amortized over $N\beta$, and bandwidth is what matters.

Formal Definition

Point-to-Point and Collective Costs

Point-to-Point Transfer

By definition, a single point-to-point message of $N$ bytes costs: (P2P)$$\begin{array}{ccc}& T_{\text{p2p}}(N)=\alpha +N\beta .& \text{<span class="tml-eqn"></span>}\end{array}$$

Broadcast

A broadcast sends a message of $N$ bytes from one root node to all $P-1$ other nodes. The optimal tree-based implementation proceeds in $\lceil {\log }_{2}P\rceil$ rounds, doubling the number of informed nodes each round. In round $r$, every node that has the data sends it to one new node.

(Bcast)$$\begin{array}{ccc}& T_{\text{bcast}}(N)=\lceil {\log }_{2}P\rceil \cdot (\alpha +N\beta ).& \text{<span class="tml-eqn"></span>}\end{array}$$

Derivation. There are $\lceil {\log }_{2}P\rceil$ rounds. In each round, every active sender transmits the full $N$-byte message to one new receiver: cost $\alpha +N\beta$ per round. Rounds are sequential, giving (Bcast).

All-Reduce: Butterfly vs Ring

Butterfly (Recursive Halving-Doubling).

The butterfly all-reduce, natural on a hypercube, proceeds in $\lceil {\log }_{2}P\rceil$ rounds. In round $r$, each node exchanges half its data with a partner at distance $2^r$. The total cost is: (AR Butterfly)$$\begin{array}{ccc}& T_{\text{AR,butterfly}}(N)=2\lceil {\log }_{2}P\rceil \cdot \alpha +2\frac{P-1}{P}N\beta .& \text{<span class="tml-eqn"></span>}\end{array}$$ The factor $2(P-1)/P$ in the bandwidth term reflects that each node effectively sends and receives the full message $N$ once (total $2N$ bytes), reduced by $1/P$ because its own chunk needs no transmission.

Ring.

From Proposition 2: (AR RING)$$\begin{array}{ccc}& T_{\text{AR,ring}}(N)=2(P-1)\alpha +2\frac{P-1}{P}N\beta .& \text{<span class="tml-eqn"></span>}\end{array}$$

Comparison.

The bandwidth terms of and are identical. The difference lies entirely in the latency term: $2\lceil {\log }_{2}P\rceil \alpha$ (butterfly) versus $2(P-1)\alpha$ (ring). For large $P$, the butterfly has $O(\log P)$ latency rounds while the ring has $O(P)$ latency rounds.

• If $\alpha \gg N\beta /P$, i.e., we are latency-bound, use the butterfly.

• If $N\beta /P\gg \alpha$, i.e., we are bandwidth-bound, both algorithms perform identically; use the ring (simpler).

In deep learning, gradient all-reduce involves large tensors ($N\sim {10}^8$–${10}^{10}$ bytes) and $P$ is modest (8–1024 GPUs), so $N\beta \gg P\alpha$ and we are almost always bandwidth-bound. Ring all-reduce is the standard choice.

Communication Costs of Common Collectives

tab:distgen:collective-costs summarizes the alpha-beta cost formulas for the most common collective operations.

Bandwidth Optimality of Ring All-Reduce

Timeline Diagram: Latency vs Bandwidth

fig:distgen:timeline illustrates how the alpha-beta model decomposes communication time into latency and bandwidth contributions for three different message sizes.

Crossover Point and Regime Analysis

The crossover between the latency-dominated and bandwidth-dominated regimes occurs at: (Crossover)$$\begin{array}{ccc}& N^{\ast }=\frac{\alpha }{\beta }.& \text{<span class="tml-eqn"></span>}\end{array}$$ For $N\ll N^{\ast }$, the latency term $\alpha$ dominates and the effective throughput $N/T(N)\approx N/\alpha$ is far below the peak bandwidth $1/\beta$. For $N\gg N^{\ast }$, the bandwidth term dominates and throughput approaches $1/\beta$.

Example: All-Reduce for a 7B Model

Insight: Latency-Bound vs Bandwidth-Bound

Exercises

Message Passing Interface (MPI) Tutorial

Every framework that trains large language models across hundreds of GPUs ultimately rests on a single substrate: the ability of one process to send a buffer of bytes to another process on a different machine, reliably, at high speed, and with well-defined semantics. The Message Passing Interface-MPI-is the standard that codifies exactly this capability. It is not glamorous, and it predates transformers by three decades. But understanding MPI is to understand the grammar of distributed computation: the primitives from which all higher-level abstractions, NCCL, PyTorch Distributed, DeepSpeed's communication engine, are ultimately assembled.

The MPI Execution Model

MPI follows the Single Program, Multiple Data (SPMD) model. Every participating process runs the same program, but they diverge in behavior because each knows its own identity-its rank- within a communicator.

When a program calls MPI_Comm_rank( MPI_COMM_WORLD, &rank), it learns its position in the process group. When it calls MPI_Comm_size( MPI_COMM_WORLD, &size), it learns the total number of participants. The canonical “hello world” in MPI-stripped to its essence-is:

MPI_Init(&argc, &argv);
MPI_Comm_rank(MPI_COMM_WORLD, &rank);
MPI_Comm_size(MPI_COMM_WORLD, &size);
printf("I am rank 
MPI_Finalize();

Between MPI_Init and MPI_Finalize, all MPI communication is valid. Outside those calls, no MPI function may be used. The runtime launches $p$ copies of this process, each printing a different rank, producing $p$ lines of output in an unspecified order since message passing does not impose a global ordering on independent events.

Point-to-Point Communication

The fundamental communication primitive in MPI is the point-to-point message: one process sends, another receives.

The blocking send and receive have signatures:

int MPI_Send(const void *buf, int count, MPI_Datatype datatype,
             int dest, int tag, MPI_Comm comm);

int MPI_Recv(void *buf, int count, MPI_Datatype datatype,
             int source, int tag, MPI_Comm comm,
             MPI_Status *status);

A concrete example: rank 0 sends a parameter tensor to rank 1.

if (rank == 0) {
    double params[N];
    // ... fill params ...
    MPI_Send(params, N, MPI_DOUBLE, 1, 0, MPI_COMM_WORLD);
} else if (rank == 1) {
    double params[N];
    MPI_Recv(params, N, MPI_DOUBLE, 0, 0,
             MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    // ... use params ...
}

The wildcard MPI_ANY_SOURCE and MPI_ANY_TAG allow a receiver to accept from any sender or with any tag, useful for work-stealing patterns where a coordinator accepts results from whichever worker finishes first.

Blocking vs. Non-Blocking Communication

The calls MPI_Send and MPI_Recv are blocking: they do not return until the operation is complete from the calling process's perspective. For MPI_Send, “complete” means the send buffer is safe to reuse (either the message has been copied into a system buffer or the receiver has begun receiving). For MPI_Recv, “complete” means the data has arrived in the receive buffer.

While blocking communication is easy to reason about, it leaves CPU and GPU cycles idle during the communication phase, exactly when we want to be computing. Non-blocking operations solve this by decoupling the initiation of communication from its completion.

MPI_Request req;

// Initiate send - returns immediately
MPI_Isend(buf, N, MPI_DOUBLE, dest, tag,
          MPI_COMM_WORLD, &req);

// Do computation here while data is in flight
compute_next_layer(...);

// Wait for send to complete before touching buf again
MPI_Wait(&req, MPI_STATUS_IGNORE);

The MPI_Request handle is the key. It represents the pending operation; MPI_Wait blocks until that specific request completes. Multiple requests can be waited on simultaneously:

MPI_Request reqs[2];
MPI_Isend(send_buf, N, MPI_DOUBLE, dest, 0,
          MPI_COMM_WORLD, &reqs[0]);
MPI_Irecv(recv_buf, N, MPI_DOUBLE, src, 0,
          MPI_COMM_WORLD, &reqs[1]);
// Overlap with computation
compute_gradients(...);
// Complete both
MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE);

Figure fig:distgen:blocking-nonblocking visualises the critical difference. In the blocking scenario, every training step has a serial “send, wait, receive” phase during which the accelerators sit idle. In the non-blocking scenario, the data transfer happens in the background of a backward pass, so the next layer's gradient computation overlaps with the previous layer's gradient synchronization. This is not a minor optimisation: at the scale of a 70B-parameter model trained on 1000 GPUs, the communication volume is so large that blocking all-reduces would reduce GPU utilization to below 40%.

Collective Operations

Point-to-point operations are the atoms of MPI. Collectives are the molecules: coordinated multi-process operations that implement common communication patterns in a single call, with implementations optimised for the underlying network topology. All processes in the communicator must call the collective.

Broadcast (MPI_Bcast). One process (the “root”) has data that every other process needs.

MPI_Bcast(buffer, count, datatype, root, comm);

After the call, every rank holds the same data. In distributed training, broadcast is used to initialize all workers with the same model parameters before training begins. The naive implementation sends $p-1$ point-to-point messages from root; an optimized implementation uses a binomial tree, reducing latency from $O(p)$ to $O(\log p)$.

Reduce (MPI_Reduce). Every process has local data; a reduction operation (sum, max, min, logical AND, etc.) combines them, and the result lands at the root.

MPI_Reduce(sendbuf, recvbuf, count, datatype,
           MPI_SUM, root, comm);

Only the root's recvbuf is meaningful after the call. In gradient aggregation, each worker computes local gradients ${\nabla }_{i}L$, and the root obtains ${\sum }_{i=0}^{p-1}{\nabla }_{i}L$.

All-Reduce (MPI_Allreduce). Like Reduce, but the result is distributed to all processes, not just the root. This is the workhorse of data-parallel training.

MPI_Allreduce(sendbuf, recvbuf, count, datatype,
              MPI_SUM, comm);

After MPI_Allreduce, every rank holds $\bar{g}=\frac{1}{p}{\sum }_{i=0}^{p-1}{\nabla }_{i}L$ (dividing by $p$ is done in the application code). Ring all-reduce, the algorithm underlying NCCL and PyTorch DDP, achieves near-optimal bandwidth of $2(p-1)/p$ of peak network throughput.

All-Gather (MPI_Allgather). Every process contributes a piece; every process ends up with all pieces concatenated. In ZeRO optimizer (Zero Redundancy Optimizer), model parameters are sharded across GPUs; before each forward pass, Allgather reconstructs the full parameter tensor on every GPU.

MPI_Allgather(sendbuf, sendcount, sendtype,
              recvbuf, recvcount, recvtype, comm);

Scatter (MPI_Scatter). The inverse of gather: root has a large array and distributes distinct chunks to each process.

MPI_Scatter(sendbuf, sendcount, sendtype,
            recvbuf, recvcount, recvtype, root, comm);

Figure fig:distgen:collectives shows the data flow for each collective. These five primitives, together with their blocked variants and the MPI_Reduce_scatter operation (reduce followed by scatter of the result), cover the vast majority of communication patterns in distributed deep learning.

Table tab:distgen:collectives-cost summarises the asymptotic communication cost of each collective on a $p$-process ring with per-byte bandwidth $\beta$ and per-message latency $\alpha$.

One-Sided Communication and RDMA

Classical two-sided communication (Send/Recv) requires both the sender and receiver to participate actively. In large-scale training, this can create synchronization bottlenecks: a parameter server must explicitly post receives before workers can push updates. One-sided communication removes this requirement.

One-sided communication maps directly onto RDMA: Remote Direct Memory Access. RDMA is a hardware capability, not an MPI abstraction: it allows a network interface card (NIC) to perform DMA operations directly into a remote machine's memory without involving either machine's CPU. The transfer bypasses the operating system kernel, the protocol stack, and the remote CPU, achieving latencies under $1\mu \text{s}$ on InfiniBand networks and bandwidths exceeding $200\text{Gb/s}$ on HDR InfiniBand.

Two network technologies enable RDMA at scale:

• InfiniBand: a purpose-built interconnect developed in the late 1990s that natively supports RDMA, with current HDR InfiniBand achieving $200\text{Gb/s}$ per port and $O(100\text{ns})$ latency. Mellanox (acquired by NVIDIA) and Intel manufacture InfiniBand adapters. Every major AI supercomputer cluster-Summit, Frontier, NVIDIA DGX SuperPOD, and most cloud HPC clusters-uses InfiniBand.

• RoCE (RDMA over Converged Ethernet): a standard that runs the InfiniBand transport layer (IB BTH headers) over standard Ethernet, making RDMA accessible on commodity switching hardware. RoCEv2 runs over UDP/IP, enabling RDMA across routed Ethernet networks. Cloud providers (AWS EFA, Azure RDMA, Google's Jupiter network) implement RoCE-style RDMA at scale. The primary operational challenge is loss sensitivity: RDMA queues require lossless transport, so Ethernet must be configured with Priority Flow Control (PFC).

Example: Gradient All-Reduce for LLM Training

We now build a minimal but complete gradient synchronization loop using MPI collectives, illustrating the pattern that underlies every data-parallel training framework.

Suppose we are training a model with $D$ parameters across $p=8$ workers. Each worker holds an identical copy of the model and a different shard of the training data.

// Each rank owns a full copy of params[D] and grads[D]
double params[D], grads[D];

for (int step = 0; step < num_steps; step++) {
    // 1. Forward pass on local data shard
    double loss = forward(params, local_data[step]);

    // 2. Backward pass: compute local gradient
    backward(params, local_data[step], grads);  // fills grads[]

    // 3. Synchronize gradients across all workers
    //    Each rank contributes grads[], receives the sum
    MPI_Allreduce(MPI_IN_PLACE,   // in-place: result overwrites input
                  grads,
                  D,
                  MPI_DOUBLE,
                  MPI_SUM,
                  MPI_COMM_WORLD);

    // 4. Scale by 1/p to get mean gradient
    for (int i = 0; i < D; i++) grads[i] /= size;

    // 5. Update parameters (SGD)
    for (int i = 0; i < D; i++)
        params[i] -= lr * grads[i];
    // All ranks now have identical params for next step
}

Exercises

NCCL and GPU Communication Libraries

MPI was designed in an era when the compute node was a CPU and memory was flat RAM. It assumes that communication is managed by the CPU and that data lives in host memory. Neither assumption holds in modern GPU clusters. In a DGX H100 node, eight H100 GPUs are connected to each other via NVLink, a dedicated high-speed GPU interconnect that bypasses the CPU entirely and achieves $900\text{GB/s}$ of bidirectional bandwidth. Between nodes, GPUs communicate over InfiniBand through RDMA.

Asking the CPU to orchestrate GPU-to-GPU communication in this environment is like asking a bicycle courier to route messages in a city that has its own pneumatic tube network: the courier adds latency, CPU memory copies, and becomes a bottleneck. NCCL (NVIDIA Collective Communications Library) is NVIDIA's answer: a library that runs collective operations directly on the GPU, using NVLink for intra-node transfers and GPU-Direct RDMA for inter-node transfers, with the CPU doing nothing more than launching the operation.

Why MPI Is Insufficient for GPU Workloads

Standard MPI implementations are not GPU-aware. When an MPI application calls MPI_Send(gpu_ptr, ...), the implementation typically:

1. Copies the GPU buffer to a CPU staging buffer (a synchronous D2H memcopy that stalls the GPU);

2. Sends the data from CPU memory over the network;

3. Copies the received data back to a GPU buffer (H2D memcopy).

Each of these steps is expensive. A 1,GB gradient tensor at 900,GB/s NVLink bandwidth transfers in $\sim$1,ms NVLink-to-NVLink, but the PCIe bus between GPU and CPU saturates at $\sim$32,GB/s, making the round-trip copy take $\sim$31,ms-a $31\times$ slowdown.

GPU-aware MPI implementations (OpenMPI with CUDA support, Cray MPICH with GPU support) address this by detecting GPU pointers and routing via RDMA directly, but they still rely on the CPU thread to initiate every communication operation, creating synchronization overhead.

NCCL's key architectural decision: the communication kernels run on the GPU. The CPU posts an NCCL operation and returns immediately; the GPU executes the collective in a stream, overlapping with other GPU kernels scheduled in the same stream or in concurrent streams.

NCCL Architecture and Evolution

NCCL has undergone substantial algorithmic evolution since its first public release (NCCL 1.x, 2015–2016).

NCCL 1.x: Flat ring.

All-reduce was implemented as a ring all-reduce with a single ring topology connecting all GPUs. Within a node, the ring traversed GPUs over PCIe or NVLink; between nodes, data went over InfiniBand. This achieved near-optimal bandwidth for large messages but suffered from $O(p)$ latency for small messages (latency scales with ring circumference).

NCCL 2.x: Tree-based algorithms.

To address small-message latency, NCCL 2.x introduced double-binary tree collectives. A binomial or binary tree achieves $O(\log p)$ latency versus $O(p)$ for a ring, making tree algorithms preferable when the message fits in a few packets. NCCL 2.x uses a dual-tree construction (two complementary binary trees) to keep all links utilized even in tree topology, avoiding the bandwidth underutilization of a single tree.

NCCL 2.x+: Hybrid algorithms.

Modern NCCL selects the algorithm based on message size, number of ranks, and network topology: ring for large messages (bandwidth-bound), tree for small messages (latency-bound), with tuning tables built from profiling on target hardware. The transition point is typically around $64\text{KB}$–$256\text{KB}$ per rank.

NCCL and NVLink.

Within a DGX node, NCCL detects the NVLink topology (an all-to-all mesh for H100, a hypercube for A100) and chooses ring orderings that maximise NVLink utilization. For 8 H100s connected by NVLink4, NCCL can achieve $\sim$700,GB/s all-reduce bandwidth, approaching the NVLink bisection bandwidth.

Ring All-Reduce: The Workhorse Algorithm

Ring all-reduce is the algorithm behind the majority of gradient synchronization in production training. We now develop it formally.

Tree All-Reduce: Latency-Optimal for Small Messages

Ring all-reduce requires $O(p)$ communication steps, each taking at least one network latency $\alpha$. For small messages (a few kilobytes of gradients from a lightweight layer), the total latency is $2(p-1)\alpha$, which grows linearly with the number of GPUs. At $p=1000$ GPUs with $\alpha =1\mu \text{s}$, that is $\sim$2,ms just in latency overhead before any data moves.

A tree all-reduce uses a binary or binomial tree topology:

1. Reduce phase: leaves send to their parent; each internal node accumulates children's values and passes the sum upward. After $O(\log p)$ steps, the root holds the global sum.

2. Broadcast phase: root sends the global sum down the tree. After another $O(\log p)$ steps, all nodes have the result.

Total latency: $O(\log p)$ steps, reducing the latency from 2,ms to $\sim 20\mu \text{s}$ at $p=1000$. The cost: bandwidth efficiency drops because internal tree nodes are bottlenecks-the root must send and receive $D$ bytes, consuming full bandwidth, while leaves send and receive only $D/p$ bytes. Tree all-reduce is therefore optimal for latency-bound (small $D$) workloads and ring all-reduce is optimal for bandwidth-bound (large $D$) workloads.

The crossover point $D^{\ast }$ satisfies: (Crossover)$$\begin{array}{ccc}& \frac{D^{\ast }}{B}=\alpha \log p⟹D^{\ast }=\alpha B\log p,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $B$ is the single-link bandwidth. For InfiniBand ($\alpha =1\mu \text{s}$, $B=25\text{GB/s}$) with $p=256$: $D^{\ast }={10}^{-6}\times 25\times {10}^9\times {\log }_{2}256=200\text{KB}$. Messages below $200\text{KB}$ prefer tree; larger prefer ring. This matches NCCL's empirical tuning thresholds.

PyTorch DistributedDataParallel and NCCL

PyTorch's DistributedDataParallel (DDP) is the standard interface for data-parallel training in Python. Its design is instructive because it exposes exactly the communication patterns we have formalized.

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# Initialize NCCL process group
dist.init_process_group(backend="nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()
device = torch.device(f"cuda:{rank}")

# Wrap model in DDP - hooks registered at init time
model = MyLLM().to(device)
ddp_model = DDP(model, device_ids=[rank])

optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=1e-4)

for step, (inputs, labels) in enumerate(dataloader):
    inputs, labels = inputs.to(device), labels.to(device)
    optimizer.zero_grad()

    # Forward (no communication)
    outputs = ddp_model(inputs)
    loss = criterion(outputs, labels)

    # Backward: DDP hooks fire as each parameter's gradient
    # is computed, launching async NCCL all-reduce for that bucket
    loss.backward()

    # By here, NCCL all-reduces are complete (DDP waited)
    optimizer.step()

Gloo: CPU-Based Alternative

Not all distributed training happens on GPU clusters. CPU-only deployments (federated learning on edge nodes, small models on commodity hardware, debugging environments without GPU) need an alternative to NCCL. Gloo is Facebook AI Research's CPU-based collective communications library, integrated into PyTorch as a backend for torch.distributed.

Gloo implements the same all-reduce semantics as NCCL but over standard Ethernet (TCP/IP) without RDMA. It uses a halving-doubling all-reduce algorithm (also called recursive halving-and-doubling or butterfly all-reduce):

• In ${\log }_{2}p$ steps, pairs of processes exchange half their data and accumulate;

• After ${\log }_{2}p$ steps, every process holds the global sum.

Gloo also supports point-to-point sends and receives, making it a drop-in replacement for NCCL in environments where GPUs are absent.

The practical usage difference is a single argument:

dist.init_process_group(backend="gloo")  # CPU
dist.init_process_group(backend="nccl")  # GPU

Comparison: NCCL, MPI, and Gloo

Table tab:distgen:backend-comparison makes the use-case boundaries clear. For LLM and diffusion model training on GPU clusters, NCCL is the unambiguous choice. For HPC simulation workloads where some GPUs and some CPU-only nodes mix in the same job, GPU-aware MPI can span both. For federated learning over wide-area networks, for rapid prototyping, or for environments where CUDA is unavailable, Gloo provides a friction-free fallback.

Connection to Generative AI: Every Gradient Sync Uses These Primitives

It is worth pausing to appreciate the scale at which these primitives operate in modern generative AI.

A GPT-4 scale training run involves an estimated $\sim$25,000 A100 GPUs running for months. At each optimizer step (every $\sim$1 second of wall time at typical batch sizes), NCCL performs all-reduce over gradients for hundreds of billions of parameters. The communication volume per step exceeds $1\text{TB}$ of data across the cluster. Keeping the GPUs busy for more than $50\%$ of wall time requires the pipelining, overlapping, and topology-aware routing that NCCL provides.

A Stable Diffusion 3 training run at scale uses tensor parallelism (splitting individual attention matrices across GPUs) in addition to data parallelism, requiring not just all-reduce but also all-to-all collectives (for rearranging activations between tensor-parallel and sequence-parallel regions). ncclAllToAll and ncclAllGather appear in the inner loop of every forward pass.

A mixture-of-experts (MoE) model introduces a unique communication pattern: the expert dispatch operation. Each token is routed to $k$ experts, which may reside on different GPUs. The tokens must be shipped to the expert GPUs (a sparse all-to-all), processed, and then the results shipped back. This is an irregular collective that standard NCCL's ncclAllToAll can handle when the routing table is known in advance, but requires custom kernels for truly dynamic routing.

Example: Profiling NCCL All-Reduce on 8 H100 GPUs

# Build and run nccl-tests all-reduce benchmark
$ ./build/all_reduce_perf -b 8 -e 2G -f 2 -g 8

# NCCL version: 2.18.3
# nThread 1 nGpus 8 minBytes 8 maxBytes 2G step factor 2
# Using devices
#   Rank  0 Pid 12345 on dgx01 device  0 [0x18] NVIDIA H100 SXM5 80GB
#   ...   (8 GPUs total)
#
#                                              out-of-place
#       size    count  type    redop    time   algbw   busbw
#        (B)                             (us)  (GB/s)  (GB/s)
           8        2  float    sum     14.0    0.00    0.00
        1024      256  float    sum     17.0    0.06    0.11
       65536    16384  float    sum     22.0    2.98    5.21
     1048576   262144  float    sum     38.0   27.59   48.28
    16777216  4194304  float    sum    157.0  106.86  186.99
   134217728 33554432  float    sum    961.0  139.66  244.40
  1073741824 268435456  float   sum   7210.0  148.92  260.61

Exercises

Communication Hiding - The Art of Multitasking

There is a moment in every distributed training run when a layer's gradients have been computed and are waiting - waiting for the all-reduce to finish, waiting for tensors to travel across InfiniBand or NVLink, waiting for the ring-reduce to complete its circuit. During that wait, the GPUs are idle. Their thousands of CUDA cores sit dark. The power draw does not drop, the memory stays allocated, the clocks keep ticking - but no useful arithmetic is happening.

This is the central inefficiency that communication hiding is designed to eliminate. The idea is disarmingly simple: instead of computing, then communicating, then computing again - overlap the two. Arrange the work so that while the network is busy shuffling gradient bytes between nodes, the GPUs are simultaneously busy computing the next batch of gradients. Neither operation waits for the other. Both proceed in parallel.

In the best case, communication disappears entirely from the critical path. The training throughput is determined purely by compute speed, just as if the distributed cluster were a single monolithic accelerator. That ideal is rarely achieved in practice, but the gap between naive sequential execution and well-engineered overlap can be a factor of two or more in wall-clock time.

This section develops the theory, analogies, and algorithms behind communication hiding, from informal intuition to formal analysis. Application to Generative AI: DDP Backward Pass applies these ideas to the specific architectures of generative model training.

An Everyday Analogy: Laundry and Cooking

Consider a Saturday morning task list: do the laundry and cook lunch. There are two strategies.

The sequential strategy. Put the laundry in the machine. Stand in the kitchen and watch the drum spin for forty-five minutes. When the cycle finishes, move the clothes to the dryer. Then start cooking. Total time: forty-five minutes of washing plus thirty minutes of cooking equals seventy-five minutes.

The overlapping strategy. Put the laundry in the machine. While it is running, cook lunch. When the timer beeps on both, you have clean clothes and a hot meal. Total time: $\max (45,30)=45$ minutes.

You did not reduce the amount of laundry. You did not cook faster. You simply stopped waiting. The washing machine ran, as machines do, without your supervision. You used that unattended time for useful work.

This is, precisely, the logic of communication hiding. The “washing machine” is the network card sending gradient tensors. The “cooking” is the GPU computing gradients for the next layer. Neither process requires the other to be finished before it can proceed. The only constraint is that the network card cannot send gradients it has not yet received from the GPU - just as you cannot hang laundry you have not yet washed.

More Analogies for Intuition

Communication hiding is a form of temporal parallelism: doing two things at the same time by exploiting the fact that they use different resources. Three more everyday examples sharpen the intuition.

Downloading while editing. You need to download a large dataset (3 GB, expected to take 8 minutes) and also review and annotate 200 existing data files. Sequential: 8 minutes of download, then annotation. Overlapping: start the download, annotate files while it proceeds. The annotation takes 6 minutes; the download takes 8. Total: 8 minutes instead of 14.

Waiting for a bus while reading. You must travel to a destination (15-minute bus wait, 20-minute ride) and must also read a 30-page report. Sequential: stand at the bus stop, ride the bus, then read. Overlapping: read while waiting and riding. The travel is on the critical path regardless; the reading fits inside it.

Pipelining in chip design. A CPU pipeline fetches the next instruction while decoding the current one, while executing the one before that. This is the canonical form of communication hiding in computer architecture: latency is hidden behind useful work, not eliminated.

In all cases the key insight is the same: some operations are blocking (they force the caller to wait for a result) while others are non-blocking (the caller can continue working while the operation proceeds in the background). Communication hiding requires rewriting blocking communication as non-blocking, and then filling the gap with useful computation.

Formal Framework

We now formalise the framework. Let a distributed training iteration consist of two phases: a computation phase of duration $T_{\mathrm{comp}}$ and a communication phase of duration $T_{\mathrm{comm}}$.

The speedup from hiding is the ratio of sequential to overlapped time:

(Hiding Speedup)$$\begin{array}{ccc}& {𝒮}_{\mathrm{hide}}=\frac{T_{\mathrm{comp}}+T_{\mathrm{comm}}}{\max (T_{\mathrm{comp}},T_{\mathrm{comm}})}.& \text{<span class="tml-eqn"></span>}\end{array}$$

When $T_{\mathrm{comm}}\ll T_{\mathrm{comp}}$ the speedup approaches 1 (communication was cheap to begin with). When $T_{\mathrm{comm}}=T_{\mathrm{comp}}$ the speedup is 2 - a doubling of throughput from hiding alone. When $T_{\mathrm{comm}}\gg T_{\mathrm{comp}}$ the speedup again approaches 1 from above: compute finishes quickly but we still wait for the network.

The Pipelining Approach

The most powerful realisation of communication hiding is pipelining: breaking both computation and communication into multiple stages and interleaving them so that stage $k$'s communication overlaps with stage $k+1$'s computation.

Consider a backward pass over $L$ layers. Each layer $\ell$ produces a gradient tensor ${𝐠}_{\ell }$ that must be all-reduced across workers. A naive implementation computes all $L$ gradient tensors, then launches a single all-reduce for the entire gradient vector. The all-reduce is $O(\text{model size})$ and blocks the next iteration.

A pipelined implementation instead launches the all-reduce for layer $\ell$'s gradient immediately after that layer's backward pass completes, while the backward pass continues to layer $\ell -1$.

Let $T_{\mathrm{AR}}^{\ell }$ denote the all-reduce time for layer $\ell$'s gradients and $T_{\mathrm{bwd}}^{\ell }$ the backward time for layer $\ell$. Under a layer-wise overlap schedule, the total iteration time is

(Layerwise TIME)$$\begin{array}{ccc}& T_{\mathrm{iter}}=\sum _{\ell =1}^L{T}_{\mathrm{bwd}}^{\ell }+\max (0,T_{\mathrm{AR}}^1-\sum _{\ell =2}^L{T}_{\mathrm{bwd}}^{\ell }),& \text{<span class="tml-eqn"></span>}\end{array}$$ where the second term captures the “tail” all-reduce for layer 1 that cannot be hidden behind further backward computation. When ${\sum }_{\ell =2}^L{T}_{\mathrm{bwd}}^{\ell }\ge T_{\mathrm{AR}}^1$, full hiding is achieved and the iteration time equals $T_{\mathrm{comp}}$.

Communication-Hiding Conjugate Gradient

Communication hiding was studied theoretically long before it became a practical concern in deep learning. In the numerical linear algebra community, iterative solvers such as conjugate gradient (CG) contain global dot products - collective operations equivalent to all-reduces - inside their inner loops. These synchronisation points were identified as the dominant bottleneck on massively parallel machines.

Ghysels and Vanroose (2014) introduced the communication-hiding conjugate gradient (CHCG) method, which reorganises the classic CG recurrences to decouple global dot products from the subsequent computation. In standard CG, the sequence of operations per iteration is:

1. Compute matrix-vector product $𝐪=𝐀𝐩$.

2. All-reduce to compute $\alpha =\rho /({𝐩}^{\top }𝐪)$.

3. Update $𝐱$ and $𝐫$ using $\alpha$.

4. All-reduce to compute $\beta ={\rho }_{\mathrm{new}}/\rho$.

5. Update $𝐩$.

Steps 2 and 4 are blocking global dot products. Between each pair of dot products, the computation is local.

CHCG recognises that the dot products in step 4 of iteration $k$ can be pipelined with the matrix-vector product in step 1 of iteration $k+1$. The reformulated algorithm reads: