ZeRO and FSDP: Memory-Efficient Data Parallelism

"With N replicas, (N-1)/N of your VRAM is a redundant copy of someone else's optimizer state. ZeRO is the trick of shipping the bytes without storing them twice."

KV, Memory-Allergic AI Agent

59.2.1 The Memory Budget of Mixed-Precision Training

Before we can shard anything, we need to know what we are sharding. The per-rank memory footprint of a transformer training step decomposes into four bucket classes. Let $P$ be the parameter count, $N$ the data-parallel world size.

The 14 bytes per parameter beyond the BF16 params are almost entirely redundant under data parallelism. Every rank performs identical AdamW updates on identical inputs, so storing 1.12 TB on each of 16 GPUs is a 17.92 TB of cluster memory for a 1.12 TB working set. ZeRO erases that redundancy stage by stage.

59.2.2 The ZeRO Progression

Rajbhandari et al. (SC'20) introduced ZeRO (Zero Redundancy Optimizer) as a series of three orthogonal sharding stages, each picking off one of the buckets above. Each later stage subsumes the earlier ones and adds more communication for more memory savings.

59.2.2.1 ZeRO Stage 1: Shard Optimizer State

The AdamW master weights, first moment, and second moment are 12 bytes per parameter and used only inside optimizer.step(). Stage 1 partitions them across the $N$ data-parallel ranks: rank $i$ stores $1/N$ of every optimizer tensor.

Forward and backward run as in DDP, producing a full BF16 gradient on every rank. Before the optimizer step, an all-reduce averages gradients (same as DDP). Then each rank applies the optimizer to its shard of the master weights, producing a $1/N$ shard of updated BF16 parameters. A final all-gather reconstructs the full BF16 parameter tensor on every rank.

Memory per rank drops from $16$ B/param to $4 + 12/N$ B/param. At $N=64$, that is roughly $4.2$ B/param: a 70B model fits in $\approx 290$ GB of state per rank, or about $3.6$ H100s. Communication grows by one all-gather (size $S$, replacing the implicit broadcast inside DDP). For most workloads, Stage 1 is nearly free.

59.2.2.2 ZeRO Stage 2: Also Shard Gradients

The gradient is fully needed only on the rank that owns the corresponding optimizer-state shard. So replace the all-reduce with a reduce-scatter: each rank ends up holding only $1/N$ of the gradient (the part corresponding to its optimizer shard). Memory per rank drops to $2 + 14/N$ B/param.

The communication math is striking: a reduce-scatter moves the same total $S(N-1)/N$ bytes as the reduce-scatter half of an all-reduce. We have saved roughly half the gradient memory ($\approx 70$ GB for a 70B model at $N=64$) by replacing two collectives with one half-as-expensive collective.

59.2.2.3 ZeRO Stage 3: Also Shard Parameters

The big one. At any moment, only the parameters of the current layer need to materialize as full tensors; the rest can stay sharded. Stage 3 (also called FSDP Full-Shard) wraps each transformer block so that:

1. Just before forward: an all-gather reconstructs the block's full parameters on every rank.
2. During forward: the block computes normally; the gathered parameters are held in memory.
3. After forward: the gathered parameters are freed; only the $1/N$ shard persists.
4. Just before backward: another all-gather reconstructs the parameters (they were freed in step 3).
5. After backward: a reduce-scatter on the gradients, producing only the $1/N$ gradient shard on each rank.

Memory per rank is $16/N$ B/param. At $N=64$, that is $0.25$ B/param: a 70B model fits in $17.5$ GB of state per rank. You can run a 70B model on 8 GPUs (one DGX node), the configuration that put open-weight 70B models in reach of every research lab. The cost is three collectives per layer per step (forward gather, backward gather, gradient scatter), against one per step for DDP. The communication budget grows by roughly $3\times$ but is still bounded by interconnect bandwidth, not by $N$.

59.2.3 PyTorch FSDP: ZeRO Re-Implemented Natively

FSDP (Zhao et al., 2023) is the PyTorch-native reimplementation of ZeRO Stage 3, designed to feel like DDP from the user's perspective. FSDP differs from DeepSpeed ZeRO in two important ways: (a) it integrates with torch.compile, AOT autograd, and the rest of the PyTorch ecosystem natively; (b) it uses explicit wrapping policies to control sharding granularity, which gives finer control over the activation-vs-comm trade-off.

59.2.3.1 Wrapping Policies

FSDP shards parameters at the granularity of wrap units. A wrap unit is one all-gather: every parameter inside the unit is gathered together. There are three common policies:

• No wrap (single FSDP unit): the entire model is one unit. One all-gather brings in the full model; trivially correct but peak memory equals the unsharded model size. Useful only for tiny models.
• Auto wrap by size: traverse the module tree, wrap submodules whose parameter count exceeds a threshold (e.g., $10^7$). Simple but blunt.
• Transformer-block wrap (the default for LLMs): each transformer block (e.g., LlamaDecoderLayer) becomes its own FSDP unit. The peak materialized memory is one block's worth of params plus the gradient working set, which for a 70B model with 80 layers is about $1.75$ GB. This is what every modern LLM training recipe uses.

59.2.3.2 A Minimal FSDP Training Loop

# fsdp_minimal.py: torchrun --nproc_per_node=8 fsdp_minimal.py
import os, functools, torch
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    ShardingStrategy, MixedPrecision, BackwardPrefetch, CPUOffload,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers import AutoModelForCausalLM
from transformers.models.llama.modeling_llama import LlamaDecoderLayer

def main():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    # 1) Build the bare model on each rank (full materialization here; meta-device
    #    init is the production trick when even that is too big; see Section 59.5).
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-70B")

    # 2) Wrap each transformer block as its own FSDP unit.
    wrap_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={LlamaDecoderLayer},
    )

    # 3) Mixed precision: parameters in bf16, gradient reduction in bf16,
    #    keep loss / norm buffers in fp32 to avoid underflow.
    mp = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.float32,
    )

    # 4) FULL_SHARD = ZeRO Stage 3. SHARD_GRAD_OP is the Stage 2 equivalent.
    model = FSDP(
        model,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        auto_wrap_policy=wrap_policy,
        mixed_precision=mp,
        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,   # overlap gather w/ compute
        cpu_offload=CPUOffload(offload_params=False),      # set True if HBM tight
        device_id=local_rank,
    )

    # 5) The optimizer sees only the local shard, internally an fp32 master copy.
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-5)

    for batch in train_loader:
        out = model(**batch)
        out.loss.backward()
        # FSDP's hooks have already reduce-scattered grads by this point.
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

