Which training states does ZeRO shard?

The bottleneck

Ordinary data parallelism replicates model and training states per device. Adam moments, variances and mixed-precision master weights can exceed forward-weight memory. ZeRO distributes those states while retaining data-parallel training.

The mechanism

Stage 1 shards optimizer states; Stage 2 also shards gradients; Stage 3 also shards parameters, gathering them when needed. With parameter, gradient and optimizer bytes P,G,O, ideal per-device state goes from P+G+O to P+G+O/n, P+(G+O)/n and (P+G+O)/n.

ZeRO-3 ideal state/device = (P + G + O) / n

Examples and evidence

Example: P=2, G=2, O=12 bytes/parameter and n=4 gives 16 replicated bytes and ideal stage totals of 7, 5.5 and 4. Precision and optimizers change these terms. These are not peak-memory estimates; activations, gather buffers and workspaces are excluded.

New constraints

Deeper sharding adds gathering, communication and scheduling pressure. Network, overlap and microbatch size matter. Fitting the model does not guarantee good throughput; choose configurations beyond theoretical state bytes.