04 — ZeRO and FSDP

Papers: ZeRO (Rajbhandari et al., Microsoft DeepSpeed, 2019 — arxiv 1910.02054), PyTorch FSDP paper (Zhao, Gu et al., Meta, 2023 — arxiv 2304.11277), plus the FSDP2 design docs on the PyTorch blog (2024).

The framing

Data parallelism replicates the whole model on every rank. Beautiful for compute utilization, terrible for memory. A 70B in bf16 is 140 GB of weights alone; add AdamW’s fp32 master weights + first/second moments and you’re at 140 + 280 + 280 + 280 ≈ 1000 GB of state. Doesn’t fit on any single GPU.

ZeRO’s insight: you don’t need every rank to hold every piece of state at every moment. Shard the state, communicate it back into shape just in time for the compute that needs it, then throw the shards back.

ZeRO stages

Stage

What’s sharded

Memory saving

Extra comm

ZeRO-1

Optimizer states

4× typical (AdamW fp32 states are huge)

None extra (grad AR still happens)

ZeRO-2

+ Gradients

8× typical

Same as DP (reduce-scatter instead of all-reduce)

ZeRO-3

+ Parameters

Nx (n = DP world size)

+1 all-gather per layer per fwd/bwd

ZeRO-1 is nearly free — the optimizer step happens once per iteration, so the state gather is amortized.

ZeRO-2 replaces the grad all-reduce with reduce-scatter (same bytes moved, but each rank only receives its shard). Also nearly free.

ZeRO-3 is where the real memory saving is and where the communication cost becomes non-trivial: before each layer’s forward, the layer’s weights must be all-gathered from every rank; used; released. Same in backward. This is ~2 × params_per_layer per layer per direction. Overlapping this with compute is the whole implementation challenge.

FSDP — PyTorch’s implementation of ZeRO-3

PyTorch’s FSDP (Fully Sharded Data Parallel) is the productionized version of ZeRO-3 for the PyTorch ecosystem. The 2023 paper (Zhao et al.) is worth reading for the co-design story:

  • Deferred initialization — you can meta-init a model too big to fit anywhere, then FSDP allocates only the local shard.

  • Flat parameter grouping — instead of one collective per parameter, group into “flat params” per unit (usually per transformer block). This coalesces communication.

  • Mixed precision inside FSDP — bf16 compute, fp32 gradient reduction is native.

  • CPU offload — optimizer states or parameters can live in CPU RAM and stream in.

FSDP1 vs FSDP2 — critical distinction

FSDP1 (2022-2023) had painful edges: flat parameters made per-parameter LR groups impossible, mixed-precision recipes were fragile, activation checkpointing composed weirdly, and you couldn’t easily inspect a shard.

FSDP2 (PyTorch 2.4+, 2024) rebuilt on DTensor (distributed tensors) instead of flat params. Each parameter is a first-class DTensor with a per-parameter shard. Concrete wins:

  • Per-parameter shard means clean introspection, clean LR groups, clean state dict.

  • Cleaner CPU offload semantics.

  • Better composition with TP (torch.distributed.tensor.parallel composes with FSDP2 via device_mesh).

  • This is the current default for open-source training at 8B–70B scale.

The API is roughly:

from torch.distributed.fsdp import fully_shard, FSDPModule
from torch.distributed._composable.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy

mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",))
model = build_model()

# Apply per transformer block for finer sharding
for block in model.blocks:
    fully_shard(block, mesh=mesh, mp_policy=MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    ))
fully_shard(model, mesh=mesh)

Compare to FSDP1’s FSDP(model, ...) monolithic wrap. FSDP2 wants per-block wrapping so the all-gather scope stays small.

Composition with TP — the 2D mesh

For 70B+ you often want FSDP × TP. Torch’s DeviceMesh expresses this:

mesh = init_device_mesh("cuda", (dp_size, tp_size), mesh_dim_names=("dp", "tp"))
# TP: shard weight tensors along tp dim
parallelize_module(block, mesh["tp"], plan={...})
# FSDP: shard remaining state along dp dim
fully_shard(block, mesh=mesh["dp"], ...)

The mesh abstraction is worth learning cold — it is the fabric of every modern PyTorch distributed setup.

When to use each

Scenario

Choose

Model + AdamW state fits, want DP speed

Plain DDP

Optimizer state doesn’t fit

ZeRO-1 (Deepspeed) or FSDP2 with optim_only policy

Optimizer + gradients don’t fit

ZeRO-2 / FSDP2 grad-shard

Model itself doesn’t fit at bf16

ZeRO-3 / FSDP2 full-shard

>100B, need cross-node scaling

FSDP2 within node + PP across nodes; or Megatron 3D (TP+PP+DP)

Extreme (400B+)

Megatron TP+PP+DP+EP is still the reference stack

The scaling behavior you’ll observe

  • FSDP2 at 8 GPUs, 1 node: almost free — NVLink hides the all-gathers.

  • FSDP2 at 16-32 GPUs, 2-4 nodes: IB starts to matter. Prefetch and communication overlap become critical. Expect 70-85% of ideal scaling.

  • FSDP2 at 100+ GPUs: you feel the limit. This is where Megatron’s TP+PP becomes worth its config complexity.

Rule of thumb: FSDP2 gives you 90% of Megatron at 10% of the config complexity, up to a few tens of GPUs.

Deepspeed vs FSDP2 in 2025

  • FSDP2 — PyTorch-native, the community default, tightly integrated with torch.compile, TP, PP (via torchtitan).

  • Deepspeed — heavier stack (Runtime + Ops + ZeRO). Still solid for its niche (ZeRO++ optimizations, MoE support via Deepspeed-MoE), and Hugging Face accelerate supports both. Mindshare has moved to FSDP2.

  • Megatron-LM — dominant at 100B+ scale. It doesn’t compete with FSDP; it composes above it.

Practical debugging

  • OOM during first backward: you forgot to fully_shard some block that’s holding onto a full-size activation. Wrap it.

  • Very slow first step: memory allocator warmup + first NCCL init. Ignore step-0 latency.

  • Communication overhead too high: check that torch.compile is on, that mixed precision is enabled (bf16 params halve every collective), that layer wrapping is per-block not per-model, and that CPU offload isn’t accidentally on.

  • Save/load headaches: use torch.distributed.checkpoint (DCP), not manual state_dict. DCP writes sharded checkpoints in parallel and can resharding on load — a lifesaver.

Read this order

  1. ZeRO paper (arxiv 1910.02054) — sections 3-5, the memory arithmetic.

  2. PyTorch FSDP paper (arxiv 2304.11277) — the co-design.

  3. FSDP2 PyTorch blog post (pytorch.org/blog “Introducing FSDP2”) — API changes.

  4. torchtitan repo — the modern reference for FSDP2 + TP + PP training runs.

Exercises

  1. Compute the memory footprint of a 13B model in bf16 with AdamW on 8 GPUs under: (a) DDP, (b) ZeRO-1, (c) ZeRO-2, (d) ZeRO-3. Show per-GPU numbers.

  2. Given a 70B model, 32 GPUs, 400 Gbps IB, 900 GB/s NVLink: sketch a 2D mesh (FSDP × TP) and justify TP degree.

  3. Take a nanoGPT-like training script and convert DDP(model) → FSDP2 with fully_shard on each block. Measure step time before and after.