ZeRO: what state partitioning proves, and what the experiments test

What the paper does

ZeRO, §§3–7 partitions three kinds of model state across data-parallel ranks instead of replicating them: optimizer state (stage 1), then gradients (stage 2), then parameters (stage 3). It changes who owns which tensor and when it is communicated, not the optimization objective.

The experiments in §10 evaluate ZeRO-100B: optimizer-plus-gradient partitioning with ZeRO-R, sometimes combined with Megatron model parallelism. Full parameter partitioning (stage 3) is analysed but not isolated experimentally. §§10.1–10.2 and Figure 2 report 400 V100 GPUs in 25 DGX-2 nodes with 800 Gbps between nodes, models up to 170B parameters, and speedups up to 10× against the stated baselines. §10.5 and Figures 6–8 compare bundles of optimizations from Table 3, and the appendix shows the parallelism and batch sizes differ between runs.

The arXiv history lists v1 on 4 October 2019 (withdrawn), v2 on 7 October and v3, the version used here, on 13 May 2020.

Ownership changes the denominator

Take a common mixed-precision layout: two-byte working weights, two-byte gradients, four-byte master weights and eight bytes of Adam moments, 16 bytes per parameter. Let $N$ be the number of trainable parameters and $d$ the number of ranks sharing the state. A state of $c$ bytes per parameter costs $cN$ per rank when replicated and $cN/d$ when evenly partitioned, so the four stages give

$$\begin{aligned} M_0&=16N, & M_1&=4N+12N/d,\\ M_2&=2N+14N/d, & M_3&=16N/d. \end{aligned}$$

Multiplying by $d$ gives the total across ranks: $16Nd$, $4Nd+12N$, $2Nd+14N$ and $16N$. Partitioning removes redundancy, but whatever stays replicated is still multiplied by the number of ranks.

The replicated terms also set a floor. Stage 2 can never bring per-rank state below $2N$, however large $d$ becomes. For seven billion parameters that floor is 14 GB, so a device budget below it is ruled out before any benchmark.

None of these formulas includes activations or the parameters gathered for computation. Stage 3 says who owns each shard; a layer still needs its full weights while it runs.

Check: does a state-memory formula settle whether a training step fits?

No. It counts only the allocations it includes. The peak also depends on which activations, gathered parameters, communication buffers and temporary tensors exist at the same moment.

Capacity and speed need separate explanations

More memory headroom can change how a job runs: room for a larger microbatch changes the work per kernel, and needing fewer model-parallel partitions changes the communication. Those are plausible reasons for a speedup, but they are not measurements of state partitioning itself.

So any comparison has to say which question it answers. The best throughput achievable on a fixed set of GPUs allows each method its own tuning. The effect of one technique requires everything else held fixed. The paper’s headline numbers are closer to the first.

Throughput also does not establish time to a target loss. If the effective batch changes, so do the number and character of the updates; a convergence claim needs the data, token budget, optimizer settings and evaluation fixed as well.

What to measure before adopting it

The paper does not give the peak memory, network contention, checkpoint behaviour or recovery cost of your configuration. Current implementations, such as DeepSpeed’s ZeRO stages and PyTorch FSDP2, have their own grouping, prefetching and residency controls, so check their documentation rather than the 2020 paper.

A controlled comparison holds model, sequence length, effective batch, precision and hardware fixed, then records peak allocated and reserved memory, step time and collective time for each policy, including the ones that fail. To isolate gathered-parameter memory, change only grouping or prefetching; to isolate activation memory, change only the microbatch or recomputation.

Memory is half the trade. §7 shows that stages 1 and 2 communicate the same volume per step as plain data parallelism, and stage 3 about 1.5 times as much.