Learn: Data Parallelism & Collective Communication

What you're learning

The baseline way to train on more than one GPU — replicate the model, split the batch, average the gradients — and the communication machinery that makes it work. By the end you should be able to: define every collective primitive precisely, derive the $2\frac{N-1}{N}M$ bandwidth cost of ring all-reduce and explain why it’s bandwidth-optimal and independent of $N$, reason about which interconnect (NVLink vs InfiniBand vs PCIe) bounds a collective, explain exactly how PyTorch DDP overlaps all-reduce with backward, and articulate the memory limit that kills plain DP for large models — which is the entire motivation for the next two lessons.

This lesson opens Tier 3 (distributed training). It assumes the single-GPU material: the memory hierarchy and HBM-bandwidth intuition, the roofline / compute-vs-communication lens (we now add a third bottleneck — the network), and the kernel work from 07-triton-and-modern-kernels. The vocabulary of collectives defined here (§4) is reused verbatim in lessons 09 (ZeRO/FSDP) and 10 (model/tensor/pipeline parallelism).

---

1. Learning map

graph TD
    A["Single-GPU training<br/>(Tiers 1-2)"] --> B["Model too big to train fast<br/>on one GPU / batch too small"]
    B --> C["Data parallelism<br/>replicate + split batch"]
    C --> D["Need: average gradients<br/>across N replicas"]
    D --> E["Collective primitives<br/>bcast, reduce, all-reduce,<br/>reduce-scatter, all-gather, all-to-all"]
    E --> F["All-reduce = the workhorse"]
    F --> G["Ring all-reduce<br/>cost = 2(N-1)/N · M<br/>bandwidth-optimal"]
    F --> G2["Tree all-reduce<br/>O(log N) latency<br/>wins at small sizes"]
    G --> H["NCCL<br/>topology-aware impl"]
    G2 --> H
    H --> I["Interconnects<br/>NVLink/NVSwitch, IB/RoCE, PCIe"]
    I --> J["PyTorch DDP<br/>buckets + hook overlap + no_sync"]
    J --> K["Scaling efficiency<br/>comms-bound regime"]
    K --> L["GPU/LLM connection:<br/>DP needs full model+optim<br/>PER GPU → fails at scale"]
    L --> M["→ ZeRO/FSDP (L09)<br/>→ model parallelism (L10)"]

    style F fill:#4a4,color:#fff
    style G fill:#44a,color:#fff
    style L fill:#a44,color:#fff

---

2. Why this matters

A single H100 does ~990 dense BF16 TFLOP/s.¹ A frontier pre-training run is $\sim 10^{25}$–$10^{26}$ FLOPs (from the $6ND$ rule of lesson 03). Divide: even at a perfect 100% MFU, one GPU would take centuries. You have no choice but to spread the work across thousands of GPUs — and the moment you do, a new bottleneck appears that the roofline of lesson 03 never modeled: the network between GPUs.

Data parallelism (DP) is the first, simplest, and still most-used axis of parallelism. It is embarrassingly parallel on the compute side — every GPU runs the identical forward/backward on its own slice of data — with exactly one synchronization point: gradients must be averaged before the optimizer step, so every replica stays bit-identical. That one sync is an all-reduce, and understanding its cost is the core of this lesson.

The three bottlenecks

Lesson 03 taught you compute-bound vs memory(HBM)-bound. Distributed training adds a third: communication-bound (network-bound). A distributed step is bounded by $\max (\text{compute time},\text{HBM time},\text{comm time})$ if perfectly overlapped, and by their sum if not. The entire game of distributed training is (a) making comm cheap (good algorithms + fast interconnect) and (b) hiding it under compute.

---

3. Data parallelism: the mechanics

The recipe, on $N$ GPUs (ranks $0\dots N-1$):

1. Replicate. Each GPU holds a full copy of the model weights $\theta$, and — critically — its optimizer state (Adam’s $m,v$ moments, plus an FP32 master copy of the weights; see lesson 03 §5). All $N$ copies start identical.
2. Shard the batch. The global batch $B$ is split into $N$ local (per-GPU) micro-batches of size $B/N$. Each GPU gets different data (via a DistributedSampler).
3. Compute local gradients. Each GPU runs forward + backward on its local batch, producing a local gradient $g_i=\frac{1}{B/N}{\sum }_{x\in {\text{batch}}_i}{\nabla }_{\theta }\ell (x)$.
4. Average via all-reduce. All ranks compute $\bar{g}=\frac{1}{N}{\sum }_{i=0}^{N-1}{g}_i$. This is mathematically identical to the gradient of the full global batch — that equivalence is the whole point.
5. Step in lockstep. Every GPU applies the same $\bar{g}$ with its optimizer. Since weights + optimizer state + gradient are now identical across ranks, the replicas remain bit-for-bit synchronized. No weight communication is ever needed — only the gradient sync.

Effective / global batch size

$B_{\text{global}}=B_{\text{local}}\times N\times A_{\text{accum}}$

where $A_{\text{accum}}$ is the gradient-accumulation factor (below). This number — not the per-GPU batch — is what governs optimization dynamics, and it’s what you must hold fixed (and re-tune the LR for) when you change GPU count. The canonical scaling rule is the linear LR scaling rule: when you multiply the global batch by $k$, multiply the learning rate by $k$ (with warmup), which let Goyal et al. train ImageNet/ResNet-50 at batch 8192 with no accuracy loss.²

More GPUs silently changes your optimization problem

