FSDP2's parameter lifecycle, and the wrap-policy bug that turns one AllGather into ten

DDP’s entire discipline is getting different data to different ranks while every rank holds an identical, full copy of the model and optimizer state. That copy has an exact, computable cost: a 7B-parameter model under Adam needs 16 bytes of persistent state per parameter, 112GB, on every single GPU, whether that GPU has 80GB or 800GB. FSDP and ZeRO are the direct answer, not a fancier default: shard that 112GB across $N$ ranks instead of replicating it whole, as already laid out stage by stage, ZeRO-1 shards the optimizer moments, ZeRO-2 adds gradients, ZeRO-3, which FSDP implements natively in PyTorch, shards the parameters themselves too. What that table doesn’t cover, and what actually determines whether an FSDP job hits its expected memory and throughput or mysteriously doesn’t, is everything that happens in the seconds around each layer’s forward and backward pass, and that’s worth a full pass on its own.

16 bytes per parameter is where planning starts, not where a memory ledger ends#

The $16/N$ bytes-per-parameter number is correct as a persistent accounting, what’s permanently resident once training has stabilized. It says nothing about the transient memory a training step actually touches at its peak, and peak, not steady-state persistent memory, is what causes an OOM. A real ledger separates the two:

Persistent (every step, steady state):
  BF16 parameter shard              2/N bytes/param
  BF16 gradient shard               2/N bytes/param
  FP32 optimizer moment m shard     4/N bytes/param
  FP32 optimizer moment v shard     4/N bytes/param
  FP32 master-weight shard          4/N bytes/param
                                    = 16/N bytes/param, matches the ZeRO-3 row exactly

Transient (materializes and releases within one step):
  Full, gathered parameters for whichever module is currently executing
  Reduce-scatter output buffer
  Saved activations for backward
  Communication staging buffers

The transient row is the one a flat 16-bytes constant hides completely. Whichever transformer block is currently in its forward pass has its full, ungathered parameters sitting in memory for the duration of that block’s compute, not its $1/N$ shard, and if that materialization overlaps with saved activations from several other in-flight blocks (prefetch, gradient checkpointing interactions, a reshard policy that retains rather than frees), peak memory is set by how much transient state is alive simultaneously, not by the sum of persistent shards. Predicting FSDP memory from the 16-byte constant alone is a planning heuristic that gets the steady-state number right and the actual OOM boundary wrong; the only reliable way to know peak memory for a specific model, wrap policy, and reshard setting is to build the ledger above for that specific configuration and measure the transient row directly, e.g. with a CUDA memory snapshot, rather than assume the constant already accounts for it.

The parameter lifecycle, one module at a time#

Parameter shards at rest (each rank holds 1/N)
        │
        ▼
AllGather before this module's forward compute
        │
        ▼
Full module parameters temporarily materialized, on every rank
        │
        ▼
Forward computation runs against the full, gathered weights
        │
        ▼
Reshard now (free the gathered copy), or retain (skip a second AllGather later)
        │
        ▼
AllGather again for backward, if resharded
        │
        ▼
Backward computation
        │
        ▼
ReduceScatter: gradients reduced, each rank keeps only its own shard
        │
        ▼
Back to sharded parameters, sharded gradients, ready for the optimizer's shard-local step

AllGather and ReduceScatter are FSDP’s matched pair, the exact same sharding idea run in each direction: AllGather reconstructs what compute needs, ReduceScatter re-shards what the optimizer needs. Every module a job trains runs this full cycle, which means the number of modules chosen as the AllGather/ReduceScatter boundary is not an implementation detail, it’s the single biggest lever on how much communication the job issues.

Wrap granularity is the lever, and this is the incident it was already worth returning to#

One of the four NCCL incidents covered earlier in this series was exactly this failure, described from the outside: an FSDP job running steadily since the very start at roughly a third of expected MFU, no hang, no crash, nothing that announces itself as a communication bug. The root cause named there, auto_wrap_policy set to wrap at the individual-module level instead of the TransformerBlock level, is worth working through properly here rather than just naming it, because it’s the cleanest illustration of what wrap granularity actually controls.

