Distributed Training: Multi-GPU and Multi-Node

India's Digital India push under the Bhashini mission commits to serving all 22 scheduled languages, which means labs such as AI4Bharat routinely need to train sequence models whose parameter counts run into the billions — far beyond anything a single accelerator can host. Suppose a team wants to train a 7-billion-parameter transformer, sized like a small but respectable open-weights model, on the single best GPU it can rent: an H100 with 80 GB of high-bandwidth memory. Before a single training step runs, arithmetic already rules this out. That arithmetic — not a vague appeal to "big models need big compute" — is the actual starting point of distributed training, and it is where this chapter begins.

The memory wall: why one GPU is not enough

Modern large-model training uses mixed-precision optimization with Adam. Four separate copies of every parameter live in GPU memory simultaneously, and each has a different precision and purpose:

• fp16 (or bf16) parameters used for the forward and backward matrix multiplies — 2 bytes per parameter.
• fp16 gradients produced by backpropagation — 2 bytes per parameter.
• an fp32 master copy of the parameters, kept because accumulating small updates in fp16 loses precision and stalls learning — 4 bytes per parameter.
• Adam's two fp32 running statistics, the first moment (momentum) and second moment (variance) of the gradient — 4 bytes each, 8 bytes per parameter.

Summing these: 2 + 2 + 4 + 8 = 16 bytes per parameter. This is exactly the figure Rajbhandari, Rasley, Ruwase, and He derive in the ZeRO paper (ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC20, 2020), where they write the total as 16Ψ bytes for Ψ parameters. For a 7-billion-parameter model:

16 bytes/param × 7×10⁹ params = 112×10⁹ bytes = 112 GB

112 GB does not fit in an 80 GB H100 — and this is before counting activation memory, which for a transformer scales with batch size, sequence length, and depth, and typically adds tens more gigabytes. A single GPU cannot even hold the optimizer state for a 7B model, let alone train it. This is the memory wall, and it is the first of two independent reasons distributed training exists. The second is the time wall: even a model that does fit on one GPU can take weeks to train on one GPU, and multiplying GPU count should ideally divide wall-clock time. Every technique in this chapter is an answer to one or both of these walls, and the central engineering tension is that fixing the memory wall (splitting a model across devices) and fixing the time wall (splitting data across devices) demand different, sometimes conflicting, communication patterns.

Data parallelism and the ring all-reduce

The simplest strategy assumes the model does fit on one GPU and attacks only the time wall. Data parallelism replicates the full model on every GPU, splits each training batch into shards — one per GPU — and lets every GPU compute a forward and backward pass on its own shard independently. The catch is that after backpropagation, every replica has a different gradient, computed from a different data shard, and all replicas must apply the same update or they drift apart into inconsistent models. So before the optimizer step, every GPU's gradient tensor must be averaged across all GPUs. This averaging is an all-reduce, and how it is implemented determines whether data parallelism scales to hundreds of GPUs or collapses under network contention.

A naive design routes every GPU's gradient through one central parameter server: each of N workers sends its gradient (size S) to the server, the server averages, and sends the result back. The server's own network link now has to absorb inbound traffic from N workers and send outbound traffic to N workers — total traffic 2NS through one node. If that link has bandwidth b, the time is 2NS/b, growing linearly with N. Double the GPU count and the bottleneck gets twice as bad, which defeats the purpose of adding GPUs.

The fix used by every modern framework (NCCL, Horovod) is ring all-reduce: workers arranged in a logical ring, each connected only to its two neighbours, no central node at all. Patarasuk and Yuan proved this pattern is bandwidth-optimal for cluster all-reduce (Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations, J. Parallel Distrib. Comput., 2009), and Sergeev and Del Balso's Horovod (2018) popularized it for deep learning frameworks. The algorithm has two phases, each taking exactly N−1 steps:

Phase 1 — reduce-scatter. Each GPU splits its local gradient tensor into N equal chunks. At each of N−1 steps, every GPU sends one chunk to its right neighbour and simultaneously receives a chunk from its left neighbour, adding the received chunk into its own matching chunk. After N−1 steps, each GPU holds the fully reduced (summed across all N GPUs) version of exactly one chunk — no GPU has the whole reduced tensor yet, but the reduction work is complete.

Phase 2 — all-gather. The same ring, N−1 more steps, but now GPUs simply relay their fully-reduced chunk onward instead of adding into it. After N−1 more hops, every GPU has received every other GPU's fully-reduced chunk, and all N GPUs hold the complete, identical, fully-reduced tensor.