Naively scaling from 8 to 512 GPUs with the same per-GPU batch multiplies the global batch by 64×. That is a different training run — gradient noise scale, effective LR, and generalization all move. Data parallelism is “just replication” on the systems side, but on the optimization side you are changing a real hyperparameter. Always reason in terms of $B_{\text{global}}$.

Gradient accumulation

If the batch size you want doesn’t fit in memory (activations blow up HBM), split the local batch into $A_{\text{accum}}$ micro-batches, run forward/backward on each, and sum the gradients in .grad without stepping. Step (and all-reduce) only after the last micro-batch. This trades wall-clock for memory: it emulates a large batch with a small memory footprint.

Crucially for DP: you must skip the all-reduce on the non-final micro-batches — averaging partial sums every micro-batch is wasted network traffic (and the accumulation is associative, so the single final all-reduce over the summed gradient gives the identical result). PyTorch exposes this as the no_sync() context manager (§8).³

Quiz: You train with $N=256$ GPUs, local batch 4, and accumulation factor 8. What is the global batch? If you drop to $N=64$ GPUs but want the same global batch and same per-GPU memory footprint, what do you change?

Answer

$B_{\text{global}}=4\times 256\times 8=8192$.
Dropping to 64 GPUs loses a factor of 4 in the $N$ term. To hold $B_{\text{global}}=8192$ at the same per-GPU batch of 4 (memory-preserving, since activation memory is set by the micro-batch size, not the accumulation count), raise accumulation from 8 to 32: $4\times 64\times 32=8192$. You keep the optimization problem identical; you just do 4× more sequential accumulation per step, so each step is ~4× slower in wall-clock but the math is unchanged. This memory-vs-time knob is why accumulation is the first thing to reach for when a batch won’t fit.

---

4. Collective communication primitives

A collective is a communication operation involving all ranks in a group, in contrast to point-to-point (one sender, one receiver). These six are the complete vocabulary you need for lessons 08–10. Notation: $N$ ranks; each rank $r$ holds a buffer; $\oplus$ is an associative reduction operator (usually sum, sometimes max/min/prod).⁴

Let each rank’s input be a vector; “chunk $j$” means the $j$-th equal slice of a buffer.

Primitive	Before → After (rank $r$)	One-line definition
Broadcast	root has $x$ → all have $x$	One rank’s buffer is copied to every rank.
Reduce	rank $r$ has $x_r$ → root has ${⨁}_r{x}_r$	Element-wise reduce all inputs; result lands on one root rank.
All-Reduce	rank $r$ has $x_r$ → all have ${⨁}_r{x}_r$	Reduce, but the result is delivered to every rank.
Reduce-Scatter	rank $r$ has vector split in $N$ chunks → rank $r$ has ${⨁}_s{x}_s[r]$ (only chunk $r$, reduced)	Reduce element-wise, but scatter the result: rank $r$ keeps only the $r$-th reduced chunk.
All-Gather	rank $r$ has chunk $y_r$ → all have $[y_0,y_1,\dots ,y_{N-1}]$	Concatenate every rank’s chunk; all ranks get the full concatenation.
All-to-All	rank $r$ has $[z_r^0,\dots ,z_r^{N-1}]$ → rank $r$ has $[z_0^r,\dots ,z_{N-1}^r]$	Transpose: rank $r$ sends chunk $j$ to rank $j$ and receives chunk $r$ from everyone.

Two identities to burn in — they drive the rest of distributed training:

$\boxed{\text{All-Reduce}=\text{Reduce-Scatter}+\text{All-Gather}}$

$\text{All-Gather}\text{is the "inverse" of}\text{Reduce-Scatter (with sum dropped)}$

The first identity is the decomposition behind ring all-reduce (§5) and behind ZeRO/FSDP (lesson 09: FSDP all-gathers sharded weights before a layer and reduce-scatters gradients after). All-to-all is the primitive behind expert (MoE) parallelism and sequence parallelism (lesson 10).

Why reduce-scatter + all-gather, not "just all-reduce"?

Because splitting the op exposes where the data lives. Reduce-scatter leaves each rank owning a distinct $1/N$ slice of the reduced result — which is exactly what you want if the next thing each rank does is operate on only its slice (ZeRO shards the optimizer step this way). All-reduce is the special case where you immediately all-gather those slices back to everyone. Seeing all-reduce as two half-ops is the single most useful lens for lessons 09–10.

graph LR
    subgraph "All-Reduce = two phases"
        A["Reduce-Scatter<br/>each rank ends with<br/>1/N of the SUM"] --> B["All-Gather<br/>broadcast the slices<br/>so all have the full sum"]
    end

---

5. Ring all-reduce and its cost

The naive ways to all-reduce are bad, and seeing why motivates the ring.

Naive #1 — reduce-to-root-then-broadcast. Every rank sends its $M$-byte buffer to rank 0; rank 0 sums and broadcasts the result back. Rank 0’s single link must carry $(N-1)M$ bytes in and $(N-1)M$ out. Cost scales with $N$ and one link is the hot bottleneck — the rest of the fabric sits idle. Terrible bandwidth utilization.

Naive #2 — every rank all-to-all. Each rank sends its buffer to every other rank; each computes the sum locally. Every rank sends $(N-1)M$. Uses all links but moves $N\times$ too much data.

The ring fixes both by (a) using every link simultaneously and (b) never sending any byte more than necessary. It realizes the $\text{all-reduce}=\text{reduce-scatter}+\text{all-gather}$ decomposition on a logical ring where rank $r$ only ever sends to $(r+1)\mathrm{mod}N$ and receives from $(r-1)\mathrm{mod}N$.⁵⁶