Wrap too fine, one FSDP unit per linear layer inside a transformer block, and the lifecycle above runs once per linear layer:

AllGather Q-proj  →  compute  →  reshard
AllGather K-proj  →  compute  →  reshard
AllGather V-proj  →  compute  →  reshard
AllGather out-proj → compute  →  reshard
AllGather FFN-gate → compute  →  reshard
...

Six or eight small, latency-sensitive collectives per block instead of one, and small collectives pay a fixed per-call overhead that a single larger collective moving the same total bytes doesn’t. That’s the direct mechanism behind “10x AllGather calls for architecturally identical work” from the earlier incident: nothing about the model changed, only how many times the same lifecycle ran per block.

Wrap too coarse, one FSDP unit for half the model, and the opposite failure appears: a single AllGather now has to materialize an enormous amount of full-precision parameter memory simultaneously, which can push the transient row of the ledger above past what’s actually available, and there’s less opportunity to overlap that one enormous gather with useful compute happening elsewhere. The correct boundary, in practice, is the natural architectural repeating unit, one TransformerBlock, matching FSDP’s communication granularity to the model’s own structure rather than to Python’s module tree. The fix in the original incident, correcting auto_wrap_policy to operate at TransformerBlock granularity, is the general answer, not a special case: measure the actual AllGather count issued per step, compare it against the number of repeating architectural units in the model, and treat a mismatch between the two as the thing to fix before looking anywhere else.

What the fix actually looks like in code. FSDP2’s fully_shard wraps whatever module you hand it, and the wrap boundary is nothing more than which module that is:

from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)

# The bug: wrapping every linear inside every block, one FSDP unit per matmul
for block in model.layers:
    for linear in (block.q_proj, block.k_proj, block.v_proj, block.o_proj, block.gate_proj):
        fully_shard(linear, mp_policy=mp_policy)      # six AllGathers per block, not one

# The fix: wrapping at the repeating architectural unit instead
for block in model.layers:
    fully_shard(block, mp_policy=mp_policy)            # one AllGather per block
fully_shard(model, mp_policy=mp_policy)                 # root wrap covers embeddings, final norm, head

Nothing here is a special API for “correct” wrapping, it’s the same fully_shard call, pointed at a different level of the module tree, which is exactly why this bug is so easy to introduce silently: both versions run without error, both train a numerically correct model, and only the AllGather count issued per step tells them apart.

Reshard policy: pay for the memory, or pay for the communication again#

Resharding immediately after forward frees the transient row fast, minimizing peak memory, at the cost of a second AllGather to reconstruct the same parameters for backward. Retaining the gathered parameters through to backward avoids that second AllGather, trading memory for communication. Neither is universally correct: a large block on a memory-constrained GPU wants immediate resharding; a small block with backward following forward almost immediately, and headroom to spare, wants retention. The right choice is read off the ledger above for the specific model and hardware, not assumed from a framework default.

fully_shard(block, mp_policy=mp_policy, reshard_after_forward=True)   # reshard immediately: less peak memory, one extra AllGather
fully_shard(block, mp_policy=mp_policy, reshard_after_forward=False)  # retain through backward: more peak memory, one fewer AllGather

reshard_after_forward also accepts an integer rather than just a boolean, resharding down to a smaller world size instead of all the way to 1/N, intra-node only, say, so the eventual backward AllGather only has to happen over that smaller group instead of the full mesh. That’s the same instinct hybrid sharding generalizes below: resharding fully isn’t the only alternative to not resharding at all, there’s a genuine middle setting between the two extremes.

Prefetch: the difference between a trace with gaps and one without#

Gathering module $i{+}1$‘s parameters while module $i$ is still computing is what keeps AllGathers from becoming pure overhead sitting on the critical path:

Well-scheduled:
Compute block i:      ███████████
Gather block i+1:       █████
Compute block i+1:                ███████████

Poorly-scheduled:
Gather:               █████
Compute:                   ███████████
Gather:                                █████
Compute:                                   ███████████

