ZeRO and fully sharded data parallelism
Draw one module's parameter and gradient lifetimes through fully sharded forward and backward execution and identify the measured memory peak
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Previous section's data-parallel state and gradient model
- Chapter 13 all-gather and reduce-scatter contracts
- Chapter 3 training-memory categories
A 12-billion-parameter transformer uses more device memory for parameters, gradients, and Adam optimizer state than one 80 GiB GPU can provide. Eight ordinary data-parallel ranks still fail because each rank stores the same complete parameter-related state. The capacity problem concerns replicated bytes, not a shortage of total bytes across the server.
PyTorch Fully Sharded Data Parallel and ZeRO-style methods divide parameter-related tensors into rank-owned ranges called shards. Rank 3, for example, can retain only its range of a parameter group between uses. Before one transformer block runs, the ranks execute an all-gather that assembles the needed full parameter group on every rank. After use, software can discard non-owned ranges. Backward later produces local gradient contributions, and reduce-scatter combines the contributions while returning only each owner's shard.
Collect two observations for one steady-state step. A device-memory timeline should show low persistent parameter storage interrupted by peaks when full parameter groups and communication buffers coexist. A profiler trace should show all-gather before the relevant forward or backward computation and reduce-scatter after gradient production. PyTorch documentation also notes that a pre-forward host gap can be an intentional rate limiter that prevents too many outstanding all-gathers from exhausting device memory.
Fully sharded data parallelism therefore trades persistent replication for scheduled communication and temporary unsharded storage. Prefetching can hide some transfer but extends tensor lifetimes. Small wrapping units free memory early but issue many small collectives. Large units transfer efficiently but create larger peaks. The executed memory timeline and communication readiness decide which trade works.
Separate the sharding stages
The ZeRO analysis separates parameter-related state into optimizer state, gradients, and parameters. Sharding optimizer state removes the largest replicated category for Adam-like configurations while leaving parameters and gradients replicated. Sharding gradients assigns final gradient ownership. Full sharding also partitions parameters between uses.
The ZeRO stages change persistent storage differently. The stages do not automatically divide activations, operator workspaces, or every communication buffer. Use the actual framework strategy and data formats to calculate each category.
Follow one FSDP2 module through forward
PyTorch FSDP2 represents persistent parameters as sharded DTensors. Before a wrapped module executes, a pre-forward hook all-gathers its parameter group so every shard rank obtains the unsharded values required by local computation. The module then runs with ordinary unsharded parameter tensors.
After forward, the full parameters can be freed and the module returns to its sharded representation. If parameters reshard after forward, backward must all-gather them again. Keeping them unsharded removes that second gather but extends their full-size lifetime through the intervening forward and backward work.
Follow the module through backward
Before the module's backward computation, FSDP ensures that full parameters are available. Backward produces unsharded local gradient contributions. A reduce-scatter combines contributions across the shard group and leaves each rank with the gradient shard for its owned parameter shard.
The optimizer updates local shards using local optimizer state. This preserves sharded ownership between steps. Pending reductions, accumulation state, and partial iterations need explicit invalidation if a worker fails.
Calculate persistent and transient memory
For S shard ranks, ideal persistent parameter-related storage approaches one/S of the sharded categories. During execution, one or more complete parameter groups can coexist with activations, gradients, workspaces, all-gather outputs, and reduce-scatter buffers.
Draw lifetimes. A current full group can remain live while the next group is prefetched. Backward can overlap a next all-gather with a previous reduce-scatter. Better communication overlap can therefore raise peak memory. The high-water mark, not the persistent average, determines whether the step fits.
Choose module group granularity
FSDP2 grouping follows how fully_shard is applied. A group all-gathers its parameters in one collective and reduce-scatters its gradients in one collective. Bottom-up layer grouping provides parameter release boundaries and opportunities to overlap the next gather with current compute.
Very fine groups create small collectives, repeated alpha cost, and more hooks. Very coarse groups gather large parameter sets, delay release, and increase peak transient memory. Select groups from parameter bytes, module compute time, executed order, and available overlap.
Schedule prefetch explicitly
A forward prefetch issues the next module's all-gather before that module's pre-hook would naturally run. Earlier submission can cover CPU launch delay and move transfer beneath current compute. Backward prefetch uses known reverse execution order to prepare an upcoming module.
Prefetch must respect the memory budget. Prefetching two future groups keeps more unsharded parameters live. Also distinguish CUDA streams from network concurrency: separate all-gather and reduce-scatter streams do not guarantee simultaneous wire progress when they use one serialized communicator.
Compose with checkpointing and hybrid sharding
Activation checkpointing changes which forward operations rerun during backward and can change the window available for parameter all-gather. Model both schedules together. A configuration that fits after activation savings can still expose more parameter communication during recomputation.
Hybrid sharded data parallelism shards within one group and replicates across another mesh axis. Hybrid sharding can keep parameter all-gather and gradient reduce-scatter inside a node while synchronizing replicated gradient shards across an inter-node replica group. The outer synchronization is not automatically less frequent; its operation, payload, and ordering depend on the implementation. Persistent storage increases relative to full global sharding, while cross-node message shape and participant count change. Calculate both levels because replica synchronization still lies on the optimizer dependency.
GiB/rank = parameters × bytes/parameter ÷ shardsThis is a floor. Add transient lifetimes before deciding whether the job fits.
Worked example: Explaining an FSDP memory spike
An eight-rank FSDP2 job has 14 GiB of persistent sharded state per rank and a 24 GiB device-memory budget for model state and activations. The job still runs out of memory near the transition between two large transformer blocks.
- Capture a memory timeline with module and collective annotations. At the failure point, the persistent 14 GiB coexists with 2.8 GiB of current full parameters, 3.1 GiB of next-block prefetched parameters, 2.2 GiB of activations, and a 1.4 GiB attention workspace.
- Add communication scratch and allocator reservation. The simultaneous total exceeds 24 GiB even though persistent state alone appears safe.
- Confirm release points. One reference from a debugging hook retains the current full parameter group beyond its module. Remove the reference and verify that post-forward reshard releases it.
- Limit forward prefetch to one next group and split the 3.1 GiB block into two groups whose compute time can cover their gathers. Recalculate startup count and transient overlap.
- Measure the revised peak and step time. Lower prefetch depth can expose communication, while finer grouping adds alpha. Select the configuration that fits with operating headroom and has the lowest complete step time.
- Test backward with activation checkpointing enabled because recomputation changes module order and buffer lifetimes. Verify loss, optimizer shards, and checkpoint save/restore after regrouping.
Conclusion: Sharding reduces ownership, not every instantaneous allocation. Peak memory is governed by the schedule of gathered lifetimes.
Common errors
- Dividing complete training memory by shard count. Activations, workspaces, allocator state, and transient gathers follow different scaling.
- Assuming a parameter is sharded while its module executes. Full values are reconstructed for the local computation unless the module itself uses another parallel layout.
- Maximizing prefetch depth. Earlier transfer extends full-parameter lifetimes and can cause an out-of-memory failure.
- Using extremely small module groups. Repeated startup and hooks can dominate communication while local computation is too short to hide it.
Reference summary
ZeRO-style methods remove replicated parameter-related state in stages: optimizer state, then gradients, then parameters. A fully sharded implementation stores only local parameter shards between uses. The runtime all-gathers a module's parameters before computation, frees or reshards the full parameters after use, and reduce-scatters gradients back to their owning ranks.
- Persistent parameter-related storage can approach one divided by the shard-group size, but activations, workspaces, and transient unsharded parameters follow separate rules.
- Resharding after forward saves memory but requires another parameter all-gather before backward. Keeping full parameters avoids that communication at a higher memory cost.
- The selected module groups define all-gather and reduce-scatter message sizes. Very small groups pay repeated startup; very large groups increase transient memory and delay readiness.
- Prefetching an upcoming all-gather can hide transfer time. Early or deep prefetching keeps more full parameter groups live at once.
- Hybrid sharding partitions within a local group and replicates between groups. This limits frequent sharding traffic to a faster domain at the cost of replication.
Knowledge check
Why can wrapping every tiny submodule independently make fully sharded training slow?
Show the answer guide
- Every independently wrapped tiny module can issue its own parameter all-gather and gradient reduce-scatter, multiplying fitted startup costs and host hooks.
- Fine groups can also provide too little module compute to hide each communication operation, while prefetching many groups increases simultaneous full-parameter lifetimes.
- Group selection must compare parameter bytes, compute time, readiness, peak transient memory, startup count, and exposed communication rather than minimizing persistent storage alone.
Measure FSDP group granularity
Compare at least three module-group sizes for one FSDP model or trace-based model. Keep shard degree, global batch, and reshard policy fixed while measuring one steady-state step.
Evidence to keep: Group membership and bytes, all-gather and reduce-scatter counts, readiness and duration traces, peak memory timeline, complete step times, and correctness or loss equivalence.
FSDP shards model state across data replicas. Tensor and sequence parallelism instead partition the computation inside individual layers.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.