59.2.3.3 Prefetching and Overlap

The naive FSDP forward pass is: gather block $k$, compute block $k$, free block $k$, gather block $k+1$, compute block $k+1$, …. Each gather is a synchronous wait. With forward_prefetch=True, FSDP issues the gather for block $k+1$ while compute on block $k$ is still running; the same trick for BackwardPrefetch.BACKWARD_PRE on the way back down. With good overlap, the communication is essentially free up to about a $30\%$ slowdown over DDP for the largest models, which is the lever that decides whether FSDP is the right choice.

59.2.3.4 CPU Offload

If HBM is still tight even after sharding, FSDP supports CPUOffload(offload_params=True): the $1/N$ parameter shards live on host memory and are paged in to HBM only during a forward or backward gather. CPU memory is plentiful (1-2 TB per node is common); the cost is PCIe traffic on every layer's gather. PCIe Gen5 x16 is $63$ GB/s, far slower than NVLink, so CPU offload roughly halves training throughput. It is a "I cannot run otherwise" feature, not a default.

59.2.4 Memory Arithmetic: Llama 70B and 405B

To make the trade-offs concrete, here is the per-rank memory math for two scales on H100s with 80 GB. We assume mixed precision (BF16 + FP32 master), AdamW, transformer-block wrap, and gradient checkpointing.

The interpretation: with 16-way FSDP Full-Shard, the per-rank state of a 70B model is $1.12 \text{ TB} / 16 \approx 70$ GB, which fits an 80 GB H100 with breathing room for activations. With 64-way sharding, the per-rank state is $17.5$ GB and you have 60+ GB for activations, KV caches, or a larger batch.

59.2.5 ZeRO++ and Recent Improvements

The three classic ZeRO stages have been refined by half a decade of follow-up work. The most impactful additions:

59.2.5.1 Quantized Weights on the Wire

ZeRO++ (Wang et al., 2023) observed that the all-gather of BF16 parameters is bandwidth-bound. By block-quantizing the parameter shards to INT4 or INT8 before transmission and dequantizing on the receiver, you reduce the wire bytes by $4\times$ at minimal accuracy loss. Quantization is per-block (one scale per ~128 elements) to preserve dynamic range. The communication-overhead reduction translates directly to throughput.

59.2.5.2 Hierarchical Communication (Hybrid-Shard)

When the data-parallel group spans multiple nodes, a flat all-gather has to cross the slow inter-node fabric for every gather. Hybrid-Shard partitions the world into intra-node groups: each node fully shards within itself ($N=8$ for an H100 server), and across nodes runs plain DDP. Memory savings are $8\times$, all collectives stay on NVLink, and inter-node traffic is the one DDP all-reduce per step. This is what Meta uses for Llama-3.