The first trace has communication fully hidden inside compute that was going to happen regardless. The second exposes every gather as pure added wall-clock time. This is the same overlap principle already covered for gradient AllReduce, worth up to roughly 30% MFU when done correctly, applied to FSDP’s AllGather instead of DDP’s AllReduce, and it’s a second, independent reason wrap granularity matters: too many tiny units gives the prefetcher too little useful compute to hide each gather behind.

Hybrid sharding: shard where bandwidth is cheap, replicate where it isn’t#

A single flat FSDP mesh across every GPU in a job means every AllGather potentially crosses node boundaries, paying inter-node bandwidth for what intra-node NVLink could have done far faster. Hybrid Sharded Data Parallel splits this into two dimensions instead: shard within a node, where bandwidth is abundant, and replicate that sharded copy across nodes, where it isn’t. This is the same underlying instinct tensor parallelism is built around, keep the traffic that has to happen constantly on the fastest physical link available, and accept a coarser-grained, less frequent form of communication (replication instead of a cross-node gather) on the slower one.

FSDP versus DeepSpeed is a stack decision, not a correctness one#

Both implement the same ZeRO-3 idea; the decision between them is almost never “which one shards correctly,” it’s which ecosystem the rest of the job already lives in. FSDP is the PyTorch-native path: DTensor-backed sharded state, direct composition with PyTorch’s own tensor- and pipeline-parallel APIs through a shared DeviceMesh. DeepSpeed brings its own integrated engine, and is the more common choice specifically when CPU or NVMe offload is a real requirement, spilling optimizer state past GPU memory entirely, rather than just sharding it across GPUs. It’s telling that real infrastructure postings ask for comfort with both rather than a single choice: Mistral’s research-engineering requirements list “DeepSpeed / FSDP / SLURM / K8s (distributed training frameworks)” together, as one combined line item, not an either-or. The expectation in a real infra role isn’t picking a favorite, it’s being able to read and reason about whichever of the two a given team already standardized on.

Common mistakes#

Treating $16/N$ bytes per parameter as the full memory picture: it’s the persistent row only, and peak memory, the number that actually determines whether a job OOMs, is set by the transient row, gathered full-precision modules, activations, and communication buffers alive at once.

Choosing wrap granularity by what’s convenient in the module tree rather than by measuring: the fix that resolved the earlier incident was aligning the wrap boundary to the model’s actual repeating architectural unit, found by counting AllGathers against expected layer count, not by guessing.

Assuming a reshard policy is a fixed best practice rather than a memory/communication trade to make per model: immediate resharding and retention are both correct in different regimes, and the regime is read off the ledger, not off a framework default.

Try it yourself#

Beginner. For a 7B-parameter model at $N=8$ with the standard BF16-compute / FP32-master-weight recipe, compute both the persistent-only memory per GPU and give one concrete example of what would need to be added to reach a realistic peak (name at least two transient contributors).

Intermediate. Given a TransformerBlock containing 6 linear sublayers, compute how many AllGather calls per block a per-linear wrap policy issues versus a per-block wrap policy, and state the ratio. Then explain why that ratio, not the absolute count, is what determines whether MFU looks like the healthy or the degraded case from the earlier incident.

Advanced. Design the CUDA memory snapshot experiment that would let you measure, for a real model, whether peak memory is dominated by the persistent or the transient row of the ledger above, and what a “retain after forward” versus “reshard immediately” run would each look like in that snapshot.

---

The one-sentence version: the 16-bytes-per-parameter number correctly plans steady-state persistent memory but says nothing about the transient row that actually causes an OOM, FSDP2 runs an AllGather-materialize-reshard-ReduceScatter cycle per wrapped module and the size of that module is the single biggest communication lever in the job, and the wrap-granularity bug flagged only as a symptom earlier in this series turns out to be exactly this lifecycle running too many or too few times for the model’s real architecture. Sharding solves the memory-replication problem DDP has; tensor and pipeline parallelism exist for the case where a single layer, or the model’s full depth, doesn’t fit no matter how the optimizer state is sharded.