ZeRO Sharding

Data parallelism replicates the optimiser state N times for no reason. ZeRO shards it, then the gradients, then the parameters themselves.

[distributed][memory]

Theory

Mixed-precision Adam holds, per parameter: an fp16 weight and gradient (2 + 2 bytes) plus fp32 weight, momentum and variance (12 bytes) — sixteen bytes where the model itself needs two. Replicating that across $N$ ranks wastes $14(N-1)\Psi$ bytes.

\begin{aligned} \text{stage 1:} \quad M_1 &= 4\Psi + \tfrac{12\Psi}{N} \\[2pt] \text{stage 2:} \quad M_2 &= 2\Psi + \tfrac{14\Psi}{N} \\[2pt] \text{stage 3:} \quad M_3 &= \tfrac{16\Psi}{N} \end{aligned}

Stages 1 and 2 keep the communication volume of plain data parallelism: a reduce-scatter of gradients plus an all-gather of updated weights is exactly one all-reduce in disguise. Stage 3 additionally all-gathers each layer’s parameters immediately before use and frees them after, adding one further $\Psi$ of traffic per step.

Implementation

python · torch ≥ 2.4 · FSDP2

import torch
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,   # reduce grads in fp32
)

# Shard each block separately so parameters are gathered
# per-block and freed again, not held for the whole model.
for block in model.blocks:
    fully_shard(block, mp_policy=policy)
fully_shard(model, mp_policy=policy)

Reducing gradients in bf16 to save bandwidth is a false economy at scale: the reduction is a sum over $N$ ranks, and bf16 has eight mantissa bits. Keep reduce_dtype at fp32. Wrapping only the top-level model rather than each block is the other common mistake — it collapses stage 3 back into a single enormous all-gather.

gradient-checkpointing mixture-of-experts-routing

References

[1]Rajbhandari et al. — ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2019)arXiv:1910.02054

[2]Zhao et al. — PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023)arXiv:2304.11277