Data parallelism and ZeRO optimization: Exercises

The master memory formula everything below leans on:

Figure — stacked memory bars (Ex 2.1–2.3 made visible). The chart below stacks the three memory components (orange = params $2Φ$ , magenta = grads $2Φ$ , violet = optimizer states $12Φ$ ) for each stage at $Φ = 7$ B, $N = 8$ . Watch what shrinks left-to-right: in ZeRO-1 only the tall violet block collapses; in ZeRO-2 the magenta block also shrinks; in ZeRO-3 all three blocks are cut to $1/ N$ . The number above each bar is the per-GPU total you compute in the Level-2 exercises.

Figure — Data parallelism and ZeRO optimization

Level 1 — Recognition

Exercise 1.1

A model has $Φ$ parameters. In standard data parallelism with mixed precision + Adam, how many bytes per GPU does the optimizer state alone occupy, and what three sub-pieces make it up?

Recall Solution

What we're asked: just recall the breakdown of the $12Φ$ term. Optimizer state $= 12Φ$ bytes, made of:

Momentum (fp32): $4Φ$
Variance (fp32): $4Φ$
Master weights (fp32 copy): $4Φ$

Why fp32? The Adam update touches tiny numbers; fp16 would underflow them to zero. The parent note earned this — the master copy keeps precision the fp16 params lose.

Exercise 1.2

Match each collective to its role: (a) AllReduce, (b) ReduceScatter, (c) AllGather — against (i) "every GPU ends with the full averaged gradient", (ii) "every GPU ends with only its own slice of the averaged gradient", (iii) "reassemble the full parameter set from partitions."

Recall Solution

(a) AllReduce → (i) full averaged gradient everywhere. Used by standard DP.
(b) ReduceScatter → (ii) each GPU keeps only its slice. Used by ZeRO-2/3 to partition gradients.
(c) AllGather → (iii) rebuild full params. Used by ZeRO-1/3.

See AllReduce and Collective Communication for the wiring of each.

Level 2 — Application

Exercise 2.1

A model has $Φ = 7 \times 1 0^{9}$ parameters. Compute the standard-DP memory per GPU in GB (use $1 GB = 1 0^{9}$ bytes).

Recall Solution

What: plug into $M_{standard} = 16Φ$ . $16 \times 7 \times 1 0^{9} = 112 \times 1 0^{9} bytes = 112 GB .$ Why it matters: a 7B model already busts an 80 GB A100 under plain data parallelism — that is the whole reason ZeRO exists.

Exercise 2.2

Same $Φ = 7 \times 1 0^{9}$ , now with ZeRO-3 across $N = 8$ GPUs. Memory per GPU in GB?

Recall Solution

What: $M_{ZeRO-3} = \frac{16Φ}{N}$ . $\frac{16 \times 7 \times 1 0 ^{9}}{8} = \frac{112 \times 1 0 ^{9}}{8} = 14 \times 1 0^{9} bytes = 14 GB .$ From 112 GB → 14 GB. That is the full $N \times = 8 \times$ reduction, because ZeRO-3 partitions every component.

Exercise 2.3

For $Φ = 7 \times 1 0^{9}$ , $N = 8$ , compute ZeRO-1 and ZeRO-2 memory per GPU (GB).

Recall Solution

ZeRO-1: $4Φ + \frac{12Φ}{N} = 28 + \frac{84}{8} = 28 + 10.5 = 38.5$ GB. (Here $4Φ = 4 \times 7 = 28$ GB and $12Φ = 12 \times 7 = 84$ GB before splitting; $84/8 = 10.5$ .)

ZeRO-2: $2Φ + \frac{14Φ}{N} = 14 + \frac{98}{8} = 14 + 12.25 = 26.25$ GB. ( $2Φ = 14$ GB; $14Φ = 98$ GB before splitting; $98/8 = 12.25$ .)

Ladder of savings: 112 → 38.5 → 26.25 → 14 GB as we go standard → Z1 → Z2 → Z3. This is exactly the shrinking-bars figure at the top of the page.

Level 3 — Analysis

Exercise 3.1

Derive the ZeRO-1 memory-reduction factor $\frac{4 N}{N + 3}$ from the two memory formulas, then evaluate its limit as $N \to \infty$ . Interpret that limit.

Recall Solution