I traced this by hand (verified programmatically, not just asserted) for N = 4 GPUs, each starting with a 4-element gradient vector:

GPU0 = [1, 2, 3, 4]
GPU1 = [5, 6, 7, 8]
GPU2 = [9, 10, 11, 12]
GPU3 = [13, 14, 15, 16]
# element-wise sum should converge to [28, 32, 36, 40] on every GPU

def ring_allreduce(local_arrays):
    N = len(local_arrays)
    csz = len(local_arrays[0]) // N
    chunk = lambda a, i: a[i*csz:(i+1)*csz]
    buf = [list(a) for a in local_arrays]

    # Phase 1: reduce-scatter (N-1 steps)
    for step in range(N - 1):
        sends = {i: ((i - step) % N, chunk(buf[i], (i - step) % N)) for i in range(N)}
        new_buf = [list(b) for b in buf]
        for i in range(N):
            dest = (i + 1) % N
            idx, vals = sends[i]
            for j in range(csz):
                new_buf[dest][idx*csz + j] += vals[j]
        buf = new_buf

    # Phase 2: all-gather (N-1 steps)
    for step in range(N - 1):
        sends = {i: ((i + 1 - step) % N, chunk(buf[i], (i + 1 - step) % N)) for i in range(N)}
        new_buf = [list(b) for b in buf]
        for i in range(N):
            dest = (i + 1) % N
            idx, vals = sends[i]
            for j in range(csz):
                new_buf[dest][idx*csz + j] = vals[j]
        buf = new_buf
    return buf

result = ring_allreduce([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]])
print(result)

Running this produces [[28, 32, 36, 40], [28, 32, 36, 40], [28, 32, 36, 40], [28, 32, 36, 40]] — all four GPUs converge on the identical, correctly-summed vector, confirmed against a brute-force element-wise sum. The diagram below shows the ring topology and the full seven-step trace (three reduce-scatter steps, three all-gather steps) that produced this result.

Now the bandwidth accounting that made the parameter server slow disappears. In each of the 2(N−1) total steps, every GPU sends and receives exactly S/N bytes (one chunk). Total per-GPU traffic:

2(N − 1)/N × S bytes

For a 1 GB (10⁹-byte) gradient tensor across N = 8 GPUs: 2×7/8×10⁹ = 1.75×10⁹ bytes = 1.75 GB per GPU — regardless of which GPU you look at, because there is no central node. At a link bandwidth of 100 GB/s (an illustrative NVLink-class figure), that is 1.75 GB ÷ 100 GB/s = 17.5 ms. The equivalent parameter-server transfer, 2×8×1×10⁹ = 16 GB through the one server link, takes 16 GB ÷ 100 GB/s = 160 ms — roughly 9× slower, and the gap widens as N grows, because ring traffic per GPU converges to a constant (2S as N → ∞) while parameter-server traffic through the hot node grows linearly with N. This is what "bandwidth-optimal" means concretely: the busiest link in the system carries the least possible traffic.

Model parallelism: splitting a single layer across GPUs

Data parallelism does nothing for the memory wall — every replica still needs the full 112 GB. When the model itself does not fit, you must split individual layers across GPUs, which Shoeybi et al. formalized as tensor parallelism in Megatron-LM (Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019). Consider a transformer MLP block computing Y = GeLU(XA)B, where X is the activation, A is the up-projection weight, and B is the down-projection weight. Split A column-wise across two GPUs, A = [A₁, A₂]. Each GPU computes X·Aᵢ independently — no communication needed for this matmul, since the columns are disjoint. GeLU is applied elementwise, so GeLU(XA₁) and GeLU(XA₂) can also be computed independently on each GPU with no communication. Now split B row-wise to match, B = [B₁; B₂]. Each GPU computes GeLU(XAᵢ)·Bᵢ, producing a partial sum of the final output — the two GPUs' partial results must be added together to get the true Y. That single addition is one all-reduce per MLP block, in the forward pass, and one more in the backward pass. Self-attention is split analogously by giving each GPU a disjoint subset of attention heads (heads are independent by construction), with one more all-reduce after the output projection.

