Distributed Training Theory

Why This Matters

A single GPU cannot train a frontier language model. GPT-4-class models have hundreds of billions of parameters and train on trillions of tokens. Even if you could fit the model on one GPU (you cannot), training would take decades.

Distributed training splits the work across hundreds or thousands of devices. But parallelism is not free: devices must communicate gradients, activations, or parameters, and this communication can bottleneck the entire system. The theory of distributed training is about understanding these tradeoffs and minimizing wasted computation.

Mental Model

Think of building a house. Data parallelism is hiring $N$ identical crews, each building a copy of the same wall using different bricks, then averaging their work. Model parallelism is splitting the blueprint so each crew builds a different section of the house. Pipeline parallelism is an assembly line: crew 1 pours foundations, passes to crew 2 for framing, crew 3 for roofing.

Each approach has different communication patterns. Data parallelism requires sharing gradients (the "average" step). Tensor parallelism requires sharing intermediate activations within each layer. Pipeline parallelism requires passing activations between stages. The art is choosing the right combination to minimize idle time and communication overhead.

Formal Setup and Notation

Let $f(\theta; x)$ be the loss for parameters $\theta$ on example $x$. We want to minimize $F(\theta) = \mathbb{E}_{x \sim \mathcal{D}}[f(\theta; x)]$ using SGD or its variants, distributed across $K$ workers.

Definition

Data Parallelism

In data parallelism, each of $K$ workers holds a complete copy of the model $\theta$. At each step:

1. Each worker $k$ samples a minibatch $B_k$ of size $b$
2. Each worker computes its local gradient $g_k = \frac{1}{b} \sum_{x \in B_k} \nabla f(\theta; x)$
3. Workers average their gradients: $\bar{g} = \frac{1}{K} \sum_{k=1}^{K} g_k$
4. Each worker updates: $\theta \leftarrow \theta - \eta \bar{g}$

The effective batch size is $Kb$. Step 3 is typically implemented via AllReduce.

Definition

AllReduce

AllReduce is a collective communication primitive that computes the sum (or average) of vectors across all workers and distributes the result to every worker. For $K$ workers each holding a vector of $d$ elements:

• Ring AllReduce: each worker sends and receives $d(K-1)/K$ elements in $2(K-1)$ steps. Total communication volume per worker: $2d(K-1)/K \approx 2d$ for large $K$.
• The communication cost is $O(d/\text{bandwidth})$ and is independent of $K$ in the bandwidth-optimal ring algorithm.

Definition

Tensor Parallelism

Tensor parallelism (TP) splits individual layers across devices. For a linear layer $Y = XW$, partition $W$ column-wise across $K$ devices:

$W = [W_1, W_2, \ldots, W_K]$

Device $k$ computes $Y_k = XW_k$ (a slice of the output). An AllGather collects the full output $Y = [Y_1, \ldots, Y_K]$. For the backward pass, a ReduceScatter distributes gradient computation.

Definition

Pipeline Parallelism

Pipeline parallelism (PP) partitions the model's layers into $S$ sequential stages, each on a different device. A microbatch flows through stage 1, then stage 2, etc. Multiple microbatches are in flight simultaneously to fill the pipeline, but pipeline bubbles (idle time) occur at startup and shutdown.

For $M$ microbatches and $S$ stages, the bubble fraction is:

$\text{bubble} = \frac{S - 1}{M + S - 1}$

which shrinks as $M$ grows relative to $S$.

Core Definitions

ZeRO (Zero Redundancy Optimizer) from DeepSpeed eliminates memory redundancy in data parallelism. Standard data parallelism replicates the full optimizer state, gradients, and parameters on every worker. ZeRO partitions these across workers:

