Data parallelism and ZeRO optimization: Question bank

Before we start, three words we lean on everywhere below, in plain language:

True or false — justify

TF1. In standard data parallelism, every GPU sees the same data batch.

False — every GPU sees a different slice of data; that is the whole point. What is identical across GPUs is the model copy, and (after synchronization) the gradients.

TF2. ZeRO-3 is a form of pipeline or tensor parallelism because it splits the model.

False — ZeRO-3 is still data parallelism: each GPU processes a different data batch and the full parameter set is logically present. It only shards storage; parameters are gathered back on demand so the math is identical to standard DP. Pipeline Parallelism splits the model by layer stage instead.

TF3. ZeRO changes the final trained model compared to standard data parallelism.

False — ZeRO is a pure memory/communication rearrangement. Every stage computes the same averaged gradient and the same update, so the resulting weights are bit-for-bit equivalent (up to floating-point ordering).

TF4. Going from ZeRO-1 to ZeRO-3 always makes training faster.

False — it makes it fit in less memory, but ZeRO-3 needs ~3× the communication of standard DP. On a slow interconnect it can be dramatically slower wall-clock even though memory drops.

TF5. Optimizer states can be partitioned (ZeRO-1) because each parameter's momentum and variance depend only on that parameter's own gradient.

True — Adam's per-parameter state updates are element-wise. GPU

r

can update its slice

θ_{r}

using only

\nabla θ_{r}

and

θ_{r}

's states, with no cross-talk.

TF6. With

N

GPUs, ZeRO-3 divides total cluster memory for model state by

N

False — it divides per-GPU memory by

N

; the total memory summed over the cluster stays at

16Φ

(one full copy, just spread out). Standard DP stored

16Φ \times N

total.

TF7. ZeRO-2 uses AllReduce to synchronize gradients.

False — ZeRO-2 uses ReduceScatter: each GPU ends up holding only the averaged gradient for its partition, not the full averaged gradient. See AllReduce and Collective Communication.

TF8. Because ZeRO-3 partitions parameters, a GPU never holds a full layer's weights during compute.

False — it temporarily AllGathers a layer's parameters right before using them, computes, then discards them. The full weights briefly materialize; they just aren't stored permanently.

TF9. Doubling

N

(GPUs) roughly halves ZeRO-3 per-GPU memory.

True —

M = 16Φ/ N

, so per-GPU model-state memory is inversely proportional to

N

. (Activation memory is separate and does not shrink this way — see edge cases.)

TF10. ZeRO removes the need for mixed precision.

False — the

2Φ

(fp16) +

12Φ

(fp32 states incl. master copy) breakdown assumes Mixed Precision Training. ZeRO partitions those numbers; it doesn't replace the reason they exist.

Spot the error

SE1. "ZeRO-1 saves memory by making gradients smaller, storing only

\frac{1}{N}

of them."

Wrong stage — ZeRO-1 partitions only optimizer states; gradients stay fully replicated at

2Φ

. Gradient partitioning is ZeRO-2.

SE2. "ZeRO-3 communication is 3× standard DP because we AllGather parameters three times per step."

The factor is 3, but not from three AllGathers — it's two AllGathers (forward + backward) of

2Φ

each plus one ReduceScatter of

2Φ

, giving

6Φ (N - 1) / N

2Φ (N - 1) / N

SE3. "Optimizer states are

12Φ

because Adam stores momentum and variance, each 4 bytes fp32."

Incomplete — that's only

8Φ

. The missing

4Φ

is the fp32 master copy of the weights, needed so tiny updates don't underflow in fp16. All three sum to

12Φ

SE4. "We keep optimizer states in fp16 to save memory."

Wrong — they are deliberately fp32. Momentum/variance and the master weights need the precision so that

η \cdot grad

(often extremely small) doesn't vanish. Saving them in fp16 would break the update. Memory-Efficient Optimizers tackle this differently.

SE5. "In standard DP with 8 GPUs holding a 2800 GB model state, total memory used is 2800 GB."

Wrong — it's

8 \times 2800 = 22, 400

GB, because the state is replicated on every GPU. Only 2800 GB is genuinely needed; the rest is redundancy — exactly what ZeRO attacks.

SE6. "ReduceScatter is more expensive than AllReduce, so ZeRO-2 costs more to communicate than standard DP."

Wrong — AllReduce is effectively ReduceScatter + AllGather. ZeRO-2 replaces the update-time AllReduce with a ReduceScatter, so its gradient-sync volume is comparable, not larger; it saves memory at similar comms cost.

SE7. "ZeRO-3 memory per GPU is

16Φ/ N

, so activations also shrink by

N

Wrong —

16Φ/ N

is only the model state. Activation memory depends on batch size and sequence length and is not partitioned by ZeRO; it needs separate tools like Activation Checkpointing.

SE8. "Since ZeRO-3 makes a 175B model fit on 64 A100s, adding more GPUs is free memory-wise."

