TL;DR: ZeRO stage 1 shards optimizer state, stage 2 also shards gradients, and stage 3 also shards parameters at rest. Under one explicit 16-byte-per-parameter recipe, 70.6B parameters at 64 ranks need 295.64 GB, 156.64 GB and 17.65 GB of persistent payload per rank respectively. These are not peak allocation or proof that a training configuration fits.
How to approach it
List the recipe's actual dtypes and copies before using a bytes-per-parameter rule. Identify which terms stay replicated, derive each stage, then add transient live state. End with the communication and runtime tradeoff of the least-sharded feasible candidate.
A strong answer
A spreadsheet says stage 3 needs less than 18 GB per GPU, yet training still runs out of memory. The spreadsheet counted stored shards, while execution also needed gathered parameters and activations. Model Memory Footprint establishes the accounting boundary; ZeRO and FSDP explains sharded ownership.
ZeRO (Zero Redundancy Optimizer) removes duplicate copies of training state across data-parallel ranks, giving each rank one shard instead of a full copy. The DeepSpeed ZeRO overview defines the incremental stages. For this worked example, assume BF16 parameters and gradients, an FP32 master parameter copy, and two FP32 Adam moments. Other recipes differ; 16 bytes is not a universal mixed-precision requirement.
W = 2 bytes of parameters
G = 2 bytes of gradients
O = 4 master + 4 first moment + 4 second moment = 12 bytes
unsharded payload = W + G + O = 16 bytes/parameter
N = 70.6e9 parameters; n = 64 data-parallel ranks
stage 1 = N×(W+G+O/n) = 295.6375 GB/rank
stage 2 = N×(W+(G+O)/n) = 156.64375 GB/rank
stage 3 = N×(W+G+O)/n = 17.65 GB/rank
sanity: 64×17.65 = 1,129.6 GB, the original total payload.
The graphic uses four ranks so ownership is easy to see; the numeric example uses 64. Stage 1 leaves four bytes per parameter replicated under this recipe. Stage 2 leaves two, not four. Their asymptotic floors for 70.6B are therefore 282.4 GB and 141.2 GB respectively, before the remaining shards. Neither floor fits one 80 GB device without additional measures.
| Ranks | Stage 3 persistent payload |
|---|---|
| 8 | 141.2 GB |
| 16 | 70.6 GB |
| 32 | 35.3 GB |
| 64 | 17.65 GB |
Even 70.6 GB below a nominal 80 GB device capacity is only a candidate: gathered units, activations, temporary kernels, communication buffers, allocator behavior and other reservations still need room. There is no universal practical floor of 32 GPUs independent of the model implementation and workload.
A simple communication model helps describe the cost. Let P be parameter bytes and G gradient bytes for a full corresponding model partition, with q=(n−1)/n. A ring reduce-scatter sends qG per rank; one parameter all-gather sends qP. If parameters are gathered once for forward and again for backward, the illustrative send total is q(2P+G). When P=G this is 1.5 times a ring gradient all-reduce's 2qG. Retaining gathered parameters, prefetching, accumulation and different dtypes change the count.
That is not a universal ZeRO-3 traffic multiplier. Stage 2 can reduce and release gradient buckets while retaining only its ownership shard; temporary full tensors or bucket buffers can still exist. Stage 3 divides stored state among ranks, but live allocations can exceed one rank’s share during execution. A device memory reading alone cannot tell you which stage is configured.
Inspect the framework configuration, parameter ownership, allocator peak and collective timeline. Compare candidates at equal numerical settings and batch. Parameter prefetch can hide waits but increases live memory. If stage 2 meets the memory requirement and runs faster, stage 3 has no automatic advantage; if stage 2 cannot fit, its speed on a smaller workload is not a valid comparison.
What interviewers probe next
- “Are results bitwise identical to DDP?” The intended update can match, but reduction order and numerical implementations can change floating-point results. Test with an appropriate tolerance and training-quality check.
- “Are activations divided by n?” Not merely because parameters are sharded. Activation memory depends on the computation, micro-batch, sequence and checkpointing policy.
- “Can optimizer sharding help before stage 3?” Yes. For a 7B version of this recipe at eight ranks, stage 1 is
7×(4+12/8)=38.5 GBbefore execution extras.
Common mistakes
- Applying the 16-byte recipe without checking the optimizer and dtypes.
- Saying stage 2 still replicates full reduced gradients at rest.
- Treating stored-shard arithmetic as a peak-memory measurement.
Key takeaways
- The stated 64-rank recipe gives 295.64, 156.64 and 17.65 GB per rank across stages 1–3.
- Stage 2's replicated floor is two bytes per parameter under this recipe.
- Gather lifetimes and activations decide whether the static estimate fits in practice.