Stop Replicating What You Can Shard: ZeRO and FSDP

The redundancy that data parallelism leaves on the table

In plain data parallelism every one of the N replicas holds the full optimizer state, the full gradients, and the full weights. After the all-reduce, every device has bit-for-bit identical gradients — so *N − 1* of those copies are pure waste. The insight behind the Zero Redundancy Optimizer (ZeRO) is that data parallelism does not require replication; it only requires that, at the moment a tensor is needed, the right device has it. Everything else can be sharded across the N ranks and gathered on demand.

M_{\text{replica}} = \underbrace{2\Psi}_{\text{fp16 weights}} + \underbrace{2\Psi}_{\text{fp16 grads}} + \underbrace{12\Psi}_{\text{fp32 master, }m,v} = 16\Psi

Mixed-precision Adam makes every data-parallel replica store 16 bytes per parameter — exactly the redundancy ZeRO removes.

The three stages of ZeRO

Stage 1 — shard the optimizer states. Each rank owns the Adam moments for only 1/N of the parameters and updates only that slice. This alone removes the largest of the four memory terms.
Stage 2 — also shard the gradients. Replace the global all-reduce with a reduce-scatter so each rank receives only the gradient slice it is responsible for updating.
Stage 3 — shard the parameters too. No rank holds the whole model; weights for a layer are all-gathered right before that layer runs and freed right after.

\begin{aligned} P_{\text{os}}\ (\text{Stage 1}) &: 4\Psi + \frac{12\Psi}{N}\\ P_{\text{os}+g}\ (\text{Stage 2}) &: 2\Psi + \frac{14\Psi}{N}\\ P_{\text{os}+g+p}\ (\text{Stage 3}) &: \frac{16\Psi}{N} \end{aligned}

Per-device memory under the three ZeRO stages: each stage shards one more term, and Stage 3 divides every term by N.

Stage 3 is where the memory really collapses: peak weight memory per device drops by a factor of N, so adding more devices lets you train a larger model, not just a larger batch. The cost is more communication — an all-gather on the forward pass and another on the backward — which is exactly why the overlap discipline from the previous guide is non-negotiable here. The widely used implementation is DeepSpeed-ZeRO, which also offers CPU and NVMe offload to push the sharded states off the GPU entirely when you are willing to trade bandwidth for capacity.

FSDP: the same idea, layer-native

Fully Sharded Data Parallel (FSDP) is the same ZeRO-Stage-3 idea expressed as a wrapping of the model's modules. You group layers into FSDP units; each unit's parameters live sharded across ranks, are all-gathered into a full tensor just before the unit's forward (or backward) runs, and are immediately re-sharded afterward so only one unit's full weights are resident at a time. Because the unit boundary controls both memory and the all-gather granularity, choosing how to wrap is the main performance knob: too fine and you pay launch overhead, too coarse and the transient full-weight spike grows.

\underbrace{2\Psi}_{\text{DP all-reduce}} \;\longrightarrow\; \underbrace{\Psi}_{\text{all-gather (fwd)}} + \underbrace{\Psi}_{\text{all-gather (bwd)}} + \underbrace{\Psi}_{\text{reduce-scatter}} = \underbrace{3\Psi}_{1.5\times}

FSDP/ZeRO-3 trades plain data parallelism's 2Ψ all-reduce for a 3Ψ all-gather + reduce-scatter pattern — 1.5× the communication.

# Wrap each transformer block as its own FSDP unit
for block in model.transformer_blocks:
    fsdp_wrap(block)               # params now sharded across ranks
fsdp_wrap(model)                   # root unit handles embeddings/head

# Forward of one block, conceptually:
#   all_gather(block.params) -> compute -> reshard(block.params)

Per-block FSDP wrapping: only one block's full weights are ever materialized.

Two memory levers you compose with sharding

Sharding attacks the weight/gradient/optimizer terms, but it does nothing for activations — for that you reach for activation recomputation, which we develop fully in the model-parallel guide. The two are complementary: ZeRO/FSDP shrink the parameter footprint while recomputation shrinks the activation footprint, and large runs use both at once. Sharding also reshapes how you save state to disk: there is no single rank that owns the whole model, so checkpoints are written as sharded checkpoints, each rank dumping its own slice in parallel.

\underbrace{O(L)}_{\text{store every layer}} \;\longrightarrow\; \underbrace{O(\sqrt{L})}_{\sqrt{L}\text{ checkpoints}} \quad (\text{one extra forward pass})

Activation recomputation keeps only √L checkpoints, cutting activation memory from O(L) to O(√L) for one extra forward pass.