Wrong — communication volume

6Φ (N - 1) / N \to 6Φ

N

grows, so per-GPU memory keeps shrinking but comms stops improving and per-GPU throughput can drop from network overhead. There's a sweet spot, not free scaling.

SE9. "The gradient of the full batch equals the sum of per-GPU gradients, so we sum in AllReduce."

Careful — it's the average (

\frac{1}{N} \sum

), matching how the loss is a mean over samples. Summing without dividing would inflate the effective learning rate by

N

Why questions

WY1. Why does ZeRO-2 use ReduceScatter instead of a plain AllReduce?

Because after ZeRO-2 each GPU only updates its own partition

θ_{r}

, so it only needs the averaged gradient for that slice. ReduceScatter delivers exactly that slice, letting the GPU discard the rest and store only

2Φ/ N

of gradients.

WY2. Why can't we partition parameters (ZeRO-3) "for free" like optimizer states?

Because a forward pass through layer

l

needs all of

θ_{l}

's weights at once, not just one GPU's slice. So ZeRO-3 must pay to AllGather them on demand — the memory saving is bought with extra communication.

WY3. Why is the optimizer state the first and most valuable thing to partition?

Because it is the largest term (

12Φ

of the

16Φ

total, ~75%). Splitting the biggest block first gives the best memory-per-communication return, which is why ZeRO-1 alone already yields ~2.9× at

N = 8

WY4. Why does averaging gradients across GPUs reduce variance rather than bias?

Each GPU computes an unbiased noisy estimate of the true gradient from its data slice. Averaging

N

unbiased estimates keeps the mean correct (no bias) but shrinks the noise by roughly

N

— same effect as a bigger batch.

WY5. Why does ZeRO-3 AllGather parameters twice per step, not once?

Once in the forward pass (to compute activations) and again in the backward pass (weights are needed to compute their gradients). Since parameters were discarded after forward, they must be re-gathered — hence the "2×" in the comms formula.

WY6. Why is ZeRO-3 only practical on fast interconnects like NVLink?

Its comms volume is ~3× standard DP. On NVLink (~300 GB/s) that overlaps with compute acceptably; on 10 Gbps Ethernet (~1 GB/s) the multi-terabyte transfers dominate and stall the GPUs, wrecking efficiency.

WY7. Why does Gradient Accumulation pair naturally with ZeRO for large models?

It grows the effective batch without more memory (accumulate several micro-batch gradients before one update), and it amortizes ZeRO's communication — you sync/gather less often per token processed, softening the comms penalty.

WY8. Why does the memory formula use

2Φ

for parameters but

4Φ

for the master copy of the same parameters?

The working copy is fp16 (2 bytes) for fast matmuls; the master copy is fp32 (4 bytes) so accumulated updates stay precise. They are the same weights at two precisions serving two jobs — compute vs. accurate accumulation.

Edge cases

EC1. What is per-GPU model-state memory when

N = 1

(single GPU)?

All ZeRO formulas collapse to

16Φ

— identical to standard DP. With one GPU there is nothing to partition and no communication, so ZeRO gives zero benefit.

EC2. As

N \to \infty

, what happens to ZeRO-3 per-GPU memory and to its communication volume?

Per-GPU model-state memory

16Φ/ N \to 0

, but communication

6Φ (N - 1) / N \to 6Φ

(a constant ceiling). So memory keeps improving while comms plateaus — beyond some

N

you pay network cost for shrinking memory return.

EC3. Does ZeRO help if the bottleneck is activation memory, not model state?

No — ZeRO partitions only parameters/gradients/optimizer states. If activations blow up your memory, you need Activation Checkpointing or smaller batches; ZeRO alone won't save you.

EC4. A parameter count

Φ

not divisible by

N

— is partitioning still valid?

Yes — partitions are just made uneven (some GPUs hold slightly more). The formulas are

1/ N

approximations; the collectives handle ragged slices, and the extra imbalance is negligible for large

Φ

EC5. On a "fat" model that already fits in one GPU's memory, is ZeRO-3 worth it?

Usually no — you'd pay 3× communication for memory you don't need. Standard DP or at most ZeRO-1 is preferable; ZeRO-3 earns its cost only when the model state genuinely can't fit otherwise, as in Large Language Models Training.

EC6. If two GPUs happen to receive identical data batches, does data parallelism break?

It doesn't crash, but it wastes work — those GPUs produce identical gradients, so averaging them adds no variance reduction. Effective batch size is smaller than intended; you should feed disjoint shards.

Recall Quick self-test

Which ZeRO stage first partitions gradients? ::: ZeRO-2 (ZeRO-1 = optimizer states only; ZeRO-3 adds parameters). Standard DP total cluster memory for model state with $N$ GPUs? ::: $16Φ \times N$ (full replication). ZeRO-3 communication multiplier vs standard DP? ::: About 3× (two param AllGathers + one gradient ReduceScatter).