The crucial property is frequency: tensor parallelism inserts an all-reduce inside every single transformer block, on both forward and backward passes — dozens of times per training step, each carrying a relatively small tensor. This makes it extremely latency-sensitive: it only works well when GPUs share a very high-bandwidth, low-latency interconnect, which in practice means keeping tensor-parallel groups inside one physical node (NVLink, roughly 900 GB/s aggregate on an H100 node) rather than spanning nodes connected by InfiniBand (roughly 50 GB/s per NIC on a 400 Gb/s NDR link) — an order of magnitude less bandwidth and materially higher latency. Put a tensor-parallel all-reduce across nodes and the frequent, small, latency-bound synchronizations dominate training time.

Pipeline parallelism and the bubble problem

An alternative way to split a model that doesn't fit is to assign whole layers to different GPUs — GPU 0 holds layers 1–8, GPU 1 holds layers 9–16, and so on — forming a pipeline. Huang et al. introduced this as GPipe (GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019). The obvious problem: if GPU 1 must wait for GPU 0 to finish layers 1–8 before it can start, then GPU 0 sits idle while GPU 1 works, and vice versa during the backward pass. GPipe's fix is to slice each mini-batch into m smaller micro-batches and stream them through the p pipeline stages, so that once the pipeline fills, every stage is processing a different micro-batch simultaneously. Even so, the pipeline has an unavoidable fill time (waiting for the first micro-batch to reach the last stage) and drain time (waiting for the last micro-batch to finish the first stage). If each stage takes time t per micro-batch, total wall-clock time to process m micro-batches through p stages is (p + m − 1)·t — this is the same fill-and-drain pattern as a systolic array. Each stage is actually busy for only m·t of that time; the rest, (p − 1)·t, is idle "bubble." The bubble fraction — the idle share of each device's time — is:

bubble fraction = (p − 1) / (p + m − 1)

For p = 4 stages and m = 8 micro-batches: (4−1)/(4+8−1) = 3/11 ≈ 27.3% of every device's time is wasted on pipeline fill/drain. GPipe's paper quotes the simpler approximation (p−1)/m, valid when m ≫ p; here that gives 3/8 = 37.5%, visibly different from the exact 27.3% because m is not much larger than p in this example — a reminder to use the exact formula unless the approximation's condition genuinely holds.

Pipeline communication looks nothing like tensor parallelism's: only activations (forward) and activation gradients (backward) cross the boundary between adjacent stages, and only p−1 times per micro-batch, not inside every layer. That communication is far less frequent and more latency-tolerant, so pipeline stages can be — and in large clusters usually are — placed on different physical nodes, connected by InfiniBand rather than requiring NVLink.

ZeRO: removing data parallelism's redundancy

Data parallelism's flaw is now obvious: every one of the N replicas stores the full 16Ψ bytes of parameters, gradients, and optimizer state, even though at any instant a replica only needs its own shard of that state to do its local computation on its own data shard. Rajbhandari et al.'s ZeRO (Zero Redundancy Optimizer) partitions this state across the N data-parallel GPUs instead of replicating it, in three increasingly aggressive stages: ZeRO-1 shards only the optimizer state (momentum, variance, fp32 master weights), ZeRO-2 additionally shards gradients, and ZeRO-3 additionally shards the parameters themselves. Under ZeRO-3, each GPU permanently holds only 1/N of the model. Applying this to the 7B-parameter, 112 GB example across a data-parallel group of N = 8 GPUs:

112 GB ÷ 8 = 14 GB per GPU

14 GB comfortably fits an 80 GB (or even a 24 GB) GPU alongside activation memory — the memory wall that blocked a single GPU is gone. The cost is that ZeRO-3 must temporarily reconstruct a layer's full weights via an all-gather immediately before that layer runs (forward and backward), then discard the reconstructed copy afterward — unlike pipeline parallelism, where each GPU owns its assigned layers outright and never needs to rebuild them. This is why ZeRO-3 communication volume, while it enables the largest memory savings, can exceed plain data parallelism's single per-step gradient all-reduce; it is a genuine tradeoff, not a strictly-better replacement, and picking ZeRO-1/2/3 is a bandwidth-versus-memory decision, not a default "always use the highest stage."

Multi-node: why topology dictates which parallelism goes where

