Home / Fundamentals / Distributed Training

Data parallelism, collectives and ZeRO/FSDP

Distributed data parallel, ring all-reduce and its cost, overlapping communication with the backward pass, and how ZeRO stages and FSDP shard optimizer state, gradients and parameters.

14 min read

Data parallelism

Each of N workers holds a full model replica, processes a different slice of the global batch, and averages gradients before the optimizer step. Mathematically it's identical to training with the full batch on one device.

Gradient accumulation runs several micro-batches before each optimizer step, giving a larger effective batch without more memory.

Collectives

Collective Result Used for
All-reduce Every rank gets the sum DP gradients, tensor-parallel outputs
Reduce-scatter Each rank gets 1/N of the sum ZeRO/FSDP gradients
All-gather Every rank gets all shards ZeRO-3/FSDP parameters
All-to-all Each rank sends a different piece to each other rank MoE token routing

Ring all-reduce = reduce-scatter followed by all-gather, each N − 1 steps of S/N bytes. Each GPU sends and receives:

2 · (N − 1) / N · S   bytes

This is nearly independent of N, which is why all-reduce scales: bandwidth cost is about 2S regardless of the number of workers (latency grows with N, which is why tree and hierarchical algorithms are used at large scale).

Example: 7B parameters, bf16 gradients (S = 14 GB), 8 GPUs: each GPU moves 2 × 7/8 × 14 ≈ 24.5 GB. At ~450 GB/s per direction over NVLink that's ~55 ms; at ~50 GB/s per GPU across InfiniBand nodes, ~0.5 s.

Overlap. Gradients are grouped into buckets and all-reduced as soon as a bucket's backward finishes, overlapping communication with the rest of the backward pass.

ZeRO and FSDP

Plain data parallelism replicates all 16 bytes per parameter on every GPU. ZeRO removes the redundancy in stages:

Stage What's sharded across N_d ranks Bytes/param per GPU Communication vs DP
DDP Nothing 16 1×
ZeRO-1 Optimizer state (12 B) 4 + 12/N_d 1×
ZeRO-2 + gradients 2 + 14/N_d 1×
ZeRO-3 / FSDP + parameters 16/N_d ~1.5×

Example, 7B model on 64 GPUs: DDP 112 GB (doesn't fit), ZeRO-1 ~29 GB, ZeRO-2 ~15.5 GB, ZeRO-3 ~1.75 GB per GPU (plus activations and temporary buffers).

Scroll sideways to see the whole diagram.

One FSDP training step (per layer unit)[forward]; Rank i → Collectives: all-gather unit k params; Collectives → Unit k: full params; Rank i → Collectives: prefetch: all-gather unit k+1; Unit k: forward, then free non-owned params; Collectives → Unit k+1: full params; Unit k+1: forward, free; [backward (reverse order)]; Rank i → Collectives: all-gather unit k+1 again; Unit k+1: backward → full gradients; Unit k+1 → Collectives: reduce-scatter gradients; Collectives → Rank i: summed gradient shard; Note: unit k repeats: all-gather, backward, reduce-scatter; [step]; Rank i → Optimizer: update own shard (AdamW); Optimizer → Rank i: new bf16 shard.Rank iholds 1/N shardCollectivesNCCLUnit ktransformer blockUnit k+1Optimizershard of fp32 + AdamFORWARDall-gather unit k paramsfull paramsprefetch: all-gather unit k+1forward, then free non-owned paramsfull paramsforward, freeBACKWARD (REVERSE ORDER)all-gather unit k+1 againbackward → full gradientsreduce-scatter gradientssummed gradient shardunit k repeats: all-gather,backward, reduce-scatterSTEPupdate own shard (AdamW)new bf16 shard
One FSDP training step (per layer unit). Each rank stores only its shard of every unit. Parameters are all-gathered just before use and freed afterwards, in both forward and backward. Gradients are reduce-scattered so each rank keeps the summed gradient for its own shard, and each rank updates only that shard and its optimizer state.

How FSDP works: parameters are sharded by layer group ("unit"). Before a unit's forward, its parameters are all-gathered, used, and freed. In backward, they're all-gathered again, gradients are computed and reduce-scattered so each rank keeps only its gradient shard, and each rank updates its own shard of the parameters and optimizer state. Prefetching the next unit's all-gather overlaps communication with compute.

Hybrid sharding (HSDP) shards within a node (or small group) and replicates across groups, keeping the frequent all-gathers on fast links.

What to remember

  • Ring all-reduce costs 2(N−1)/N · S per GPU.
  • ZeRO-1/2 save memory at no extra communication; ZeRO-3/FSDP shards everything for ~1.5× communication.
  • Overlap and bucketing hide much of the communication.