Skip to content
mlmentorship

FSDP and ZeRO: budget shards and temporary gathers

Derive persistent bytes at each sharding stage, trace one gathered unit, and test communication and activation-memory changes.

Published · 8 min read ·Role-specific ·Advanced
Reading depth

Summary

Sharding removes replicated stored state; computation can still require a temporary full parameter unit. Budget both the persistent shards and the buffers live during computation.

Task: place one billion parameters across four ranks. Assume BF16 parameters and gradients, FP32 Adam moments, and a separate FP32 master stored with optimizer shards. Count decimal GB, not GiB, and initially exclude activations and temporary buffers. How many persistent bytes remain on each rank?

ComponentBytes per parameter
Parameters (BF16)2
Gradients (BF16)2
Adam first moment (FP32)4
Adam second moment (FP32)4
Separate master weights (FP32, assumed here)4
Total16

Without separate masters, this convention totals twelve bytes, not sixteen. Framework precision and storage choices must be checked rather than inferred from the parameter dtype alone. The figure traces one fully sharded unit; box widths distinguish ownership from temporary reconstruction, not byte proportions.

Learning objective

Sharded storage still temporarily gathers the full current unit

Sharded resting storage still gathers the current unit temporarilyFour ranks each retain one shard at rest. The active wrapped unit gathers for computation, then returns to sharded storage.One rank’s shard grows temporarilyinto one active unitRest: retain one quarter1---Compute: gather this wrapped unit1234Afterward: retain the same quarter1---
Derivation
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ᵢ
narrow boxes show what rank i owns across the step. Wide boxes show the temporary full unit gathered for compute. This policy frees parameters after forward, so backward needs another gather. Reduce-scatter leaves Gᵢ; the local optimizer updates Pᵢ. Activations and communication buffers add to the peak. Original schematic checked against the ZeRO paper, PyTorch FSDP documentation, and DeepSpeed ZeRO documentation.
Check: Why does parameter sharding not mean a rank never temporarily holds the full current unit?
Read it this way: Four ranks each retain one shard at rest. The active wrapped unit gathers for computation, then returns to sharded storage.
Does sharded resting storage eliminate the active unit’s temporary gather?

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

Let be parameter count and the data-parallel group size. Optimizer storage here includes eight moment bytes and four master-weight bytes per parameter.

Stage 1: shard optimizer state

Parameters and gradients stay replicated; optimizer storage divides across ranks: bytes. For the fixture, GB, versus the replicated GB. Each rank updates its owned slice, then shares updated parameters as required by the implementation.

Stage 2: shard optimizer state + gradients

Only parameters stay replicated: GB here. Reduced gradient shards support local optimizer updates. The large- reduction tends to eightfold under this sixteen-byte convention; at four ranks it is only .

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

Ideal persistent storage becomes GB. All three formulas return sixteen GB at . FSDP and ZeRO-3 implement related full-sharding approaches, not identical APIs or peak-memory policies.

The lifecycle is gather, forward, free, gather again, backward, reduce-scatter, and local update. Prefetch and retained units can overlap those windows and increase live memory.

Trade a gather for retention

New task: reduce communication when memory has headroom. Suppose a wrapped unit contains million BF16 parameters: its full gathered buffer is GB. Assume the four-GB persistent shard remains allocated alongside two-GB activations and GB of other buffers. Conservatively counting a separate full gather gives GB.

Check: retaining this unit through backward avoids its second gather but keeps that buffer live. An ideal four-rank ring gather sends and receives GB per rank. Retaining four equal units instead raises this ledger to GB, exceeding an eight-GB device. Use when: measured communication warrants a policy whose complete live-buffer set fits.

Change the dominant memory

New task: activations grow to six GB while local batch and sequence length stay fixed across rank comparisons. At four ranks, the same ledger becomes GB. At eight ranks, persistent state halves, but GB still exceeds eight GB.

Check: a hypothetical checkpointing change reducing activations to two GB gives five GB at eight ranks, with recomputation cost. Use when: activations dominate; state sharding alone does not divide that term.

Failure and transfer

A giant wrapping unit can exceed capacity even when persistent shards are small. Reconstruct: why can smaller units be slower?

Answer

Smaller gathers reduce individual buffers but increase collective launches and can weaken efficient overlap. Choose from measured peak liveness and step time, not persistent bytes alone.

Tradeoffs

  • Wrapping determines gather size, launch count, and parameter lifetime; prefetch can make multiple full units overlap.
  • Gather retention trades traffic for live memory; allocator reuse and storage aliasing change the actual peak.
  • Shard counts can be uneven, especially with small tensors or padding.
  • Hybrid groups shard within a selected group and replicate across groups; use that sharding degree, not total GPU count.
  • Checkpoint formats and reshard-on-load policies must preserve parameter identity, optimizer state, and update count.

When to use what

ConstraintStrategy
Replicated state fits across the target data-parallel workersConsider DDP without state 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 largeRevisit wrapping or consider tensor parallelism
Depth or topology needs another splitConsider pipeline parallelism

Common pitfalls

FSDP does not inherently partition every activation dimension. Use activation checkpointing or an appropriate sequence/context layout for the actual dominant term. Tensor parallelism partitions a layer’s computation; it has different layout and collective contracts. The numerical ledgers above are hypothetical accounting examples, not measured allocator peaks or performance claims.

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