graph LR
    G0["GPU 0"] --> G1["GPU 1"]
    G1 --> G2["GPU 2"]
    G2 --> G3["GPU 3"]
    G3 --> G0

The logical ring: each GPU sends only to its clockwise neighbor. Every link is used every step — no hot spot.

The reduce-scatter phase

Split each rank’s $M$-byte buffer into $N$ chunks of $M/N$ bytes. Run $N-1$ steps. In step $k$, each rank sends one chunk to its successor and adds the chunk it receives into its own buffer. After $N-1$ steps, rank $r$ holds the fully reduced chunk $r$ (the sum over all ranks of that chunk), and every other chunk is partially reduced. Bytes sent per rank:

$(N-1)\cdot \frac{M}{N}$

The all-gather phase

Now each rank owns one complete reduced chunk. Run another $N-1$ steps of the same ring pattern, but now each rank forwards the complete chunk it received (no reduction). After $N-1$ steps, every rank has all $N$ reduced chunks — i.e. the full reduced buffer. Bytes sent per rank:

$(N-1)\cdot \frac{M}{N}$

Total cost

Bytes sent per GPU across both phases:

$\boxed{{\text{data}}_{\text{ring}}=2(N-1)\cdot \frac{M}{N}=2\frac{N-1}{N}M}$

Under the standard $\alpha$–$\beta$ cost model (a message of $n$ bytes costs $\alpha +n\beta$, where $\alpha$ = per-message latency and $\beta =1/\text{bandwidth}$ = per-byte time),⁷ the ring all-reduce time is:

$T_{\text{ring}}=\underset{\text{latency term}}{\underbrace{2(N-1)\alpha }}+\underset{\text{bandwidth term}}{\underbrace{2\frac{N-1}{N}M\beta }}$

Now the two payoffs:

(1) Bandwidth-optimal and $N$-independent. As $N\to \infty$, $\frac{N-1}{N}\to 1$, so the bandwidth term $\to 2M\beta$ — a constant independent of $N$. Whether you have 8 GPUs or 8192, each GPU pushes ~$2M$ bytes over its link. This is provably the minimum: every rank must receive at least $M(N-1)/N$ bytes of “new” data to know the sum, and send that much, so $2\frac{N-1}{N}M$ is a lower bound the ring meets exactly.⁶ More GPUs does not cost more bandwidth per GPU — that is why ring all-reduce is the backbone of large-scale DP.

(2) But the latency term grows with $N$. The $2(N-1)\alpha$ term scales linearly in $N$: $2(N-1)$ sequential hops, each paying the fixed latency $\alpha$. For small messages (where $M\beta \ll \alpha$) at large $N$, this dominates and the ring becomes latency-bound.

Contrast: tree all-reduce

A tree (or recursive-doubling) all-reduce reduces up a binary tree and broadcasts down, completing in $\sim 2\lceil {\log }_{2}N\rceil$ steps instead of $2(N-1)$:

$T_{\text{tree}}\approx 2\lceil {\log }_{2}N\rceil \alpha +(\text{bandwidth term})$

The latency term is now logarithmic in $N$ instead of linear — a huge win at large scale for small messages. The tradeoff is that a simple tree isn’t bandwidth-optimal (interior links carry more than $2\frac{N-1}{N}M$). NCCL’s answer is double binary trees: two complementary trees run at once so that every rank is a leaf in one tree and interior in the other, balancing the load and recovering near-full-bandwidth utilization while keeping $O(\log N)$ latency.⁸

The crossover: ring for big messages, tree for small

Ring minimizes bandwidth (flat in $N$) but its latency grows linearly in $N$. Tree minimizes latency ($\log N$) but historically at a bandwidth cost. So: small messages / large $N$ → tree; large messages → ring. NCCL picks automatically per collective, per message size, per topology (§6). The published crossover is roughly in the low-tens-of-MB range, and tree tends to win beyond ~32–64 nodes for latency-bound sizes.⁸⁹

Quiz: Derive why ring all-reduce bandwidth cost is $2\frac{N-1}{N}M$ and explain in one sentence why it beats the naive reduce-to-root for large $N$. Then: at $N=8$ vs $N=512$, how does per-GPU bandwidth cost change?

Answer

Derivation. All-reduce = reduce-scatter + all-gather. Split the $M$-byte buffer into $N$ chunks of $M/N$. Reduce-scatter is $N-1$ ring steps, each sending one chunk ($M/N$ bytes) → $(N-1)M/N$ sent per GPU, after which each GPU owns one fully-summed chunk. All-gather is another $N-1$ steps sending $M/N$ each → another $(N-1)M/N$. Total sent per GPU $=2(N-1)M/N=2\frac{N-1}{N}M$. $\square$

Why it beats reduce-to-root. In reduce-to-root the root’s single link must move $\sim (N-1)M$ each way — cost grows linearly with $N$ and every other link idles; the ring spreads the traffic so every link moves the same $\sim 2M$ regardless of $N$, saturating the whole fabric.

$N=8$ vs $512$. $2\frac{7}{8}M=1.75M$ vs $2\frac{511}{512}M=1.996M$. Essentially unchanged (both $\to 2M$). Per-GPU bandwidth cost is flat in $N$ — the defining property. (The latency term $2(N-1)\alpha$, by contrast, grows 73× from $N=8$ to $512$, which is exactly when you’d switch to a tree.)

The ring in one picture

Ring all-reduce as two ring passes. The reduced buffer is never fully materialized on any single link — each of the $2(N-1)$ steps moves only $M/N$ bytes, and the total per-GPU traffic is bandwidth-optimal.

---

6. NCCL: the implementation

