Distributed Training

One model. Too big for one GPU. Many GPUs that must agree on every gradient. This article is about the engineering that makes that work — and the communication bill it racks up.

What it is (intuition first)

A modern frontier model has hundreds of billions to trillions of parameters. Just storing one for training is brutal: a trillion-parameter model needs roughly 14 terabytes of GPU memory once you count weights, gradients, and optimizer state — but a single high-end accelerator holds only 64–192 GB. So you cannot train it on one chip. You distribute the work across thousands of GPUs.

There are two distinct problems, and people constantly conflate them:

1. I have enough memory, I just want to go faster. → Replicate the model on N GPUs, give each a different slice of the data, average the gradients. This is data parallelism. It's the easy one.
2. The model doesn't even fit on one GPU. → You must cut the model itself into pieces that live on different GPUs. This is model parallelism, and it comes in two flavors: tensor parallelism (cut each layer's matrices) and pipeline parallelism (put different layers on different GPUs).

Real frontier training uses all three at once — that's "3D parallelism" — plus a memory trick called ZeRO/FSDP that shards the bookkeeping. The whole game is a balance: every way you split the model adds communication between GPUs, and communication is slow relative to compute. The art is hiding that communication behind useful work.

The mental model to hold onto: compute is cheap, moving bytes between chips is expensive. Every design decision below is downstream of that one fact.

Why it matters

• It is the entire reason large models exist. GPT-3, Llama, DeepSeek, Gemini — none of them could be trained without these techniques. The algorithm (a transformer + Adam) is almost boringly standard; the systems engineering is what's hard and what differentiates labs.
• It sets the cost. Training cost ≈ (FLOPs required by [[scaling-laws|scaling laws]]) ÷ (FLOPs your cluster actually delivers). That denominator — Model FLOPs Utilization (MFU) — is decided almost entirely by how well your parallelism strategy hides communication. A good 3D-parallel setup hits ~50% MFU; a naive one can drop to single digits and 10× your bill.
• It's a hardware-software co-design problem. The right strategy depends on your interconnect. NVLink inside a node moves ~900 GB/s; InfiniBand between nodes is far slower. Tensor parallelism is fine within a node and disastrous across nodes. You can't pick a strategy without knowing the wiring.

To make it concrete: the SC'21 Megatron-LM work trained a 1-trillion-parameter model at 502 petaFLOP/s on 3072 GPUs — about 52% of theoretical peak per GPU. That 52% is the headline result; everything in this article is in service of pushing that number up.

How it works (real mechanics)

0. The baseline: what does training actually need in memory?

Take a model with Ψ parameters trained in mixed precision (fp16/bf16 compute, fp32 master weights — the standard recipe) with the Adam optimizer. Per parameter you store:

The last three rows — 12Ψ bytes — are the optimizer states. Total: 2Ψ + 2Ψ + 12Ψ = 16Ψ bytes (ZeRO, Rajbhandari et al. 2019). For a 7.5B model that's ~120 GB before activations — already over one GPU. This 16Ψ breakdown is the number every memory-optimization technique attacks.

1. Data parallelism (DP) — the easy axis

Replicate the full model on each of N_d GPUs. Split the global batch into N_d shards. Each GPU does a full forward+backward on its shard, producing local gradients. Then a single all-reduce averages gradients across all GPUs so every replica applies the identical update.

for each step:
    local_grad = backward(forward(model, my_shard))
    all_reduce(local_grad)            # average across all N_d GPUs
    optimizer.step(model, local_grad) # identical update everywhere

• Pro: trivially simple, near-linear speedup if communication hides behind compute.
• Con: every GPU stores the full 16Ψ. DP solves throughput, not memory. The all-reduce volume is ~2Ψ bytes per step (ring all-reduce), independent of N_d.

2. Gradient accumulation — a free virtual batch size

Sometimes the per-GPU batch you need (for [[scaling-laws|compute-optimal]] training, large batches are desirable) doesn't fit in memory. Gradient accumulation runs K micro-batches sequentially, summing gradients, and only steps the optimizer (and all-reduces) once every K micro-batches:

for k in range(K):                    # K micro-batches
    grad += backward(forward(model, micro_batch[k]))
all_reduce(grad); optimizer.step()    # one update per K micro-batches

Effective batch size = per_micro_batch × K × N_d — with the memory of a single micro-batch. It also amortizes the all-reduce over K micro-batches (less frequent communication) and, as we'll see, fills pipeline bubbles. It's the cheapest knob in the whole toolbox.

3. ZeRO / FSDP — shard the redundancy, keep DP's simplicity

In plain DP, all N_d replicas store identical copies of the 16Ψ state. That's N_d-fold redundancy. ZeRO (Zero Redundancy Optimizer, DeepSpeed) and its PyTorch-native twin FSDP (Fully Sharded Data Parallel) partition that state across the DP group, reconstructing pieces on-the-fly only when needed:

