Data parallelism and ZeRO optimization: Worked examples

Everything below rests on the parent's master equations. Let us re-write them once, in plain words, so no symbol is used before it is named.

The scenario matrix

Every question this topic can throw at you falls into one of these cells. The worked examples below are each tagged with the cell(s) they cover.

#	Cell class	The tricky bit it tests
A	Baseline vs stage	plug $Φ, N$ into $M_{DP}, M_{1}, M_{2}, M_{3}$ and compare
B	Degenerate $N = 1$	one GPU — does ZeRO save anything? (sign/zero-input analogue)
C	Limit $N \to \infty$	what floor does each stage approach?
D	Tiny model	when is ZeRO pointless? (activations dominate)
E	Feasibility solve	given GPU memory, solve for the minimum $N$
F	Communication cost	$C_{3}$ vs $C_{DP}$ , the 3× trade-off
G	Efficiency word problem	overlap compute & comms → wall-clock %
H	Exam twist	mixed: choose the cheapest stage that fits a constraint

Figure — Data parallelism and ZeRO optimization

The figure above plots $M_{1}, M_{2}, M_{3}$ against $N$ for a fixed model. Notice cells B and C are just the left edge ( $N = 1$ ) and the right tail ( $N \to \infty$ ) of these curves. Keep it in view.

Worked example Example 1 — Cell A (baseline vs every stage)

A model has $Φ = 1 0^{9}$ parameters (1 billion). You have $N = 8$ GPUs. Give $M_{gpu}$ for standard DP and for ZeRO-1, -2, -3 in GB.

Forecast: DP costs $16$ GB. Guess: does ZeRO-3 get you near $2$ GB or near $8$ GB?

Standard DP: $M_{DP} = 16Φ = 16 \times 1 0^{9}$ bytes $= 16$ GB. Why this step? $1 0^{9}$ bytes $= 1$ GB exactly by our convention, so $16Φ$ bytes reads directly as $16$ GB. This is the "everyone stores everything" cost.
ZeRO-1: $M_{1} = 4Φ + \frac{12Φ}{8} = 4 + 1.5 = 5.5$ GB. Why this step? Only the $12Φ$ optimizer block is split; the $4Φ$ (params+grads) stays replicated.
ZeRO-2: $M_{2} = 2Φ + \frac{14Φ}{8} = 2 + 1.75 = 3.75$ GB. Why this step? Now gradients ( $2Φ$ ) are also sliced, leaving only params ( $2Φ$ ) replicated.
ZeRO-3: $M_{3} = \frac{16Φ}{8} = 2$ GB. Why this step? Everything is sliced by $N$ ; nothing is replicated.

Verify: Monotone decreasing $16 > 5.5 > 3.75 > 2$ — each stage strictly beats the previous. ✓ Units: bytes/ $1 0^{9}$ = GB. ✓

Worked example Example 3 — Cell C (limit

$N \to \infty$ ) As you add infinitely many GPUs, what memory floor does each stage approach? (Ignore activations.)

Forecast: Which stage can be driven to near zero model-state memory, and which cannot?

$N \to \infty lim M_{1} = 4Φ + N \to \infty lim \frac{12Φ}{N} = 4Φ + 0 = 4Φ$ . Why this step? $\frac{12Φ}{N} \to 0$ because numerator is fixed while $N$ grows without bound (this is the meaning of a limit: the value you get arbitrarily close to). The $4Φ$ term has no $N$ , so it survives.
$N \to \infty lim M_{2} = 2Φ + 0 = 2Φ$ .
$N \to \infty lim M_{3} = N \to \infty lim \frac{16Φ}{N} = 0$ . Why this step? Here every term carries $\frac{1}{N}$ , so the whole expression vanishes.

Verify: For $Φ = 1 0^{9}$ : floors are $4$ GB, $2$ GB, $0$ GB. Sanity: ZeRO-1 can never drop below the $4Φ$ of replicated params+grads; only ZeRO-3 has no replicated floor. This is exactly why ZeRO-3 (with Activation Checkpointing for the activations it ignores) is the tool for trillion-parameter models. ✓ The curves in figure s01 flatten to these dashed floors.

Worked example Example 4 — Cell D (tiny model, ZeRO is pointless)

A small model has $Φ = 5 \times 1 0^{7}$ (50 million params). Its activations need $M_{act} = 12$ GB (big batch, long sequences). Compare total per-GPU memory for DP vs ZeRO-3 on $N = 8$ . Is ZeRO worth it?

Forecast: Model state is tiny here. Guess: does ZeRO-3 shrink the total by a meaningful fraction?