NCCL (NVIDIA Collective Communications Library, pronounced “nickel”) is the de-facto implementation of these collectives on NVIDIA GPUs, and the backend PyTorch, Megatron, DeepSpeed, and JAX all call for GPU communication.⁴ You essentially never hand-roll a ring; you call ncclAllReduce (directly or via torch.distributed).

What NCCL actually does for you:

• Topology discovery. At init, NCCL detects the physical fabric — NVLink/NVSwitch connectivity, PCIe switch trees, NIC-to-GPU affinity, InfiniBand/RoCE links — and builds a set of rings and (double-binary) trees mapped onto it so that ring hops follow fast links wherever possible.⁴¹⁰
• Algorithm + protocol selection. Per collective, per message size, per cluster shape, NCCL chooses the algorithm (Ring, Tree, NVLS/NVLSTree, CollNet/CollNetChain, PAT) and, orthogonally, the wire protocol (Simple, LL, LL128). You can pin these with NCCL_ALGO / NCCL_PROTO for debugging, but the auto-tuner is usually right.⁴⁹
• In-network reduction (SHARP). With NVLS (NVLink SHARP, Hopper+) and CollNet (InfiniBand SHARP), the switch ASIC itself performs the summation. GPUs send their contribution once and receive the result once — per-GPU volume drops toward $\frac{N-1}{N}M$, roughly half of ring’s $2\frac{N-1}{N}M$, because the reduction happens in the fabric instead of round-tripping through GPU memory.⁹¹¹
• Channels. A collective is split across multiple parallel channels, each a ring/tree driven by a subset of SMs, to saturate the available links.

Topology-awareness is not optional at scale

A topology-unaware single ring laid across 4 GPUs split over two PCIe switches would route bulk traffic over the slow inter-switch PCIe link, throttling the whole collective to PCIe speed. NCCL instead builds a hierarchical collective — reduce within an NVLink island first, then across islands over the NIC — so the fast links carry the bulk. This is why “which GPUs are wired to which” (§7) directly determines collective throughput.¹⁰

(NCCL’s public release line is well past the early 2.x series — recent versions add NVLS FP8 reductions, PAT for cross-datacenter topologies, etc. Treat specific version behavior as fast-moving and check the release notes.)⁴

---

7. Interconnects: the physical fabric

The collective’s cost model has $\beta =1/\text{bandwidth}$ — so the actual interconnect sets the constant. There is a hierarchy of links, each ~an order of magnitude apart, and a collective is only as fast as the slowest link it must traverse.

Link	Scope	Per-GPU bandwidth	Notes
NVLink 4 + NVSwitch (H100)	intra-node (within a server)	900 GB/s bidirectional	18 NVLink4 links × 50 GB/s; 4 on-board NVSwitch3 chips form a non-blocking crossbar so every GPU can push 900 GB/s to any other simultaneously.1211
NVLink 5 (Blackwell)	intra-node	~1800 GB/s bidirectional	~2× NVLink4; scales NVLink domains to 72 GPUs (GB200 NVL72).13
InfiniBand NDR (ConnectX-7)	inter-node (across servers)	400 Gb/s = 50 GB/s per port	Standard HGX H100: one CX-7 NIC per GPU → 400 Gb/s GPU-to-fabric each; Quantum-2 switches.14
InfiniBand XDR (ConnectX-8)	inter-node	800 Gb/s = 100 GB/s per port	Current frontier for B200/GB300; Quantum-X800 switches, SHARP v4.15
RoCE (RDMA over Converged Ethernet)	inter-node	comparable to IB (400/800 GbE)	Ethernet alternative to InfiniBand; NVIDIA Spectrum-X targets this.14
PCIe 5.0	host↔GPU / GPU↔GPU w/o NVLink	~64 GB/s bidirectional (x16)	The slow path. If a collective is forced onto PCIe, it dominates the cost.12

The headline gap: NVLink is ~18× the per-GPU bandwidth of NDR InfiniBand (900 vs 50 GB/s). That single ratio drives the design of every large-scale parallelism strategy.

The intra-node / inter-node cliff is the central fact of distributed training

Communication inside an 8-GPU node (NVLink, 900 GB/s) is ~18× faster than communication between nodes (NDR IB, 50 GB/s/GPU). So you place the chattiest parallelism (tensor parallelism, lesson 10 — all-reduces on the critical path of every layer) inside a node over NVLink, and the least chatty (data/pipeline parallelism — one sync per step) across nodes over IB. Getting this mapping wrong can cost you 2–5× in throughput. This “hierarchical parallelism” placement is the payoff of understanding both collectives and the fabric.

graph TD
    subgraph "Node 0 (NVLink 900 GB/s, non-blocking)"
        A0["GPU0"] --- A1["GPU1"] --- A2["GPU..."] --- A3["GPU7"]
    end
    subgraph "Node 1 (NVLink 900 GB/s)"
        B0["GPU0"] --- B1["GPU1"] --- B2["GPU..."] --- B3["GPU7"]
    end
    A0 -. "InfiniBand NDR<br/>400 Gb/s/GPU<br/>(~18x slower)" .- B0
    style A0 fill:#4a4,color:#fff
    style B0 fill:#4a4,color:#fff

Quiz: You must all-reduce a 500 MB gradient buffer across (a) 8 GPUs in one node over NVLink4, (b) 8 GPUs spread one-per-node over NDR InfiniBand. Estimate the bandwidth-limited time for each, ignoring latency. What does the ratio tell you?

Answer

