Vol. XIV · Parallelism Plate 5 of 11ZeRO, FSDP, fully sharded data parallel, optimiser state sharding, reduce-scatter, all-gather, memory redundancyCE

Sharding the States.

Every data-parallel replica stores the same 16 bytes per parameter. Store each byte once, and fetch it when needed.

Plain data parallelism is wasteful: 64 replicas hold 64 identical copies of every weight, gradient and optimiser moment. Rajbhandari and colleagues’ ZeRO (2020) removes the redundancy in stages. Stage 1 gives each replica one n-th of the optimiser state and master weights; after the gradients are summed, each updates only its shard. Stage 2 also shards the gradients, replacing the all-reduce with a reduce-scatter so each replica receives only the summed shard it needs. Stage 3 shards the BF16 weights too, all-gathering each layer’s weights just before its forward and again before its backward pass, then discarding them.

Stages 1 and 2 cost no extra communication: a reduce-scatter plus an all-gather of the updated weights is exactly an all-reduce (plate XIV.3). Stage 3 adds one more all-gather, so it moves 1.5 times as many bytes — but brings memory per device down to 16N/n. PyTorch’s implementation of stage 3 is called fully sharded data parallel (FSDP). The workshop reproduces the ZeRO paper’s own example and applies it to Llama 3.

Below: left, model-state memory per device at each ZeRO stage for the chosen model and number of data-parallel devices, split into weights, gradients and optimiser state. Right, the bytes each device sends per step.

What you will learn
Workshop
Model states per device for N parameters on n data-parallel devices: 2 bytes of BF16 weights, 2 of BF16 gradients and 12 of FP32 master weights and Adam moments, each sharded or not according to the stage. Right: bytes each device sends per step, (n − 1)/n of 4N for an all-reduce (or reduce-scatter plus all-gather) and of 6N for stage 3’s extra all-gather. The 7.5 B, 64-device case is the ZeRO paper’s Figure 1.

Store each byte once

The ZeRO paper’s example is a 7.5-billion-parameter model on 64 devices. Without sharding each holds 120.0 GB; sharding the optimiser state brings it to 31.4 GB, adding the gradients to 16.6 GB, and adding the weights to 1.9 GB — the paper’s 120, 31.4, 16.6 and 1.9 GB. Stages 1 and 2 are free in communication; stage 3 sends 1.5 times as many bytes, because each layer’s weights are gathered twice, once for the forward pass and once for the backward.

For Llama 3 70B on 64 devices, ZeRO-1 still leaves 295.4 GB per device, ZeRO-2 156.5 GB, and only ZeRO-3 fits, at 17.6 GB — with activations still to come. The 405-billion-parameter model on 512 devices gets to 12.7 GB with full sharding. Which stage to use is a trade between memory and traffic: DeepSeek-V3 used ZeRO-1 inside its data-parallel dimension because pipeline and expert parallelism had already divided the model, while FSDP is the common choice when data parallelism is the main or only axis.

ZeRO also has offloading variants that move optimiser state to CPU memory or NVMe, trading PCIe bandwidth for capacity; they suit fine-tuning on few GPUs more than pre-training at scale.

Terms

ZeRO
Zero Redundancy Optimizer: sharding model states across data-parallel devices.
ZeRO-1, -2, -3
Sharding the optimiser state; then also gradients; then also weights.
FSDP
Fully sharded data parallel: PyTorch’s implementation of ZeRO stage 3.
Redundancy
Identical copies of the same states on every replica.
ZeRO stores each byte of model state once across the data-parallel group; stages 1–2 are free in communication, stage 3 costs half as much again for the largest saving.

Sources

  1. Rajbhandari, Rasley, Ruwase and He, “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models”, SC 2020.
  2. Zhao et al., “PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel”, VLDB 2023.
  3. DeepSeek-V3 Technical Report, arXiv:2412.19437, section 3.2 (ZeRO-1 data parallelism).