• ZeRO Stage 1: partition optimizer states (e.g., Adam's $m$ and $v$). Saves $4\times$ memory for Adam (from 12 bytes/param to 4 bytes/param per worker).
• ZeRO Stage 2: additionally partition gradients. Workers compute full gradients locally but only store their partition.
• ZeRO Stage 3: additionally partition parameters. Workers must AllGather parameters before each forward/backward step, trading communication for memory.

Gradient compression reduces communication volume:

• Gradient quantization: send gradients in low precision (e.g., FP16 or INT8) instead of FP32, halving or quartering bandwidth.
• Gradient sparsification: send only the top-$k$ gradient components. The error from dropped components is accumulated locally and sent in future rounds.
• PowerSGD: approximate the gradient matrix with a low-rank decomposition, reducing communication from $O(d)$ to $O(d \cdot r / \min(m,n))$ for rank $r$.

Main Theorems

Theorem

Linear Speedup of Data-Parallel SGD

Statement

For data-parallel SGD with $K$ workers, each using local batch size $b$, the effective batch size is $B = Kb$. After $T$ steps, the convergence rate for a smooth non-convex objective satisfies:

$\frac{1}{T} \sum_{t=1}^{T} \mathbb{E}\|\nabla F(\theta_t)\|^2 \leq O\left(\frac{F(\theta_0) - F^*}{\eta T} + \eta L \frac{\sigma^2}{Kb}\right)$

With learning rate $\eta = O(\sqrt{Kb/T})$, this gives:

$O\left(\sqrt{\frac{(F(\theta_0) - F^*)L\sigma^2}{KbT}}\right)$

This is a factor of $\sqrt{K}$ better than single-worker SGD with batch size $b$ and the same number of steps. Equivalently, to reach the same accuracy, data-parallel SGD needs $K$ times fewer steps, giving linear speedup.

Intuition

Each worker contributes independent gradient estimates. Averaging $K$ independent estimates reduces variance by a factor of $K$, just like averaging $K$ random variables. Less variance means you can take larger steps or reach the same accuracy faster. The speedup is linear until the batch size becomes so large that the variance reduction saturates (the "critical batch size").

Proof Sketch

Start from the smoothness inequality: $F(\theta_{t+1}) \leq F(\theta_t) - \eta \nabla F(\theta_t)^T \bar{g}_t + \frac{\eta^2 L}{2} \|\bar{g}_t\|^2$. Take expectations. The averaged gradient $\bar{g}_t$ has variance $\sigma^2 / (Kb)$ (each of $K$ workers averages $b$ independent samples). The bias is zero since each worker computes unbiased gradients. Telescope over $T$ steps and optimize the learning rate.

Why It Matters

Linear speedup is the best you can hope for from parallelism. This theorem says data-parallel SGD achieves it, up to a critical batch size beyond which increasing $K$ gives diminishing returns. Understanding this limit is essential for choosing the right number of workers and avoiding wasted compute.

Failure Mode

Linear speedup breaks down when: (1) the batch size $Kb$ exceeds the critical batch size $B_{\text{crit}} \approx \sigma^2 / \|\nabla F\|^2$, after which additional workers reduce variance below what matters; (2) communication overhead exceeds computation savings; (3) large batch sizes require learning rate warmup and careful tuning to maintain convergence.

Proposition

Communication-Computation Tradeoff

Statement

In data-parallel training with ring AllReduce, the total time per step is:

$T_{\text{step}} = C + \frac{2d}{W}$

where $C$ is the local computation time (forward + backward pass) and $2d / W$ is the AllReduce communication time (sending $2d$ values at bandwidth $W$). The computation-to-communication ratio is:

$\text{ratio} = \frac{C}{2d / W} = \frac{CW}{2d}$

Training is efficient when this ratio is much greater than 1 (computation dominates). It is communication-bound when the ratio is less than 1.

Intuition

Every step, you do some math (forward and backward pass, time $C$) and some networking (share gradients, time $2d/W$). If the network is slow relative to the GPU, workers spend most of their time waiting. Making each worker do more computation per communication round (larger local batch, gradient accumulation) improves the ratio.

Proof Sketch

Ring AllReduce for $d$ parameters across $K$ workers requires each worker to send and receive $d(K-1)/K \approx d$ values in each of two phases (reduce- scatter and all-gather). At bandwidth $W$, this takes $2d/W$ time. Computation and communication can overlap partially (compute gradients for layer $l+1$ while communicating gradients for layer $l$), but the non-overlapping portion still imposes the bound.

Why It Matters

This ratio determines whether adding more GPUs actually helps. For a model with $d = 10^9$ parameters in FP16 ($2 \times 10^9$ bytes), AllReduce takes $4 \times 10^9 / W$ seconds. At 100 Gbps bandwidth ($\approx 12.5$ GB/s), this is about 0.32 seconds. If the forward/backward pass takes 0.5 seconds, communication takes 39% of step time, a significant overhead. Faster interconnects (InfiniBand, NVLink) or gradient compression are needed.

Canonical Examples

Example

Scaling from 1 to 256 GPUs

A model with $d = 7 \times 10^9$ parameters (7B) in FP16 uses 14GB for parameters alone. Adam optimizer states add another 28GB (FP32 copy + $m$ + $v$). Total: 42GB per worker in standard data parallelism.

With ZeRO Stage 1: optimizer states are split across $K$ workers. Each worker stores $28/K$ GB of optimizer state plus 14GB of parameters and gradients. At $K = 8$, this is $3.5 + 14 = 17.5$ GB per worker, fitting on 24GB GPUs.

With ZeRO Stage 3: everything is split. Each worker stores $(14 + 14 + 28)/K$ GB. At $K = 8$, this is 7GB per worker, but each forward and backward pass requires AllGather operations to reconstruct parameters.

Common Confusions

Watch Out

Data parallelism does not change the model

Data parallelism with synchronized gradients produces exactly the same updates as single-GPU training with the same effective batch size $Kb$. The only difference is speed. Asynchronous data parallelism, where workers do not wait for all gradients, does change the dynamics and can hurt convergence.

Watch Out

Pipeline parallelism is not just splitting layers

Naive pipeline parallelism has the pipeline bubble problem: only one stage is active at a time, giving utilization $1/S$. The key innovation (GPipe, PipeDream) is microbatching: split the minibatch into $M$ microbatches and overlap their execution. With $M \gg S$, utilization approaches 100%, but memory increases because activations for all in-flight microbatches must be stored.

Watch Out

3D parallelism is not optional at scale

Frontier models use all three forms simultaneously: data parallelism across groups, tensor parallelism within a node (fast NVLink), and pipeline parallelism across nodes. The 3D configuration is a system design choice driven by hardware topology, not just model architecture.

Summary

• Data parallelism gives linear speedup up to a critical batch size; beyond that, adding workers wastes compute
• AllReduce communication cost is $O(d/W)$, independent of worker count in the bandwidth-optimal case
• Tensor parallelism splits layers within a node (requires high bandwidth); pipeline parallelism splits layers across nodes (tolerates lower bandwidth)
• ZeRO eliminates memory redundancy: Stage 1 partitions optimizer states, Stage 2 adds gradients, Stage 3 adds parameters
• The computation-to-communication ratio determines whether training is GPU-bound or network-bound

Exercises

ExerciseCore

Problem

A model has $d = 1 \times 10^9$ parameters in FP32. Each worker has 100 Gbps network bandwidth. How long does AllReduce take, assuming the ring algorithm?

Hint

Reveal Solution

ExerciseAdvanced

Problem

With $K = 64$ workers and local batch size $b = 16$, the effective batch size is $1024$. If the critical batch size is $B_{\text{crit}} = 2048$, are we in the linear speedup regime? What happens if we double $K$ to $128$?

Hint

Reveal Solution

ExerciseResearch

Problem

Explain the memory-communication tradeoff in ZeRO Stages 1-3. For each stage, what memory is saved per worker and what additional communication is required?

Hint

Reveal Solution

References

Canonical:

• Dean et al., Large Scale Distributed Deep Networks (NIPS 2012)
• Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (SC 2020)

Current:

• Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (2020)
• Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (SC 2021)

Next Topics

Distributed training connects to:

• Scaling laws that determine how many GPUs and how much data you need
• Mixed-precision training for reducing communication and memory further