Ring all-reduce moves $2\frac{N-1}{N}M\approx 2\cdot \frac{7}{8}\cdot 500\text{ MB}=875$ MB per GPU either way (algorithm is identical; $N=8$).
(a) NVLink4 at 900 GB/s: $0.875\text{ GB}/900\text{ GB/s}\approx 0.97$ ms.
(b) NDR IB at 50 GB/s/GPU: $0.875/50\approx 17.5$ ms.
Ratio ≈ 18×, exactly the bandwidth ratio. Same algorithm, same bytes — the fabric is the entire difference. Takeaway: keep high-frequency collectives on NVLink; this is why an 8-GPU node is the natural unit and why crossing node boundaries is the thing you architect around.

---

8. PyTorch DDP: making it fast in practice

torch.nn.parallel.DistributedDataParallel (DDP) is the standard implementation of §3. Wrapping a model in DDP gives you data-parallel training with three key optimizations that turn the naive “backward, then all-reduce” into something that hides most of the communication.³¹⁶

(1) Gradient bucketing

Calling a separate all-reduce for each of thousands of parameter tensors would drown in per-message latency ($\alpha$) and kernel-launch overhead. DDP instead coalesces gradients into buckets (default 25 MB, bucket_cap_mb) and all-reduces one bucket at a time.³¹⁶ Fewer, larger messages amortize $\alpha$ and push each transfer into the bandwidth-optimal regime of the ring (§5).

Parameters are assigned to buckets in roughly the reverse order of model.parameters(), because backward produces gradients last-layer-first — so the reverse order makes buckets fill in approximately the order gradients become ready, letting communication start as early as possible.³

(2) Backward / all-reduce overlap via autograd hooks

This is the crucial performance idea. DDP registers an autograd post-hook on each parameter’s AccumulateGrad node at construction. When a gradient is computed during backward, its hook fires and marks it ready. When every gradient in a bucket is ready, DDP kicks off an asynchronous all-reduce on that bucket immediately — while the backward pass is still computing gradients for earlier layers.³¹⁶ Communication of the later layers thus overlaps with computation of the earlier ones. At the end of backward, DDP waits on the outstanding all-reduces and writes the averaged results into .grad.

All-reduce order must be identical across ranks

DDP always issues bucket all-reduces in bucket-index order, not in the (possibly rank-varying) order buckets happen to finish. If ranks issued collectives in different orders, they’d mismatch and either produce wrong results or hang (a classic deadlock — two ranks each waiting on a collective the other hasn’t launched). This is why NCCL collectives must be called in the same order on every rank.³

As each bucket’s gradients finish (last layer first), its all-reduce launches asynchronously and runs concurrently with the still-computing earlier layers. Ideally only the final bucket’s all-reduce is left “exposed” at the end of backward.

(3) no_sync for gradient accumulation

Under the model.no_sync() context manager, DDP disables the all-reduce for that backward pass — gradients just accumulate in .grad locally. You run all but the last accumulation micro-batch inside no_sync(), then the final micro-batch outside it triggers a single all-reduce over the summed gradient.³ This is the §3 accumulation optimization: one network sync per optimizer step, not per micro-batch.

for i, micro_batch in enumerate(micro_batches):
    is_last = (i == len(micro_batches) - 1)
    ctx = contextlib.nullcontext() if is_last else model.no_sync()
    with ctx:
        loss = model(micro_batch).loss / len(micro_batches)
        loss.backward()          # all-reduce fires ONLY on the last iter
optimizer.step(); optimizer.zero_grad()

Scaling efficiency and the comms-bound regime

Define scaling efficiency = (throughput on $N$ GPUs) / ($N\times$ throughput on 1 GPU). Perfect DP is 1.0; overlap and fast interconnect push you toward it. The step is roughly:

$T_{\text{step}}\approx \max (T_{\text{compute}},T_{\text{comm}})\text{(with perfect overlap)}$

You are compute-bound (good, high efficiency) when $T_{\text{compute}}>T_{\text{comm}}$, and comms-bound (efficiency collapses) when communication can no longer hide under compute. What pushes you into the comms-bound regime:

• Small per-GPU batch / short sequences → $T_{\text{compute}}$ shrinks but the gradient all-reduce ($\approx 2M$ bytes, fixed by model size) doesn’t → comm is exposed. (This is why “just add more GPUs” hits diminishing returns: each GPU’s compute share drops while its comm stays fixed.)
• Slow interconnect (across-node NDR instead of NVLink, or worse, PCIe) → $T_{\text{comm}}$ balloons 18×+.
• Large model → $M$ grows, and eventually the all-reduce of the full gradient can’t be hidden. (This, plus memory, is the wall in §9.)

Quiz: DDP overlaps all-reduce with backward. (a) Which layer's gradient is all-reduced first, and why does that ordering help? (b) After all this overlap, which single all-reduce is typically still "exposed" (not hidden)? (c) Why does shrinking the per-GPU batch hurt scaling efficiency even though the all-reduce size is unchanged?

Answer

(a) The last layer’s gradient finishes first (backward runs output→input), so its bucket is ready first and its all-reduce launches first — precisely while the earlier layers are still computing their gradients. Reverse-order bucketing is designed to exploit this: comm starts at the earliest possible moment and has the whole rest of backward to hide under.
(b) The first layer’s bucket (the last to become ready). Its all-reduce can only start when backward is essentially done, so there’s little/no remaining compute to overlap it with — it sticks out at the tail as exposed communication.
(c) All-reduce bytes $\approx 2M$ are set by the model size, independent of batch. Shrinking the per-GPU batch shrinks $T_{\text{compute}}$ (less math per step) but leaves $T_{\text{comm}}$ fixed. So the fixed comm has less compute to hide behind; the exposed fraction grows and efficiency drops. Formally, the ratio $T_{\text{comm}}/T_{\text{compute}}$ rises → you slide toward the comms-bound regime.

