Why FSDP exists: the redundancy in plain data parallelism

Standard DistributedDataParallel keeps a full copy of the model on each of N GPUs. Under mixed-precision training with Adam, every parameter drags along a lot of state: an fp16 weight (2 bytes), an fp16 gradient (2), an fp32 master copy (4), and Adam’s fp32 momentum and variance (4 + 4) — roughly 16 bytes per parameter. DDP replicates all of it everywhere, so per-GPU memory is ≈ 16P regardless of N. More GPUs buy throughput, never headroom.

The observation FSDP exploits is that this replication is pure redundancy. If you shard the parameters, gradients, and optimizer states so each rank stores only its slice, the persistent footprint drops to ≈ 16P / N. The price is that no rank holds a complete layer anymore, so the full weights have to be reconstructed on demand and thrown away again. FSDP is the machinery that makes that reconstruction cheap, overlapped, and correct.

Advertisement

The sharding unit: FlatParameter and the wrapping policy

FSDP does not shard individual tensors. Instead it groups the parameters of a wrapped submodule, flattens them into a single 1-D FlatParameter, and splits that buffer evenly across the N ranks. Each rank physically stores one contiguous 1/N slice. Working on one flat buffer means one collective per unit rather than a storm of tiny per-tensor operations — collectives love large contiguous payloads.

What counts as a ‘unit’ is set by the wrapping policy. You wrap submodules as nested FSDP instances; each wrapped module becomes one FlatParameter and one gather/reshard granularity. For transformers the idiomatic choice is transformer_auto_wrap_policy, which makes each TransformerBlock its own unit. This granularity is the single most important tuning decision: the unit size directly bounds the transient memory spike (the whole unit is gathered at once) and sets how many collectives run per step. Wrap too coarsely and you gather the whole model; too finely and the collectives get small and inefficient.

Advertisement

The forward cycle: all-gather, compute, reshard

In the forward pass FSDP walks the units in execution order. Just before a unit runs, it issues an all-gather: every rank contributes its 1/N slice of the FlatParameter, and each rank ends up with the full, unsharded weights for that one unit. The module’s forward then runs normally on the complete parameters, producing activations exactly as it would on a single device.

The moment the unit finishes, FSDP reshards — it frees the gathered full weights and keeps only the local slice again. Memory for the full parameters of that unit is reclaimed before the next unit is gathered. The consequence is the key invariant: at any instant only one unit’s worth of full parameters lives in GPU memory (a little more with prefetch), not the whole model. FSDP is essentially trading a persistent 16P footprint for a persistent 16P/N footprint plus a small, moving 2 × unit_size window of gathered weights that slides across the network layer by layer.

The backward cycle: all-gather, grads, reduce-scatter

Backward is the mirror image, with one extra collective. Because the forward resharded each unit, the full weights are gone by the time gradients flow back, so FSDP must all-gather the unit’s parameters again before it can compute that unit’s gradients. The autograd engine then produces the gradient with respect to the full parameters, just as on one device.

Those full gradients are then dispersed with a reduce-scatter: the collective both sums the gradients across all ranks (the averaging that data parallelism requires) and scatters the result so each rank keeps only the reduced gradient slice matching its own parameter slice. This is the crucial difference from DDP, which uses an all-reduce so every rank holds the whole averaged gradient. FSDP only ever needs the shard it owns, so reduce-scatter is exactly right. After it completes, the full parameters and full gradients are freed, the optimizer updates the local 1/N slice in place, and the rank’s footprint returns to 16P/N.

Per-unit memory and communication, in numbers

Two budgets govern FSDP. The persistent budget is the sharded state that lives on the GPU all step long; the transient budget is the gathered full weights of whatever units are currently active.

persistent  ≈ 16P / N            (fp16 w+g, fp32 master+Adam, sharded)
transient   ≈ 2 × U × (1 + prefetch)  (fp16 gathered params of the live unit)
where  U = params in the largest wrapped unit,  P = total params,  N = ranks

On the wire, the cost is three model-sized collectives per step. A bandwidth-optimal ring all-gather or reduce-scatter moves ≈ P per rank; an all-reduce moves ≈ 2P. FSDP pays all-gather(fwd) + all-gather(bwd) + reduce-scatter(bwd) ≈ 3P, versus DDP’s single all-reduce ≈ 2P. So full sharding costs about 1.5× the communication of DDP — the price of turning a fixed 16P memory bill into a 16P/N one — extra traffic that prefetch and overlap exist to hide.

A worked example: a 7B model on 8 GPUs

Take a 7-billion-parameter model, Adam, mixed precision, on N = 8 GPUs of 80 GB. DDP would need 16 × 7e9 = 112 GB of state per GPU — it simply will not load. FSDP shards that:

persistent = 112 GB / 8            = 14 GB per GPU
unit size  = 7B / 32 blocks       ≈ 0.22B params per TransformerBlock
transient  = 2 bytes × 0.22e9 × 2 ≈ 0.9 GB  (live unit + one prefetched)
peak state ≈ 14 + 0.9        ≈ 15 GB, leaving ~65 GB for activations

The model now fits with room to spare. Notice how the wrapping policy shows up in the arithmetic: because each of the 32 transformer blocks is its own unit, the transient spike is set by one 0.22B block, not the full 7B. Wrap the whole model as a single unit and that transient term becomes 2 × 7e9 = 14 GB gathered on every GPU at once — you would be back to holding the entire model and the memory win would evaporate. Communication is about 3 × 14 GB per rank each step, against DDP’s 28 GB.

Overlap: prefetching to hide the collectives

That 1.5× communication would cripple throughput if the GPU sat idle during every all-gather. FSDP avoids the stall by overlapping communication with computation. While a unit’s forward is computing, FSDP can already issue the all-gather for the next unit (forward_prefetch), so the weights arrive just as they are needed. In backward, backward_prefetch = BACKWARD_PRE gathers the next unit’s parameters before the current unit’s gradient computation finishes.

The overlap works only when there is enough compute per unit to hide the collective behind it — another reason unit size matters. Very small units produce collectives too short to overlap and too numerous to schedule; very large units blow the transient budget. FSDP runs collectives on a separate CUDA stream and uses limit_all_gathers (rate limiting) to stop prefetching so far ahead that the in-flight gathered weights themselves exhaust memory. Tuned well, step time approaches compute-bound despite moving more bytes than DDP.