Skip to content
mlmentorship

FSDP and ZeRO: sharding optimizer state, gradients, and parameters

How modern training scales beyond a single GPU's memory by partitioning the optimizer state, gradients, and parameters across the data-parallel group.

Published · 6 min read ·Role-specific ·Advanced

Visual quick review

Visual first · depth when needed

Distinguish the parameter, gradient, and optimizer shards that persist on each data-parallel rank from the full wrapped unit that exists only temporarily for forward and backward compute.

Preparing the visual…

Summary

ZeRO stages progressively partition optimizer state, gradients, and model parameters across data-parallel GPUs. Fully Sharded Data Parallel (FSDP) partitions all three persistent state types. Earlier ZeRO stages leave some state replicated.

Training memory has four big consumers:

ComponentBytes per parameter (BF16 + FP32 master + Adam)
Parameters (BF16)2
Gradients (BF16)2
Adam first moment (FP32)4
Adam second moment (FP32)4
Master weights (FP32, optional)4
Total~12–16 bytes/param

A 7B-parameter model therefore needs about 84–112 GB of persistent training state before activations and temporary buffers. The lower value omits separate FP32 master weights. The exact value depends on optimizer and framework storage. Ideal full sharding spreads persistent state across GPUs, but each GPU also needs transient gathered parameters, communication buffers, and activations.

Learning objective

Separate persistent shards from the temporary full unit needed for compute.

Lifecycle of one wrapped unit on one rank under FSDP full sharding At rest, rank i stores only parameter shard P i, a slot for gradient shard G i, and optimizer shard O i. Before forward, an all-gather reconstructs the current wrapped unit's full parameters P zero through P three temporarily. Forward compute runs, then the full parameters are freed while P i remains. Before backward, another all-gather reconstructs the full unit. Backward compute produces gradients, and reduce-scatter leaves gradient shard G i on this rank. The local optimizer shard O i uses G i to update P i. Activations and communication buffers are additional memory and are not shown. RANK i · ONE WRAPPED UNIT · FULL_SHARD Persistent rank-local state PARAMETER SHARD Pᵢ GRADIENT SLOT Gᵢ OPTIMIZER SHARD Oᵢ all-gather parameter shards Temporary forward window FULL CURRENT UNIT ON THIS RANK P₀ | P₁ | P₂ | P₃ → forward compute reshard · free full parameters BETWEEN COMPUTE WINDOWS keep Pᵢ, not full P all-gather parameter shards again Temporary backward window FULL CURRENT UNIT ON THIS RANK backward compute → full gradients reduce-scatter gradients · free full unit local step: Oᵢ + Gᵢ updates Pᵢ
Read it this way: read the narrow shard boxes as what rank i owns across the step. The wide boxes are temporary peaks: full parameters for only the active wrapped unit are all-gathered for forward and again for backward, then freed. Reduce-scatter returns only Gᵢ, so the local optimizer shard can update Pᵢ. Activations and communication buffers still add to peak memory. Original schematic checked against the ZeRO paper, PyTorch FSDP documentation, and DeepSpeed ZeRO documentation.

ZeRO (Rajbhandari et al., 2019, DeepSpeed) and PyTorch FSDP implement this idea. They are the standard for any training run that doesn’t fit in a single GPU’s memory and doesn’t need full tensor or pipeline parallelism.

The three stages (ZeRO-1/2/3)

Stage 1: shard optimizer state

Each GPU holds the full parameters and gradients but only of the Adam moments. After the backward pass, each GPU updates its slice and then all-gathers updated parameters. Memory reduction: up to about 4× for large groups, depending on master-weight storage.

Stage 2: shard optimizer state + gradients

Same as Stage 1 plus gradients are reduced-scattered (each GPU keeps its slice) instead of all-reduced. Memory reduction: up to about 8× for large groups under the 16-byte assumption.

Stage 3 (FSDP): shard optimizer state + gradients + parameters

Each GPU holds only its persistent slice of the parameters. It all-gathers one wrapped unit before computing that unit. If the full parameters are freed after forward, they must be gathered again for backward. Keeping them through backward saves that second gather but raises peak memory. Memory reduction: nearly × for persistent model state. Peak memory also includes at least one gathered unit and temporary buffers.

PyTorch FSDP and DeepSpeed ZeRO-3 are common implementations of full parameter sharding.

Tradeoffs

  • Memory vs. communication: each stage trades more communication for less memory.
  • Sharding granularity: FSDP can wrap individual layers (“auto-wrap policy”) so all-gathers cover only one layer’s parameters at a time, capping peak unsharded memory.
  • Mixing with tensor parallelism: FSDP shards across the data-parallel dimension. Very large runs often combine it with tensor parallelism on the fastest links. The exact placement follows the cluster topology.

When to use what

ConstraintStrategy
Fits on 1 GPUDDP (no sharding)
Persistent optimizer state is the first limitZeRO-1
Optimizer state and gradients are the limitZeRO-2
Stored parameters are also the limitFSDP / ZeRO-3
One gathered layer is too largeAdd tensor parallelism
Depth or topology needs another splitConsider pipeline parallelism

Common pitfalls

  • FSDP alone does not shard activation dimensions. Each GPU still holds activations for its data-parallel slice. Use activation checkpointing or sequence/context parallelism when those activations dominate.
  • All-gather overhead at small layer size. Wrapping every linear layer separately can dominate runtime; wrap at transformer-block granularity instead.
  • Confusing sharding with tensor parallelism. Sharding (FSDP) splits state across data-parallel ranks and reconstructs it for compute. Tensor parallelism splits the compute of a single layer; the math is different.

Found an error or stale claim? Report a correction or send a private note.