What: ratio of standard to ZeRO-1 memory. $Reduction = \frac{16Φ}{4Φ + \frac{12Φ}{N}} .$ Why divide? "How many times smaller" is exactly old÷new. Cancel $Φ$ (it's in every term): $= \frac{16}{4 + \frac{12}{N}} = \frac{16 N}{4 N + 12} = \frac{4 N}{N + 3} .$ Limit: as $N \to \infty$ , $\frac{4 N}{N + 3} \to 4$ . Interpretation: ZeRO-1 can never beat $4 \times$ no matter how many GPUs you add, because the params+grads ( $4Φ$ ) stay replicated forever. Only the $12Φ$ optimizer state ever leaves. That ceiling is why ZeRO-2 and ZeRO-3 exist.

Exercise 3.2

Verify that ZeRO-3 moves 3× the communication of standard DP. Then say in one sentence why three-times, structurally.

Recall Solution

Standard DP: one AllReduce of gradients. One payload ( $2Φ$ ) crosses the ring's $\frac{N - 1}{N}$ fraction: $C_{std} = \frac{N - 1}{N} \times 2Φ = \frac{2Φ ( N - 1 )}{N} .$ ZeRO-3: AllGather params in the forward pass ( $\frac{N - 1}{N} \cdot 2Φ$ ) + AllGather params again in the backward pass ( $\frac{N - 1}{N} \cdot 2Φ$ ) + ReduceScatter grads once ( $\frac{N - 1}{N} \cdot 2Φ$ ): $C_{Z3} = 2 \cdot \frac{N - 1}{N} (2Φ) + \frac{N - 1}{N} (2Φ) = \frac{6Φ ( N - 1 )}{N} .$ Ratio: $\frac{6Φ ( N - 1 ) / N}{2Φ ( N - 1 ) / N} = 3$ . ✓ Why 3×: we now move parameters twice (forward, backward) plus gradients once, versus a single gradient movement in plain DP.

Figure — the 3× gap. The plot below draws $C_{std}$ (orange) and $C_{Z3}$ (violet) as functions of $N$ , in units of $Φ$ bytes. Both curves rise and flatten toward horizontal ceilings — $2Φ$ for standard DP, $6Φ$ for ZeRO-3 — because $\frac{N - 1}{N} \to 1$ . The permanent vertical gap between the ceilings is the 3× communication penalty you just derived.

Exercise 3.3

For $Φ = 175 \times 1 0^{9}$ , $N = 64$ , confirm the parent's numbers: memory/GPU and communication volume/step (in TB, $1 TB = 1 0^{12}$ bytes).

Recall Solution

Memory: $\frac{16 \times 175 \times 1 0 ^{9}}{64} = \frac{2800 \times 1 0 ^{9}}{64} = 43.75 \times 1 0^{9}$ bytes $= 43.75$ GB. ✓ Communication: $C = \frac{6Φ ( N - 1 )}{N} = \frac{6 \times 175 \times 1 0 ^{9} \times 63}{64}$ . $= \frac{66150 \times 1 0 ^{9}}{64} = 1033.6 \times 1 0^{9} bytes \approx 1.03 TB per direction .$ Counting both send+receive on the link this rounds to the parent's $\approx 2.06$ TB of moved bytes.

Level 4 — Synthesis

Exercise 4.1

You have $Φ = 40 \times 1 0^{9}$ and eight 80 GB A100s ( $N = 8$ ). Activations need $20$ GB per GPU. Which ZeRO stage is the cheapest-communication one that still fits? Show each stage's model-state memory and add activations.

Recall Solution

Budget: $80 - 20 = 60$ GB for model state.

Standard: $16Φ = 640$ GB. ✗
ZeRO-1: $4Φ + \frac{12Φ}{8} = 160 + 60 = 220$ GB. ✗
ZeRO-2: $2Φ + \frac{14Φ}{8} = 80 + 70 = 150$ GB. ✗
ZeRO-3: $\frac{16Φ}{8} = 80$ GB. ✗ (80 > 60!)

Nothing fits on data parallelism alone. Correct synthesis answer: even ZeRO-3 needs 80 GB of state + 20 GB activations = 100 GB > 80 GB. You must combine with another technique — e.g. Activation Checkpointing to shrink the 20 GB, or offload optimizer state (Memory-Efficient Optimizers), or add Pipeline Parallelism to raise the effective $N$ . The lesson: ZeRO stages don't live alone.

Exercise 4.2

Same setup but you add Activation Checkpointing cutting activations from $20$ GB to $5$ GB, and you may raise $N$ . What is the smallest $N$ (power of 2) so ZeRO-3 fits in 80 GB?

Recall Solution

Need $\frac{16Φ}{N} + 5 \leq 80 \Rightarrow \frac{640}{N} \leq 75 \Rightarrow N \geq 8.53$ . Smallest power of 2: $N = 16$ . Check: $\frac{640}{16} + 5 = 40 + 5 = 45$ GB ≤ 80 GB. ✓ Synthesis point: checkpointing frees space, but the model state still forced us to double the GPU count. Two levers pulled together.

Level 5 — Mastery

Exercise 5.1

Prove the general reduction factor for any ZeRO memory formula of the form $M = a Φ + \frac{b Φ}{N}$ (with $a + b = 16$ ) is $\frac{16 N}{a N + b}$ , and find its $N \to \infty$ ceiling. Then read off the ceilings for Z1, Z2, Z3.

Recall Solution

Derivation: $Reduction = \frac{16Φ}{a Φ + \frac{b Φ}{N}} = \frac{16}{a + \frac{b}{N}} = \frac{16 N}{a N + b} .$ Limit: take $N \to \infty$ properly (never divide $16$ by $a = 0$ blindly). Divide top and bottom by $N$ : the term $\frac{b}{N} \to 0$ , leaving $lim_{N \to \infty} \frac{16}{a + \frac{b}{N}} = \frac{16}{a} when a > 0.$ The $a = 0$ case is genuinely different and must be handled separately. If $a = 0$ the denominator is $\frac{b}{N}$ , so the ratio is $\frac{16}{b / N} = \frac{16 N}{b}$ , which grows without bound as $N \to \infty$ — there is no finite ceiling. So:

ZeRO-1 ( $a = 4$ ): ceiling $16/4 = 4 \times$ .
ZeRO-2 ( $a = 2$ ): ceiling $16/2 = 8 \times$ .
ZeRO-3 ( $a = 0$ ): no ceiling — reduction $= \frac{16 N}{16} = N$ , pure $N \times$ scaling. This is the mathematical reason ZeRO-3 is the endgame: nothing stays replicated ( $a = 0$ ), so nothing caps the win.

Figure — reduction ceilings. The plot below shows the reduction factor vs $N$ for each stage. ZeRO-1 (orange) and ZeRO-2 (magenta) flatten toward the dotted horizontal lines at $4 \times$ and $8 \times$ — their finite $a > 0$ ceilings. ZeRO-3 (violet) rises as a straight line $N$ with no ceiling, because $a = 0$ . The picture is the $16/ a$ result made visual: a bigger replicated part $a$ means a lower flat ceiling.

Exercise 5.2

Define compute–communication overlap efficiency as $eff = \frac{t _{compute}}{t _{compute} + t _{comm}}$ where $t_{comm} = C / B$ . For the parent's 175B / 64-GPU case: $C = 2.06$ TB, bandwidth $B = 300$ GB/s, $t_{compute} = 10$ s. Compute efficiency, then find the $B$ needed to reach 80% efficiency.

Recall Solution

Comm time: $t_{comm} = \frac{2.06 \times 1 0 ^{12} bytes}{300 \times 1 0 ^{9} bytes/s} = \frac{2060}{300} \approx 6.87$ s (parent rounds to ~7). Efficiency: $\frac{10}{10 + 6.87} = \frac{10}{16.87} \approx 0.593 = 59.3%$ . ✓ (matches parent's ~59%.)

For 80%: need $\frac{10}{10 + t _{comm}} = 0.8 \Rightarrow 10 = 8 + 0.8 t_{comm} \Rightarrow t_{comm} = 2.5$ s. Then $B = \frac{C}{t _{comm}} = \frac{2.06 \times 1 0 ^{12}}{2.5} = 8.24 \times 1 0^{11} = 824$ GB/s. Mastery reading: to lift efficiency from 59% → 80% you need ~2.75× the bandwidth. Efficiency is bandwidth-bound, which is exactly why hardware (NVLink 4, InfiniBand NDR) drives ZeRO-3's viability.

Exercise 5.3

A student claims: "ZeRO-3 with $N = 64$ gives the same throughput as standard DP with $N = 64$ , just less memory." Steel-man then refute using the numbers of Ex 5.2 versus standard DP's communication $C_{std} = \frac{2Φ ( N - 1 )}{N}$ .

Recall Solution

Standard DP comm bytes: $\frac{2 \times 175 \times 1 0 ^{9} \times 63}{64} = \frac{22050 \times 1 0 ^{9}}{64} \approx 344.5 \times 1 0^{9}$ bytes $\approx 0.345$ TB (per direction), so counting the round trip ≈ $0.69$ TB — one third of ZeRO-3's ~2.06 TB, matching the 3× factor. Standard DP comm time: $\approx \frac{0.69 \times 1 0 ^{12}}{300 \times 1 0 ^{9}} = 2.29$ s ⇒ eff $= \frac{10}{12.29} = 81.4%$ . Refutation: ZeRO-3 (59%) has lower throughput per step than standard DP (81%) at the same $N$ — the extra parameter gathers cost time. The student conflated memory wins with speed. ZeRO-3's win is that it lets you train at all when the model wouldn't otherwise fit; standard DP simply can't run a 175B model on 80 GB cards. Throughput is traded for feasibility.

Recall Self-test recap

Ceiling of ZeRO-1 memory reduction as N→∞ ::: $4 \times$ (params+grads stay replicated) Ceiling of ZeRO-2 ::: $8 \times$ Ceiling of ZeRO-3 ::: unbounded, pure $N \times$ Why ZeRO-3 costs 3× communication vs standard DP ::: AllGather params twice + ReduceScatter grads once vs one AllReduce What ZeRO never partitions ::: activations (need Activation Checkpointing / Gradient Accumulation)