---

9. GPU/LLM connection: why DP alone hits a wall

Data parallelism is the baseline, and for models that comfortably fit on one GPU it’s often all you need. But it has a hard, structural limitation that is the entire reason lessons 09 and 10 exist:

DP replicates everything on every GPU

Every DP replica must hold the full model weights + full gradients + full optimizer state, plus activations. DP shards the data, never the model. So if the model’s state doesn’t fit on one GPU, DP cannot help you — adding GPUs adds identical copies, not more room.

Do the memory arithmetic for mixed-precision Adam (from lesson 03 §5): per parameter you store a BF16 weight (2 B), BF16 gradient (2 B), and in FP32 the master weight (4 B) + Adam’s $m$ (4 B) + $v$ (4 B) = 16 bytes/param for weights+grads+optimizer.¹⁷ For a 7B model that’s $7\times 10^9\times 16=112$ GB — already over an 80 GB H100, before a single activation. A 70B model needs ~1.1 TB of state. No amount of data parallelism fixes this: each replica needs the whole 1.1 TB.

The escape is to stop replicating and start sharding the model state itself — which is exactly what the collectives of §4 enable:

• ZeRO / FSDP (lesson 09) shards the optimizer state, gradients, and weights across the DP group, reconstructing each layer’s weights on the fly with an all-gather and averaging gradients with a reduce-scatter (the two halves of the all-reduce you just learned!). Memory per GPU drops ~$N\times$; the price is extra communication. This is “data parallelism that also shards the model.”
• Tensor / pipeline / expert parallelism (lesson 10) splits individual layers or the layer stack across GPUs, using all-reduce (tensor parallel), point-to-point (pipeline), and all-to-all (expert/MoE). These put chatty collectives on the critical path — which is why §7’s “keep it on NVLink” rule matters so much.

Real frontier training composes all of these (“3D/4D parallelism”). But every one of them is built from the six primitives in §4 and bounded by the fabric in §7. Master this lesson and the next two are just recompositions of it.

---

Practice problems

Practice problem 1: All-reduce time and comms-bound crossover

Problem: A 7B model trains with DDP. Gradient buffer is 7B params in BF16 = 14 GB. You run pure DP on 64 H100s across 8 nodes (8 GPUs/node), so the all-reduce spans NDR InfiniBand (assume the whole collective is bottlenecked at 50 GB/s/GPU effective). (a) Compute the ring all-reduce time per step (ignore latency). (b) If one training step’s compute is 400 ms, are you compute- or comms-bound, and what’s the best-case scaling efficiency contribution from this collective? (c) Name two changes that would flip the verdict.

Worked solution:
(a) Per-GPU bytes $=2\frac{N-1}{N}M=2\cdot \frac{63}{64}\cdot 14\text{ GB}\approx 27.6$ GB. At 50 GB/s: $27.6/50\approx 0.55$ s $=550$ ms.
(b) $T_{\text{comm}}=550$ ms $>T_{\text{compute}}=400$ ms → comms-bound. Even with perfect overlap, $T_{\text{step}}\approx \max (400,550)=550$ ms, and $\ge 150$ ms of comm is exposed. Efficiency $\approx 400/550\approx 73\%$ — you’re paying 27% to the network.
(c) (i) Bigger per-GPU batch / longer sequences → raise $T_{\text{compute}}$ above 550 ms so comm hides fully. (ii) Faster fabric or in-network reduction — XDR (100 GB/s/GPU) roughly halves the transfer to ~275 ms, or NVLS/SHARP halves the volume to $\frac{N-1}{N}M$; either makes you compute-bound. (Also valid: FP8 gradients halve $M$; or switch to ZeRO so you reduce-scatter instead of all-reduce, cutting per-GPU volume.)

Practice problem 2: Does the model even fit — and when do you abandon DP?

Problem: You have 80 GB H100s. Using mixed-precision Adam (16 bytes/param for weights + grads + optimizer state), what’s the largest model whose state (ignore activations) fits on one GPU? For a 30B model, how much state per GPU under pure DP, and how many GPUs does DP need to make it fit? What’s the actual fix?

Worked solution:
Max model: $80\text{ GB}/16\text{ B/param}=5\times 10^9=$ 5B params — and that leaves zero room for activations, so realistically ~3–4B. This is why even “small” LLMs need help.
30B under DP: state $=30\times 10^9\times 16=480$ GB per GPU. That’s 6× an H100. Number of DP GPUs that fixes it: none — DP replicates the full 480 GB on every GPU, so adding GPUs never reduces per-GPU state. DP is the wrong tool the moment state exceeds one GPU.
The fix: shard the state. ZeRO-3 / FSDP (lesson 09) spreads the 480 GB across the group: on 16 GPUs, ~30 GB/GPU of sharded state — now it fits, at the cost of all-gathering weights per layer. This is the concrete handoff from this lesson to the next: DP’s replication wall → sharding.

---

Practice & resources

Hands-on