At ZeRO-3, per-GPU memory falls to 16Ψ/N_d — linear in the number of devices, which is what unlocks models that don't fit on one GPU without model parallelism. DeepSpeed's empirical demo: a 1.5B model's optimizer states drop from 18 GB on one device to 2.25 GB across 8 (the 8× of ZeRO-1).

How ZeRO-3/FSDP actually runs (the part the abstracts hide):

• Forward: before computing layer i, all-gather that layer's sharded parameters from all peers → full weights → compute → immediately free the gathered weights.
• Backward: all-gather again, compute gradients, then reduce-scatter the gradients so each GPU keeps only its shard.

The cost is replacing DP's one all-reduce with per-layer all-gather (forward) + all-gather + reduce-scatter (backward) — roughly 1.5× the communication volume for a large memory win. The reason it's still fast is the next idea.

4. Communication–computation overlap (the secret sauce)

A naive ZeRO-3 step would all-gather layer i, wait, compute layer i, all-gather layer i+1, wait… — communication and compute serialized, GPUs idle half the time. FSDP/ZeRO instead prefetch: while the GPU computes layer i, the network is already all-gathering layer i+1's parameters. Communication runs concurrently with compute and is hidden behind it.

compute(layer_i)  ───────────────►  (GPU busy)
all_gather(layer_{i+1}) ──────────►  (network busy, overlapped)
                                     then compute(layer_{i+1})…

This overlap is the difference between a strategy that scales near-linearly and one that's communication-bound. The governing intuition: comm is free if it finishes before the compute it overlaps with does. When compute/interconnect ratio is bad (e.g. sharding across slow inter-node links), overlap can't hide it and throughput collapses — which is why ZeRO-3 is happiest within fast NVLink domains.

5. Tensor parallelism (TP) — cut the matrices (Megatron)

When a single layer's weights are too big, tensor parallelism splits the matrix multiplies within a layer across GPUs. Megatron-LM does this cleverly to minimize sync points.

For the transformer MLP Y = GeLU(XA)·B:

• Split A column-wise → A = [A₁, A₂]. Each GPU computes GeLU(X·Aᵢ) independently (GeLU is element-wise, no sync needed).
• Split B row-wise → B = [B₁; B₂]. Each GPU computes its partial Yᵢ = GeLU(X·Aᵢ)·Bᵢ, then a single all-reduce sums the partials.

Self-attention splits analogously: each head's QKV go to different GPUs (column-parallel), the output projection is row-parallel. The result: one all-reduce in the forward pass and one in the backward pass, per layer (per MLP and per attention block). That's minimal sync — but it's a blocking all-reduce on the critical path every layer.

• Pro: lets a single huge layer fit; low count of communications.
• Con: those all-reduces are large and synchronous → TP demands NVLink-class bandwidth. You keep TP inside a node (TP degree ≤ GPUs-per-node, typically 8) and never stretch it across slow inter-node links.

6. Pipeline parallelism (PP) — cut the layers, and the bubble

Put layers 1–8 on GPU A, 9–16 on GPU B, etc. — a pipeline of stages. The catch: stage B is idle until stage A produces its output, and idle again waiting for the backward pass. That idle time is the pipeline bubble.

The fix is to split the batch into m micro-batches and feed them through the pipeline staggered (GShard/PipeDream-style 1F1B schedule), keeping every stage busy most of the time. With p pipeline stages and m micro-batches, the bubble fraction is:

bubble fraction = (p − 1) / m

More micro-batches m → smaller bubble (this is exactly where gradient accumulation pays off: K accumulation steps are your micro-batches filling the pipe). Megatron's PTD-P added an interleaved 1F1B schedule — each GPU holds several non-contiguous layer chunks — cutting the bubble further and lifting throughput >10% at comparable memory.

• Pro: cheap communication — only activations cross stage boundaries (point-to-point), not weights or gradients. Tolerates slower inter-node links.
• Con: the bubble, plus load-balancing layers across stages is fiddly.

7. 3D parallelism — compose all three (PTD-P)

Frontier training stacks them so each axis is matched to the right hardware level:

Tensor parallel   →  WITHIN a node  (needs NVLink, blocking all-reduce/layer)
Pipeline parallel →  ACROSS nodes   (cheap point-to-point activations, tolerates IB)
Data parallel     →  ACROSS the rest (+ ZeRO sharding for memory)

A 1T-parameter run might be TP=8 (one node), PP=16 (16 nodes per pipeline replica), DP=many (replicate the whole pipeline). Megatron's PTD-P composition is what achieved 502 PFLOP/s on 3072 GPUs at 52% per-GPU peak for a trillion parameters (SC'21). The naming "3D" is literal: each GPU is addressed by a (tp, pp, dp) coordinate.

Key ideas & tradeoffs