59.2.5.3 Overlap and Pipelining Improvements

FSDP2 (PyTorch 2.3+, 2024) and DTensor-based sharding (2025) replaced the original wrapper-based implementation with a per-parameter sharding model that prefetches and overlaps more aggressively. Combined with torch.compile, FSDP2 closes most of the throughput gap with hand-tuned Megatron-DeepSpeed configurations.

59.2.6 FSDP vs Tensor Parallelism: When to Choose Which

FSDP and Megatron-style tensor parallelism (Section 59.3) both make a too-large model fit. They are not interchangeable; the right choice depends on cluster topology and model shape.

The rule of thumb in 2026: use tensor parallelism within a node (TP $\le$ 8), use FSDP across nodes. Tensor parallelism exploits NVLink's bandwidth advantage; FSDP exploits the fact that data-parallel scaling is bandwidth-bounded, not latency-bounded, so it tolerates InfiniBand. The 3D-parallel composition in Section 59.4 makes this concrete.

59.2.7 Activation Checkpointing and FSDP

FSDP shards parameters and gradients, but it does not shard activations. The forward pass produces $\mathcal{O}(B L d \cdot n_\text{layers})$ activation memory which is fully materialized on every rank. For a Llama-3 70B model at $B=4, L=8192$ this is roughly 160 GB per rank, well over an 80 GB H100's budget.

Activation (gradient) checkpointing recomputes activations during backward rather than storing them. Wrap the transformer block in torch.utils.checkpoint and only the block's input is stored; the intermediate activations are recomputed from that input during backward. Memory drops by roughly $4\times$ per checkpointed block (the savings depend on what the block computes internally); compute cost rises by about $33\%$ (the recomputed forward).

FSDP and activation checkpointing compose with two subtleties:

• Selective checkpointing. Recomputing every layer is wasteful; recomputing every layer except the most compute-intensive (FFN) is the typical pattern. PyTorch's checkpoint_sequential and Megatron's recompute_method=selective pick the right layers automatically.
• Order of operations with FSDP. The activation checkpointing wrapper must be inside the FSDP wrapping. If you wrap with FSDP first and then with checkpoint, the recompute phase re-triggers parameter all-gathers, which is what you absolutely want to avoid. The standard pattern is FSDP(checkpoint_wrapper(LlamaDecoderLayer)), not the reverse.

59.2.8 Meta-Device Initialization: Skipping the OOM at Startup

For 70B+ models, even constructing the model on a single rank to wrap it with FSDP overflows memory. AutoModelForCausalLM.from_pretrained("llama-70b") materializes 140 GB of BF16 parameters; one rank cannot hold that.

The standard workaround is the meta device: construct the model with device="meta", which allocates only tensor metadata (shape, dtype) but no actual storage. Then FSDP's wrap pass shards the meta tensors, materializes only the local shard on each rank's GPU, and finally loads parameters shard-by-shard from disk. Total per-rank materialization is $\approx P/N$ bytes, never the full $P$.

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from transformers import AutoModelForCausalLM

# Build the model graph on the meta device (no real allocation).
with torch.device("meta"):
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3-70B", torch_dtype=torch.bfloat16
    )

# FSDP's wrap pass materializes only this rank's shard on GPU.
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    auto_wrap_policy=wrap_policy,
    sync_module_states=True,    # rank 0 broadcasts param state to others
    param_init_fn=lambda mod: mod.to_empty(device="cuda"),
)

# Finally, stream the actual weights from disk to the local shards.
load_sharded_state_dict(model, checkpoint_path)

59.2.9 DeepSpeed-ZeRO vs PyTorch FSDP: 2026 Status

The two stacks have converged in capability but diverge in ecosystem fit:

In 2026, FSDP2 is the default in new training stacks (Megatron-Core uses it under the hood for the DP axis); DeepSpeed remains the path of least resistance when starting from Hugging Face Trainer recipes. Both reach 35-45% MFU on H100s at 70B scale; the gap between them is now within measurement noise.

pip install accelerate
# Interactive wizard writes ~/.cache/huggingface/accelerate/default_config.yaml
accelerate config
#   -> choose: multi-GPU, FSDP, FULL_SHARD, BACKWARD_PRE,
#      TRANSFORMER_BASED_WRAP (auto-wraps each transformer block),
#      bf16 mixed precision

# Same train.py, now FSDP across 8 GPUs on one node:
accelerate launch --num_processes=8 --use_fsdp \
    --fsdp_sharding_strategy=FULL_SHARD \
    --fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP \
    --mixed_precision=bf16 train.py