1. Run a real multi-GPU DDP job. Follow the official PyTorch DDP tutorial and launch with torchrun --nproc_per_node=<N> train.py. Wrap your model in DistributedDataParallel, use a DistributedSampler, and confirm loss curves match single-GPU (they should, up to batch-size effects). Then add gradient accumulation with no_sync() and verify the gradient is bit-identical to accumulating without DP. pytorch.org/tutorials/intermediate/ddp_tutorial.html
2. Benchmark all-reduce bandwidth with nccl-tests. Build github.com/NVIDIA/nccl-tests and run ./build/all_reduce_perf -b 8 -e 8G -f 2 -g <N>. Read the busbw (bus bandwidth) column — it reports the achieved $2\frac{N-1}{N}M/T$ and lets you see the ring/tree crossover as message size grows, and whether you’re hitting NVLink vs IB limits. Sweep NCCL_ALGO=Ring vs Tree to see the small-vs-large-message tradeoff of §5 directly.
3. Profile the overlap. Run nsys profile on a few DDP steps (as in lesson 03 §8) and look at the NCCL all-reduce kernels on the timeline. Confirm they overlap with backward compute, and measure the “exposed” tail (the first-layer bucket). Try setting bucket_cap_mb smaller/larger and watch the overlap change.

Where to get multi-GPU access

• Colab / Kaggle: free tiers give a single GPU (T4/P100; Colab Pro adds A100/L4) — enough to write DDP code and test it in single-process mode, but not multi-GPU. Kaggle offers 2× T4 in some sessions — the cheapest way to run a genuine 2-GPU all-reduce. Neither has NVLink/IB, so treat their numbers as functional, not representative.
• Cloud (real multi-GPU): Lambda, CoreWeave, RunPod, and the hyperscalers (AWS p5/p5e = 8× H100 with NVLink + EFA; GCP a3; Azure ND H100 v5) rent 8-GPU NVLink nodes and multi-node IB clusters by the hour. An 8× H100 node is the natural unit for exercising NVLink-vs-IB collectives.
• HPC / academic: national labs and university clusters (Slurm + IB) are often the cheapest path to genuine multi-node scaling experiments.

Reading order

1. Primary — DDP internals: Li et al., “PyTorch Distributed: Experiences on Accelerating Data Parallel Training,” VLDB 2020 (arXiv:2006.15704) — the bucketing + hook-overlap design, with scaling measurements.¹⁶
2. Official design note: PyTorch “Distributed Data Parallel” notes — the authoritative description of buckets, AccumulateGrad hooks, no_sync, ordering constraints.³
3. Ring all-reduce intuition: Andrew Gibiansky, “Bringing HPC Techniques to Deep Learning” (the Baidu ring all-reduce write-up) — the clearest walk-through of the reduce-scatter/all-gather ring.⁵
4. Bandwidth-optimality proof: Patarasuk & Yuan, “Bandwidth optimal all-reduce algorithms for clusters of workstations,” JPDC 2009.⁶
5. NCCL trees: Sylvain Jeaugey, “Massively Scale Your Deep Learning Training with NCCL 2.4” (double binary trees, the $O(\log N)$ latency win).⁸
6. Hardware: NVIDIA H100 whitepaper (NVLink4/NVSwitch)¹² and the ConnectX-7 NDR datasheet.¹⁴

---

What’s next

Two threads, in order:

1. 09-zero-and-fsdp (next) — the direct answer to §9’s wall. ZeRO/FSDP shard weights, gradients, and optimizer state across the DP group, trading the all-reduce for a reduce-scatter + all-gather (the very decomposition of §5). You’ll see DP’s replication turn into sharding, and the memory-vs-communication tradeoff made quantitative.
2. 10-model-parallelism (ahead) — when even a single layer is too big, split the layer across GPUs (tensor parallel, all-reduce every layer) or the stack across GPUs (pipeline parallel). The §7 “keep chatty collectives on NVLink” rule becomes the central design constraint. All six §4 primitives reappear.
3. Revisit 07-triton-and-modern-kernels — the kernels you learned to write run inside each DP replica; distributed training is the layer wrapped around them. Compute-comm overlap (§8) is the systems-level analogue of the memory-compute overlap you tuned in kernels.

---

Reference pages: data-parallelism, collective-communication, nccl | Topic: gpu-systems-for-llms | Filed: 2026-09-02

---

References

Footer: gpu-systems-for-llms

Footnotes

1. NVIDIA H100 Tensor Core GPU Architecture Whitepaper and datasheet — H100 SXM5 dense BF16 ≈ 989–990 TFLOP/s; HBM3 at 3.35 TB/s. https://resources.nvidia.com/en-us-tensor-core [established] ↩

2. Priya Goyal et al., “Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour,” arXiv:1706.02677 (2017). Introduces the linear LR-scaling rule with warmup; trains ResNet-50 at batch 8192 without accuracy loss. https://arxiv.org/abs/1706.02677 [established] ↩

3. PyTorch documentation, “Distributed Data Parallel” design note and torch.nn.parallel.DistributedDataParallel API reference — default bucket_cap_mb = 25 MB; buckets filled in reverse model.parameters() order; all-reduce issued in bucket-index order on all ranks (mismatched order → wrong results or hang); no_sync() disables sync for gradient accumulation. https://docs.pytorch.org/docs/stable/notes/ddp.html [established] ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸

4. NVIDIA Collective Communications Library (NCCL) documentation — collective definitions (broadcast, reduce, all-reduce, reduce-scatter, all-gather, all-to-all), topology auto-detection, and NCCL_ALGO/NCCL_PROTO selection (tokens Ring, Tree, CollnetChain, CollnetDirect, NVLS, NVLSTree, PAT; protocols Simple, LL, LL128). Current documented release line is 2.31.x (Sept 2026). https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/ [recent] ↩ ↩² ↩³ ↩⁴ ↩⁵