Put the three mechanisms together and a clear placement rule emerges, matching the "3D parallelism" (tensor + pipeline + data, sometimes called PTD-parallelism) that Narayanan et al. describe for training GPT-scale models on real GPU clusters (Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21, 2021): tensor parallelism, with its frequent, latency-sensitive all-reduces inside every layer, stays confined to the GPUs inside one node connected by NVLink — typically matching the 8 GPUs in one server. Pipeline parallelism, whose communication is sparse and tolerant of latency, is used to span across nodes over InfiniBand, assigning contiguous blocks of layers to different machines. Data parallelism, needing only one all-reduce per training step (which frameworks further hide by overlapping it with the backward pass, starting each layer's gradient all-reduce as soon as that layer's backward computation finishes rather than waiting for the whole backward pass to complete), replicates this entire tensor-parallel-and-pipeline-parallel unit across as many additional node groups as the cluster provides, again over InfiniBand. Three different parallelism axes, three different communication frequencies, deliberately matched to three different points on the bandwidth-latency spectrum available in a real data center.

Code: what DDP and FSDP actually launch

PyTorch's DistributedDataParallel (DDP) is the ring-all-reduce data parallelism described above; FullyShardedDataParallel (FSDP) is PyTorch's implementation of ZeRO's sharding.

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

def setup():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def main():
    local_rank = setup()
    model = build_transformer().to(local_rank)   # assumed helper, not shown
    model = DDP(model, device_ids=[local_rank])
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

    for batch in train_loader:                    # assumed helper, not shown
        batch = {k: v.to(local_rank) for k, v in batch.items()}
        loss = model(**batch).loss
        loss.backward()       # DDP fires ring all-reduce here, overlapped with backward
        optimizer.step()
        optimizer.zero_grad()

if __name__ == "__main__":
    main()

This script is launched once per GPU by torchrun, which sets LOCAL_RANK, RANK, and WORLD_SIZE in each process's environment before main() runs:

torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
  --master_addr=10.0.0.1 --master_port=29500 train.py

Run identically on the second node with --node_rank=1, this launches 16 processes total (8 per node × 2 nodes), each owning one GPU, all joining the same NCCL process group. Swapping to ZeRO-3-style sharding changes only the wrapper:

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

model = build_transformer().to(local_rank)          # assumed helper, not shown
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,   # ZeRO-3 equivalent
    device_id=local_rank,
)

FULL_SHARD shards parameters, gradients, and optimizer state exactly as ZeRO-3 does, all-gathering each layer's weights just before it runs and freeing them immediately after — trading extra communication for the 8× memory reduction computed earlier.

Common misconception: "more GPUs trains proportionally faster"

The intuition that doubling GPU count halves training time treats communication as free and convergence as insensitive to how you split the computation. Neither holds. Pipeline parallelism's bubble fraction shows GPUs sitting idle purely from the mechanics of filling and draining a pipeline — that idle time exists no matter how fast the GPUs themselves compute. Tensor parallelism's per-layer all-reduce means adding GPUs to a tensor-parallel group past the point where NVLink bandwidth is saturated makes each layer slower, not faster, because the synchronization, not the matmul, becomes the bottleneck — this is exactly why tensor parallelism is capped at one node's GPU count in practice rather than scaled arbitrarily. And even where pure data parallelism scales communication well, increasing GPU count while keeping per-GPU batch size fixed increases the global batch size, which changes the optimization problem: Goyal et al. showed that scaling ResNet-50 training from a batch size of 256 to 8,192 across many GPUs required linearly scaling the learning rate and a short warmup period just to match single-GPU accuracy (Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, 2017); without that adjustment, more GPUs can finish more steps per second while needing more steps (or converging to a worse optimum) to reach the same accuracy, meaning wall-clock time to a target accuracy does not fall proportionally at all. "More GPUs" is not automatically "faster training" — it changes the communication topology, the idle-time structure, and the optimization dynamics all at once, and each of those three effects can silently cancel the benefit of the additional hardware.

Active recall

Attempt each question before reading its answer.