• Memory vs. communication is the master tradeoff. DP = max memory, min comm. ZeRO-3 = min memory, ~1.5× comm. You buy memory with bandwidth.
• Match the parallelism axis to the interconnect. TP on NVLink (fast, blocking), PP across InfiniBand (slow-tolerant, point-to-point), DP/ZeRO wherever's left. Getting this mapping wrong is the #1 cause of bad MFU.
• ZeRO/FSDP vs. tensor parallelism are different answers to "model too big." ZeRO-3 keeps the DP programming model (shard + reconstruct, gather whole layers). Megatron TP changes the math (split the matmuls). ZeRO-3 is simpler and composes with everything; Megatron gives finer compute/memory IO control. Modern stacks (FSDP2, TorchTitan, Megatron-Core) compose both.
• Overlap or die. Every sharding/parallelism scheme is only as good as its prefetching. The same algorithm with vs. without overlap can differ 2× in throughput.
• Gradient accumulation is the universal lubricant. It scales batch size, amortizes all-reduces, and fills pipeline bubbles — for free, modulo sequential micro-batch latency.
• The frontier in 2025–26 is lower-precision + MoE. Llama 4 Behemoth (Meta, training since late 2025): ~2T total params / 288B active across 16 experts ([[mixture-of-experts|MoE]]), trained in FP8 on 32,000 H100s. FP8 halves the bytes moved over the network — communication reduction has become a first-class precision decision (see [[quantization]]).

Honest caveats

• arXiv abstracts don't expose the mechanics; the deep numbers here come from the papers' bodies + canonical secondary sources (DeepSpeed docs, the Baseten 2026 stack survey, the nanotron Ultra-Scale Playbook). The headline figures (502 PFLOP/s, 52% MFU, ZeRO's 8× optimizer-state reduction, 1T@14TB, Llama 4 FP8/32k-H100) are quoted from those sources; treat second-decimal precision as approximate. I could not retrieve the full body text of the ZeRO and Megatron PDFs or the Ultra-Scale Playbook via fetch — those pages returned metadata only — so the per-layer all-reduce counts and the (p−1)/m bubble formula are stated from the established literature, which is consistent across DeepSpeed/Megatron docs, not lifted verbatim from those specific PDFs. Verify exact constants against the primary PDFs before citing in a paper.
• MFU numbers are workload- and hardware-specific. "52% of peak" is for that 2021 run on those GPUs. Your mileage varies enormously with model shape, sequence length, and cluster topology. There is no universal "good" number.
• These techniques don't change the loss — only whether you can afford to compute it. Distributed training is (almost) mathematically equivalent to single-GPU training; the parallelism is invisible to the model. (Caveat: large-batch training via accumulation/DP interacts with optimizer dynamics — learning-rate warmup and batch-size scaling matter.)
• The field moves fast and frameworks churn. DeepSpeed ZeRO, PyTorch FSDP→FSDP2, Megatron-Core, TorchTitan all overlap and compete. The concepts (DP/TP/PP, shard-the-16Ψ, overlap comm) are stable; the APIs are a moving target. The 2026 OSS survey's blunt verdict: "open-source LLM training is a mess."
• MoE adds its own parallelism axis (expert parallelism) and its own headache (all-to-all routing + load balancing) not covered in depth here — see [[mixture-of-experts]] and [[deepseek-architecture]].

How it connects to OpenAlice + the Academy ladder

OpenAlice does not train trillion-parameter models — but the ideas here are the systems backbone that everything in the Academy's training track sits on, and several library articles are direct rungs:

• [[scaling-laws]] tells you how much compute a target loss needs; distributed training is how you deliver that compute. They're two halves of one equation: scaling laws set the numerator (required FLOPs), distributed training sets the denominator (FLOPs/sec actually achieved via MFU). Read scaling-laws first to know why you'd ever need 3072 GPUs.
• [[nanogpt]] is the single-file pedagogical entry point — Karpathy's nanoGPT includes a working DDP (data-parallel) path, the gentlest possible first contact with this whole topic. If "all-reduce the gradients" felt abstract above, nanoGPT's ~20 lines of DDP make it concrete.
• [[llm-c]] is the next rung up: GPT-2/GPT-3 training in raw C/CUDA with hand-written mixed-precision kernels and multi-GPU/multi-node distributed training, where you watch MFU and the communication overlap directly instead of behind a framework. It's the bridge from "I understand DP conceptually" to "I see the bytes move."
• Adjacent rungs: [[mixture-of-experts]] and [[deepseek-architecture]] (expert parallelism, the frontier MoE recipe), [[quantization]] (FP8/low-precision training — now a communication-reduction lever), [[flash-attention]] (the kernel that makes the per-GPU compute efficient enough for overlap to matter).

Suggested Academy ladder: [[nn-zero-to-hero]] → [[nanogpt]] (meet DDP) → [[scaling-laws]] (why scale) → distributed-training (this article — how to scale) → [[llm-c]] (do it in CUDA) → [[mixture-of-experts]] / [[deepseek-architecture]] (the frontier recipe). That path takes a beginner from "what's a gradient" to "I understand how Llama 4 was trained on 32,000 GPUs" without skipping a step.