5. Andrew Gibiansky, “Bringing HPC Techniques to Deep Learning” (2017, Baidu SVAIL) — the ring all-reduce (reduce-scatter + all-gather) popularized for DL; per-GPU communication $2\frac{N-1}{N}M$, independent of $N$. http://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/ [established] ↩ ↩²

6. Pitch Patarasuk & Xin Yuan, “Bandwidth optimal all-reduce algorithms for clusters of workstations,” Journal of Parallel and Distributed Computing 69(2):117–124 (2009). Proves the ring achieves the bandwidth lower bound $2\frac{N-1}{N}M$ per node. https://doi.org/10.1016/j.jpdc.2008.09.002 [established] ↩ ↩² ↩³

7. Rajeev Thakur, Rolf Rabenseifner & William Gropp, “Optimization of Collective Communication Operations in MPICH,” IJHPCA 19(1):49–66 (2005). The $\alpha$–$\beta$ (latency–bandwidth) cost model for collectives, ring vs recursive-doubling/halving all-reduce cost analysis. https://doi.org/10.1177/1094342005051521 [established] ↩

8. Sylvain Jeaugey (NVIDIA), “Massively Scale Your Deep Learning Training with NCCL 2.4” (2019) — double binary trees give $O(\log N)$ latency all-reduce while retaining near-full bandwidth, outperforming the ring for large node counts / small-to-medium messages. https://developer.nvidia.com/blog/massively-scale-deep-learning-training-nccl-2-4/ [established] ↩ ↩² ↩³

9. NCCL algorithm/protocol overview: ring is bandwidth-optimal (per-GPU volume $2\frac{N-1}{N}M$); double-binary tree is latency-optimal ($O(\log N)$); CollNet/SHARP and NVLS perform in-network (in-switch) reduction, cutting per-GPU volume toward $\frac{N-1}{N}M$. NCCL 2.18 added NVLS (NVLink SHARP) for Hopper; later versions added FP8 reductions and PAT. See NCCL docs⁴ and “Demystifying NCCL,” arXiv:2507.04786 (2025). https://arxiv.org/html/2507.04786v1 [recent] ↩ ↩² ↩³

10. NVIDIA Technical Blog, “NCCL Deep Dive: Cross Data Center Communication and Network Topology Awareness” (2025) — NCCL discovers NVLink/NVSwitch, PCIe trees, and NIC/IB affinity and builds hierarchical rings/trees; the fabric-ID mechanism for topology-aware algorithm choice. https://developer.nvidia.com/blog/nccl-deep-dive-cross-data-center-communication-and-network-topology-awareness/ [recent] ↩ ↩²

11. NVIDIA NVSwitch (Hot Chips 2022): NVSwitch3 has 64 NVLink4 ports (1.6 TB/s per direction per chip, 3.2 TB/s bidirectional) and an embedded ALU for in-network SHARP reductions (sum/broadcast). A 256-GPU NVLink domain provides 57.6 TB/s all-to-all bandwidth. https://hc34.hotchips.org/ [established] ↩ ↩²

12. NVIDIA H100 / Hopper architecture (Hot Chips 2022 NVSwitch presentation and H100 whitepaper): each H100 has 18 fourth-gen NVLink links × 50 GB/s = 900 GB/s bidirectional, connected via 4 on-board NVSwitch3 chips forming a non-blocking crossbar (every GPU can send 900 GB/s to any other). PCIe 5.0 x16 ≈ 64 GB/s bidirectional is the fallback path. https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/ [established] ↩ ↩² ↩³

13. NVIDIA Blackwell (GB200/B200) uses fifth-gen NVLink at ~1.8 TB/s (1800 GB/s) bidirectional per GPU (~2× NVLink4) and scales the NVLink domain to 72 GPUs (GB200 NVL72). NVIDIA Blackwell architecture materials. https://www.nvidia.com/en-us/data-center/technologies/blackwell-architecture/ [recent] ↩

14. NVIDIA ConnectX-7 NDR 400G InfiniBand adapter datasheet — 400 Gb/s (= 50 GB/s) per port via 4× 100G-PAM4 lanes; standard HGX H100 ships one CX-7 per GPU for 400 Gb/s GPU-to-fabric each; Quantum-2 switches (64×400G). RoCE (RDMA over Converged Ethernet) is the Ethernet alternative at comparable rates. https://www.nvidia.com/content/dam/en-zz/Solutions/networking/infiniband-adapters/infiniband-connectx7-data-sheet.pdf [established] ↩ ↩² ↩³

15. NVIDIA ConnectX-8 SuperNIC + Quantum-X800: InfiniBand XDR = 800 Gb/s (= 100 GB/s) per port via 4× 200G-PAM4 lanes; Quantum-X800 switch has up to 144×800G ports and SHARP v4. Current shipping frontier for B200/GB300 systems (2026). https://www.nvidia.com/en-us/networking/infiniband-adapters/ [recent] ↩

16. Shen Li et al., “PyTorch Distributed: Experiences on Accelerating Data Parallel Training,” Proc. VLDB Endowment 13(12), arXiv:2006.15704 (2020). Describes DDP’s gradient bucketing, AccumulateGrad autograd-hook-triggered async all-reduce, backward/comm overlap, and scaling results. https://arxiv.org/abs/2006.15704 [established] ↩ ↩² ↩³ ↩⁴

17. Samyam Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” arXiv:1910.02054 (SC 2020). Documents the ~16 bytes/param mixed-precision Adam state (2B fp16 weight + 2B fp16 grad + fp32 master weight + fp32 momentum + fp32 variance) that DP replicates on every GPU, motivating sharding. https://arxiv.org/abs/1910.02054 [established] ↩