Distributed Training Internals
The real mechanics of DDP and FSDP2 — what all-reduce and all-gather actually do, why sharding trades memory for communication, and how mixed precision avoids losing updates to rounding
Distributed Training Internals
TL;DR
DDP gives every GPU a full model copy and synchronizes gradients once per step via all-reduce — cheap communication, but the model must fit on one GPU. FSDP2 shards parameters, gradients, and optimizer state across GPUs and reconstructs each layer's full parameters on demand via all-gather — far lower peak memory, at the cost of communication on every layer's forward and backward pass, not just once per step.
| Property | Value |
|---|---|
| Level | Advanced |
| Reading time | ~20 minutes |
| Prerequisites | Alignment & DPO Math |
| You will understand | What DDP and FSDP2 actually do at the communication level, and how mixed precision avoids losing gradient updates to rounding |
DDP: A Full Copy Per GPU
In Distributed Data Parallel (DDP), every GPU holds an identical, complete copy of the model. Each processes a different shard of the current batch, computing its own forward and backward pass independently — up to that point, the GPUs know nothing about each other.
import torch.distributed as dist
# after each GPU computes its local gradients:
for param in model.parameters():
dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
param.grad /= world_size # average, not sumAll-reduce is the operation that makes this consistent: every GPU's gradient tensor is combined (summed, here) across all participants, and the result is written back to every GPU — not gathered to one and redistributed, but computed so each participant ends up holding the identical combined value. After the divide, every GPU has the exact same averaged gradient, so every GPU's subsequent optimizer step produces the exact same parameter update, and the model copies stay identical across the whole run without ever explicitly syncing the weights themselves.
DDP's communication cost is fixed and cheap: one all-reduce per step, regardless of model depth. This is why DDP is the default whenever the model actually fits on a single GPU — the communication overhead doesn't scale with how many layers the model has.
FSDP2: Sharding What DDP Duplicates
DDP's requirement — every GPU holds the whole model, its gradients, and its optimizer state — becomes the bottleneck once a model is too large for that to fit anywhere. Fully Sharded Data Parallel (FSDP2) removes the duplication: each GPU permanently holds only its own shard (slice) of every parameter, gradient, and optimizer state tensor.
What one GPU holds, DDP vs FSDP2
DDP
FSDP2
The catch: a layer can't actually compute its forward pass on a 1/N-sized sliver of its own weights. So immediately before each layer runs, FSDP2 performs an all-gather — every GPU's shard of that layer's parameters is collected and assembled into the full parameter tensor, temporarily, on every GPU. The layer computes normally. Then the full parameter is freed again, and only the GPU's own shard remains resident.
Layer N's forward pass, under FSDP2:
1. all-gather: reconstruct layer N's full parameters on every GPU (temporary)
2. compute: run the forward pass normally, using the full parameters
3. free: discard the reconstructed full parameters, keep only this GPU's shard
4. repeat for layer N+1, and again on the way back during backwardThis is a real trade, not a free lunch. DDP pays one all-reduce per optimizer step. FSDP2 pays an all-gather (and, during backward, a reduce-scatter to re-shard the computed gradients) for every layer, on both the forward and backward pass. The memory savings are large and often the only way to fit a big model at all — but the communication volume is correspondingly higher, which is why FSDP2 should be reached for because the model genuinely doesn't fit, not by default.
Why FSDP2 Uses DTensor
FSDP2's implementation represents every sharded parameter as a DTensor — a distributed tensor that knows its own shape, its sharding layout, and how to gather or reduce itself — rather than the older FSDP1 approach of flattening a whole layer's parameters into one opaque buffer.
Treating each parameter as its own DTensor is what lets FSDP2 compose cleanly with methods like LoRA that only train a subset of parameters. Because sharding is tracked per-parameter rather than per-flattened-buffer, freezing the base model's parameters while sharding and training only a small set of LoRA parameters is a straightforward composition, not a special case FSDP has to be taught separately. This is a real, current implementation detail — not a minor footnote — because it's exactly what makes "LoRA fine-tuning across many GPUs" a practical, supported combination rather than an awkward workaround.
Mixed Precision: Why a Master Copy of the Weights Exists
Training in bf16/fp16 halves memory and roughly doubles throughput on hardware that supports it — but the weights themselves usually aren't kept purely in that low precision.
Each training step, under mixed precision:
1. Forward + backward pass computed in bf16/fp16 (fast, half the memory)
2. Gradient update applied to an FP32 "master" copy of the weights
3. The updated FP32 master weights are cast back down to bf16/fp16
for the next step's forward passWhy the master copy has to be FP32: small updates get silently lost to rounding otherwise. A gradient update might be many orders of magnitude smaller than the weight it's adjusting. In bf16/fp16, that tiny delta can round to exactly zero when added to a much larger number — the update simply vanishes, and the parameter never actually changes despite the optimizer computing a real gradient for it every step. FP32 has enough precision to represent updates that small relative to the current value, so the master copy accumulates them correctly even though the forward/backward math ran in lower precision.
torch.compile: Orthogonal to Parallelism Strategy
torch.compile traces a model's forward and backward pass and compiles them into fused, optimized kernels — reducing the overhead of launching many small individual GPU operations and cutting memory-bandwidth usage by fusing operations together. It's a genuine speedup technique, but it answers a different question than DDP/FSDP2 do: those decide how work is split across GPUs; torch.compile decides how efficiently a single GPU executes the operations it's given. The two combine — compiling a model doesn't change whether it needs DDP or FSDP2, and choosing DDP or FSDP2 doesn't change whether compiling helps.
Choosing Between Them
DDP vs FSDP2 — the actual decision
DDP
RecommendedThe model, its gradients, and its optimizer state all fit comfortably on one GPU. Cheapest communication pattern, simplest mental model, the right default.
FSDP2
The model doesn't fit on one GPU at all — sharding is the only way to train it, not a throughput optimization. Worth the added communication cost because there's no alternative, not because it's generically "better."
Concept Checks
Check yourself
Why does DDP's communication cost stay the same regardless of how many layers the model has, while FSDP2's doesn't?
Because DDP performs exactly one all-reduce per optimizer step, synchronizing the already-computed full gradient tensor once — a fixed cost independent of depth. FSDP2 has to reconstruct each layer's full parameters via an all-gather immediately before that layer computes, on both the forward and backward pass, so its communication volume scales with the number of layers, not just the number of steps.
Why is a parameter update sometimes lost entirely when training purely in fp16, without a master weight copy?
Because fp16 has limited precision relative to a value's magnitude, and a gradient-scaled update can be many orders of magnitude smaller than the weight it's supposed to adjust. Adding that tiny delta to the much larger fp16 weight can round to exactly the original value, silently discarding the update. An FP32 master copy has enough precision to represent and accumulate updates that small, which is why it's maintained even though the forward/backward computation itself runs in lower precision.
Why does representing each sharded parameter as its own DTensor matter specifically for combining FSDP2 with LoRA?
Because LoRA freezes most of the base model's parameters and only trains a small additional set, and per-parameter DTensor tracking lets FSDP2 shard and manage each parameter's distributed state independently — frozen base parameters and trainable LoRA parameters can be handled according to their own needs rather than requiring the whole layer's parameters to be treated as one indivisible buffer, which was FSDP1's flattened-buffer approach.
Key Concepts Recap
| Concept | One-line summary |
|---|---|
| DDP | Full model copy per GPU, gradients synchronized once per step via all-reduce |
| All-reduce | Combines a tensor across all GPUs so every GPU ends up with the identical result |
| FSDP2 | Parameters, gradients, and optimizer state sharded across GPUs, reconstructed on demand |
| All-gather | Reconstructs a layer's full parameters from every GPU's shard, temporarily |
| The FSDP2 trade | Much lower peak memory, at the cost of communication on every layer, not just once per step |
| DTensor | A distributed tensor tracking its own sharding — what lets FSDP2 compose cleanly with LoRA |
| Master weights (FP32) | The full-precision copy that prevents small updates from rounding away to nothing |
torch.compile | Speeds up single-GPU execution; orthogonal to the DDP/FSDP2 parallelism decision |
Next
With training itself understood mechanically, the next page turns to reading what a training run is actually telling you: Evaluating & Debugging Training.
Alignment & DPO Math
The Bradley-Terry model of preference, why classic RLHF needs a separate reward model, and the full derivation of DPO's reward-model-free loss
Evaluating & Debugging Training
Reading loss curves as a real diagnostic skill, systematically tracing NaN loss to its cause, and the mechanics behind vanishing gradients and dead ReLUs