Model state DP: $16Φ = 16 \times 5 \times 1 0^{7} = 8 \times 1 0^{8}$ bytes $= 0.8$ GB. Why this step? Convert first so both numbers are in GB.
Total DP $= 0.8 + 12 = 12.8$ GB. Why this step? Total memory = state + activations; activations are not partitioned by ZeRO.
Model state ZeRO-3: $\frac{16Φ}{8} = 0.1$ GB. Total $= 0.1 + 12 = 12.1$ GB.
Savings fraction $= \frac{12.8 - 12.1}{12.8} = \frac{0.7}{12.8} \approx 5.5%$ . Why this step? Percent tells us the practical benefit, not the raw byte count.

Verify: Only $\approx 5.5%$ saved, yet ZeRO-3 costs 3× communication. Not worth it. Lesson: ==ZeRO pays off only when model state dominates memory; when activations dominate, reach for Gradient Accumulation or Activation Checkpointing instead.== Units check: GB throughout. ✓

Worked example Example 5 — Cell E (feasibility solve — solve for minimum

$N$ ) GPT-3, $Φ = 175 \times 1 0^{9}$ . Each A100 has $80$ GB, and you must leave $20$ GB free for activations, so model state must fit in $60$ GB. Using ZeRO-3, what is the minimum number of GPUs $N$ ?

Forecast: Example 3's ZeRO-3 floor is $0$ , so some $N$ works. Guess: is it closer to 16, 32, or 64?

Set up the constraint: $M_{3} = \frac{16Φ}{N} \leq 60$ GB. Why this step? We want the memory at most the budget — an inequality, not an equation, because more GPUs than the minimum also fit.
Total state $16Φ = 16 \times 175 \times 1 0^{9} = 2.8 \times 1 0^{12}$ bytes $= 2800$ GB. Why this step? This is the quantity being divided; compute it once.
Solve $\frac{2800}{N} \leq 60 \Rightarrow N \geq \frac{2800}{60} = 46.67$ . Why this step? Divide both sides — flipping to solve for $N$ . Because $N$ must be a whole number of GPUs, round up.
$N_{m i n} = 47$ . (In practice you round to a power/multiple of the node size, e.g. $48$ or $64$ .)

Verify: At $N = 47$ : $M_{3} = 2800/47 = 59.6$ GB $\leq 60$ ✓. At $N = 46$ : $2800/46 = 60.9$ GB $> 60$ ✗ (fails). So $47$ is genuinely the boundary. This matches the parent's "43.75 GB at $N = 64$ " (indeed $2800/64 = 43.75$ , comfortably under $60$ ). ✓

Worked example Example 6 — Cell F (communication cost & the 3× rule)

For $Φ = 175 \times 1 0^{9}$ , $N = 64$ , compute the per-GPU communication volume $C$ for standard DP and for ZeRO-3, and confirm the ratio.

Forecast: The parent claims exactly 3×. Guess: will the raw byte numbers really be in a clean 1:3 ratio?

$C_{DP} = \frac{2Φ ( N - 1 )}{N} = \frac{2 \times 175 \times 1 0 ^{9} \times 63}{64}$ . Why this step? Standard DP does one AllReduce of the fp16 gradients ( $2Φ$ bytes), and AllReduce moves $\frac{N - 1}{N}$ of the payload per GPU.
$C_{DP} = \frac{350 \times 1 0 ^{9} \times 63}{64} = 3.445 \times 1 0^{11}$ bytes $\approx 0.34$ TB.
$C_{3} = \frac{6Φ ( N - 1 )}{N} = 3 \times C_{DP} = 1.033 \times 1 0^{12}$ bytes $\approx 1.03$ TB. Why this step? ZeRO-3 = two AllGathers of params (fwd+bwd, $2 \times 2Φ$ ) + one ReduceScatter of grads ( $2Φ$ ) = $6Φ$ worth, versus DP's $2Φ$ .
Ratio $= C_{3} / C_{DP} = 3$ exactly. Why this step? The $\frac{N - 1}{N}$ factor is identical in both, so it cancels, leaving $6/2 = 3$ .

Verify: $6/2 = 3$ independent of $N$ and $Φ$ — the trade-off is always "3× communication for $N \times$ memory". ✓ (Parent quotes $2.06$ TB using $C = 6Φ \cdot 2 \cdot (N - 1) / N$ ; that counts send and receive, double our one-directional $1.03$ TB — same physics, factor-of-2 bookkeeping.)

Worked example Example 7 — Cell G (efficiency word problem)