1. A lab trains a 13-billion-parameter model with mixed-precision Adam. (a) What is the total parameter + gradient + optimizer-state memory footprint? (b) Under ZeRO-3 with a data-parallel group of 4 GPUs, what is the per-GPU footprint? (c) Does that fit on 4×A100 GPUs with 40 GB each?
2. A gradient tensor is 2 GB (2×10⁹ bytes) and the data-parallel group has 16 GPUs. What is the ring all-reduce communication volume per GPU? How does the per-GPU volume compare, proportionally, to the earlier example of an 8-GPU group moving a 1 GB tensor, and why does it not simply double even though both N and S changed?
3. A pipeline has 6 stages and processes 24 micro-batches per mini-batch. Compute the exact bubble fraction. Also compute GPipe's approximate formula (p−1)/m and explain why the two values are close here but were not close for p=4, m=8.
4. Starting from the p=4, m=8 pipeline example (bubble fraction 27.3%), the team doubles the micro-batch count to m=16 to shrink the bubble. Trace every consequence: what happens to the bubble fraction, what happens to peak activation memory in GPipe's simple "all-forward-then-all-backward" schedule, and what further change might be forced as a result?
5. Match each communication pattern to the parallelism strategy it belongs to, and say whether it should be confined to one node or can span nodes: (a) one all-reduce per training step, moderate size, easily overlapped with backward computation; (b) an all-reduce inside every transformer block's forward and backward pass, small messages, highly latency-sensitive; (c) point-to-point activation handoffs only at stage boundaries, larger messages, only p−1 hops per micro-batch.
6. Why does ZeRO-3 need an all-gather before every layer's forward and backward computation, when plain pipeline parallelism never needs to reconstruct a layer's weights at all?

Answers

1. (a) 16 bytes/param × 13×10⁹ = 208×10⁹ bytes = 208 GB. (b) 208 GB ÷ 4 = 52 GB per GPU. (c) No — 52 GB exceeds a 40 GB A100's memory even before activations are added; this configuration needs either more GPUs in the data-parallel group (e.g., 8, giving 26 GB/GPU) or 80 GB GPUs.
2. Per-GPU volume = 2(16−1)/16 × 2×10⁹ = 1.875 × 2×10⁹ = 3.75×10⁹ bytes = 3.75 GB. The earlier N=8, S=1GB case gave 1.75 GB. The ratio (3.75/1.75 ≈ 2.14) is not simply 2× (from doubling S) or 2× (from doubling N) because the N-dependent factor 2(N−1)/N itself changes: it rises from 2×7/8=1.75 to 2×15/16=1.875 — closer to its asymptote of 2 as N grows, but the effect is small since both values are already near 2. The dominant driver here is S doubling; N's effect is second-order once N is reasonably large.
3. Exact: (6−1)/(6+24−1) = 5/29 ≈ 17.2%. Approximate: (p−1)/m = 5/24 ≈ 20.8%. The two are closer here than in the p=4, m=8 case because m=24 is 4× larger than p=6 (m ≫ p holds better), whereas m=8 was only 2× p=4 — the approximation (p−1)/m drops the "+p−1" term from the denominator, and that term matters less the larger m is relative to p.
4. Bubble fraction: (4−1)/(4+16−1) = 3/19 ≈ 15.8%, down from 27.3% — the pipeline is more efficient. But GPipe's simple schedule runs all m forward passes before any backward pass, so it must hold every micro-batch's activation checkpoints in memory simultaneously until the backward sweep begins; doubling m from 8 to 16 roughly doubles peak activation memory. If that no longer fits in GPU memory, the team is forced to shrink the per-micro-batch batch size to compensate, which lowers each matmul's arithmetic intensity and GPU utilization — partially offsetting the bubble-fraction gain they set out to capture. (The production fix, used by PipeDream and Megatron-LM's interleaved schedule, is a 1F1B — one-forward-one-backward — schedule that caps in-flight micro-batches at p rather than m, avoiding this memory blowup entirely.)
5. (a) is data parallelism's gradient all-reduce — can span nodes, since one synchronization per step tolerates InfiniBand latency and overlaps with backward compute. (b) is tensor parallelism's intra-layer all-reduce — must stay within one node on NVLink, since its frequency and small message size make it latency-bound. (c) is pipeline parallelism's stage-boundary handoff — can span nodes, since communication happens only p−1 times per micro-batch rather than inside every layer.
6. ZeRO-3 shards the parameters themselves across the data-parallel group, so no single GPU permanently owns a complete copy of any layer's weights — each GPU holds only a 1/N shard. To actually run a layer's forward or backward computation, the full weight tensor must exist somewhere, so it is temporarily reconstructed via all-gather immediately before use and discarded afterward to keep memory low. Pipeline parallelism never shards weights this way: each stage's GPU is permanently assigned a fixed set of whole layers and keeps their full weights resident for the entire run, so there is nothing to reconstruct.

Think About It

Think about this: How would you explain distributed training: multi-gpu and multi-node to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind distributed training: multi-gpu and multi-node, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.