Fully Sharded Data Parallel (FSDP) is PyTorch’s native answer to a hard limit: a model whose parameters, gradients, and optimizer states do not fit on one GPU. Plain data parallelism (DDP) replicates all of that state on every worker, so the per-GPU memory never shrinks no matter how many GPUs you add. FSDP takes the ZeRO-3 idea — shard the parameters, gradients, and optimizer states across the data -parallel workers so each rank owns only 1/N of them — and wraps it in a concrete execution model built on flattened parameter buffers, an explicit wrapping policy, and a gather–compute–reshard rhythm that reconstructs each layer just in time. This piece walks that model from the sharding unit up: how the FlatParameter and wrapping policy set the granularity, what the forward and backward collectives move, the per-unit memory and communication arithmetic with a worked example, and the knobs — prefetch, CPU offload, hybrid sharding — that separate FSDP from ZeRO-3 in the abstract.

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.

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.

Advertisement

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.

CPU offload: trading bandwidth for capacity

When even 16P/N will not fit, FSDP can push the sharded state off the GPU entirely with CPUOffload(offload_params=True). The parameter shards, gradient shards, and optimizer states live in CPU RAM; FSDP copies a shard up to the GPU over PCIe only when its unit is about to be gathered, and the optimizer step itself runs on the CPU.

This buys a large capacity increase — host RAM is far cheaper and more plentiful than HBM — at the cost of PCIe transfer latency and slower CPU-side updates. It is the right knob when the alternative is not training at all: a model that overflows GPU memory even after full sharding, or a setup with few GPUs. The transfers overlap partly with compute like the collectives, but PCIe is far slower than NVLink, so offload usually cuts throughput noticeably. Treat it as a capacity lever, not a speed one — reach for it only when you must fit the model.

Hybrid sharding (HSDP): shard inside the node, replicate across

Full sharding all-gathers over every rank, and when those ranks span many nodes the collectives cross the slow inter-node fabric. Hybrid Sharding (HSDP), selected with ShardingStrategy.HYBRID_SHARD, splits the difference: it shards the model within a group — typically the 8 GPUs of one node linked by fast NVLink — and replicates that shard across groups.

The payoff is a communication pattern matched to the hardware. The expensive all-gather and reduce-scatter stay intra-node on the fast link, while the only cross-node traffic is a gradient all-reduce between replica groups — the same cheap collective DDP uses. Memory per GPU is 16P / group_size (shard just enough to fit within a node), and you deliberately stop sharding beyond that point to avoid dragging all-gathers over the network. HSDP is the practical sweet spot at large scale: shard just enough to fit the model in a node’s HBM, then scale out by replication, not ever-wider sharding.

FSDP vs raw ZeRO-3: the execution model and its knobs

ZeRO-3 is the idea — partition parameters, gradients, and optimizer states — and FSDP is one concrete implementation of it, with its own mechanics. The distinguishing choices are the FlatParameter (a single flattened, evenly-sharded buffer per wrapped unit, rather than per-parameter partitioning) and the nn.Module wrapping policy as the explicit control over sharding granularity, tied directly into autograd hooks that fire the gather and reshard.

The knob that most reveals the execution model is reshard_after_forward. With FULL_SHARD (true) the unit is resharded after forward and re-gathered in backward — strict ZeRO-3, minimum memory, that third collective. With SHARD_GRAD_OP (false) the gathered weights are kept between forward and backward, so the backward all-gather disappears — ZeRO-2-like, less communication for more memory. Alongside it sit sharding_strategy (FULL_SHARD / SHARD_GRAD_OP / HYBRID_SHARD / NO_SHARD), auto_wrap_policy, backward_prefetch, cpu_offload, and use_orig_params. Read FSDP as ‘ZeRO-3 plus these levers’ — granularity, reshard timing, prefetch, offload, hybrid — to reason about its memory and speed.

FSDP is PyTorch’s ZeRO-3: it shards parameters, gradients, and optimizer states so each rank holds only 1/N of the state, turning a fixed 16P per-GPU memory bill into a 16P/N one. The engine is a gather–compute–reshard rhythm over FlatParameter units set by the wrapping policy: forward all-gathers each unit, computes, and reshards; backward all-gathers again, computes gradients, and reduce-scatters them to the owning shard. That costs about 1.5x DDP’s communication, which prefetch overlap is designed to hide. The unit size is the pivotal knob — it bounds the transient gathered-weight spike, so per-block wrapping keeps the peak tiny while whole-model wrapping erases the win. Reach for CPU offload when even the sharded state overflows HBM, and for HSDP to shard within a node and replicate across it. Read FSDP as ZeRO-3 plus levers — granularity, reshard timing, prefetch, offload, hybrid — and its memory and speed follow.