Your cluster: ZeRO-3, $C = 2.06$ TB/step (both directions, from Example 6's parent figure). Interconnect gives $300$ GB/s per GPU. Compute takes $10$ s/step. If comms and compute cannot overlap, what fraction of wall-clock is useful compute? Then: if a smarter engine overlaps 60% of comms behind compute, what is the new efficiency?

Forecast: No-overlap first. Guess: below or above 60% efficiency?

Comm time $t_{c} = \frac{2.06 \times 1 0 ^{12}}{300 \times 1 0 ^{9}} = \frac{2060}{300} = 6.87$ s. Why this step? time $=$ bytes ÷ bandwidth; units $\frac{bytes}{bytes/s} =$ s. ✓
No overlap: total $= 10 + 6.87 = 16.87$ s. Efficiency $= \frac{10}{16.87} = 59.3%$ . Why this step? Useful compute is the $10$ s numerator; everything else is overhead in the denominator. Matches the parent's $\approx 59%$ .
60% overlap: hidden comm $= 0.6 \times 6.87 = 4.12$ s runs free behind compute; only $t_{c} - 4.12 = 2.75$ s is exposed. Why this step? Overlapping means comm happens during the $10$ s of compute at no extra wall-clock, up to the amount that fits.
Total $= 10 + 2.75 = 12.75$ s. Efficiency $= \frac{10}{12.75} = 78.4%$ .

Verify: Efficiency rose from $59.3%$ to $78.4%$ — overlap helps, as it should ( $59.3 < 78.4 < 100$ ). This is exactly why frameworks prefetch the next layer's params during the current layer's compute; see Pipeline Parallelism and Mixed Precision Training for the same overlap idea elsewhere. Units all seconds. ✓

Worked example Example 8 — Cell H (exam twist: cheapest stage that fits)

$Φ = 1 0^{10}$ (10 B params), $N = 4$ GPUs, each $24$ GB, activations need $6$ GB (so model state budget $= 18$ GB). Which is the cheapest ZeRO stage (least communication) that fits? Stage cost order: DP < 1 = 2 < 3 (ZeRO-1/2 add negligible comm over DP; ZeRO-3 triples it).

Forecast: Guess before computing: will plain DP already fit, or must you climb to ZeRO-2 or -3?

$16Φ = 16 \times 1 0^{10} = 1.6 \times 1 0^{11}$ bytes $= 160$ GB total state. Why this step? The common quantity all stages divide from.
DP: $M_{DP} = 160$ GB $> 18$ . ✗ Fails.
ZeRO-1: $M_{1} = 4Φ + \frac{12Φ}{4} = 40 + 30 = 70$ GB $> 18$ . ✗
ZeRO-2: $M_{2} = 2Φ + \frac{14Φ}{4} = 20 + 35 = 55$ GB $> 18$ . ✗
ZeRO-3: $M_{3} = \frac{16Φ}{4} = 40$ GB $> 18$ . ✗ Even ZeRO-3 fails! Why this step? The exam twist: sometimes no stage fits at this $N$ . You must add GPUs or use Memory-Efficient Optimizers to shrink the $12Φ$ block.
Find the real minimum $N$ for ZeRO-3: $\frac{160}{N} \leq 18 \Rightarrow N \geq 8.9 \Rightarrow N_{m i n} = 9$ . Why this step? Same solve-for- $N$ move as Example 5; round up.

Verify: At $N = 9$ : $160/9 = 17.8$ GB $\leq 18$ ✓; at $N = 8$ : $160/8 = 20$ GB $> 18$ ✗. So the honest answer is "ZeRO-3, but only once you have $\geq 9$ GPUs — 4 is not enough for any stage." A trap that punishes anyone who assumes ZeRO-3 always fits. ✓

Recall Self-test (reveal after guessing)

At $N = 1$ , how much memory does ZeRO-3 save versus standard DP? ::: None — all schemes equal $16Φ$ ; ZeRO is pure partitioning and needs $N > 1$ to help. As $N \to \infty$ , what floor does ZeRO-1 approach and why? ::: $4Φ$ — params ( $2Φ$ ) and gradients ( $2Φ$ ) stay replicated in stage 1; only the $12Φ$ optimizer block is sliced away. ZeRO-3 buys $N \times$ memory savings for how many times the communication? ::: Exactly $3 \times$ (the $6Φ$ vs $2Φ$ ratio, independent of $N$ and $Φ$ ). When is ZeRO a bad choice? ::: When activations, not model state, dominate memory — then use Activation Checkpointing / Gradient Accumulation instead.

See also: Large Language Models Training for where these numbers get used at trillion scale, and AllReduce and Collective Communication for the mechanics of the $\frac{N - 1}{N}$ factor.