ZeRO Optimizer & Memory Sharding
The hidden memory cost of ordinary data parallelism
In standard data parallelism, every device holds a full, redundant copy of the model's parameters, gradients, and optimizer state — for Adam, that's two additional moment buffers per parameter, on top of the parameters and gradients themselves. For a model with $\Phi$ parameters in mixed precision, a conservative estimate puts per-device memory at roughly $16\Phi$ bytes (fp16 params + fp16 gradients + fp32 params + fp32 Adam moments), *replicated identically on every single device*. With enough data-parallel replicas, this redundancy becomes the actual memory bottleneck — not the compute, but the sheer waste of every device storing information that, in aggregate across the cluster, is massively duplicated.
ZeRO's core idea: shard, don't replicate
ZeRO (Zero Redundancy Optimizer), developed for the DeepSpeed training library, eliminates this redundancy by sharding the optimizer states, gradients, and (in its most aggressive stage) the parameters themselves across data-parallel devices, instead of replicating them fully on each one. Each device only permanently stores its own shard; when a different device's shard of parameters is needed for a computation, it's temporarily gathered via communication, used, and then released. The net effect: the same aggregate memory that used to be N-fold redundant across N devices is now split roughly N ways, enabling training of models the aggregate cluster memory can hold even though no single device could hold a full redundant copy.
ZeRO's three stages
| Stage | What's sharded | Memory saved vs. standard data parallelism | Extra communication |
|---|---|---|---|
| ZeRO-1 | Optimizer states only | Up to ~4x less memory | Minimal added overhead |
| ZeRO-2 | Optimizer states + gradients | Up to ~8x less memory | Slightly more |
| ZeRO-3 | Optimizer states + gradients + parameters | Scales ~linearly with number of devices | Most — parameters gathered on demand per layer |
# DeepSpeed config enabling ZeRO Stage 3 (full parameter, gradient,
# and optimizer state sharding) with CPU offload for additional savings.
ds_config = {
"train_batch_size": 256,
"bf16": {"enabled": True},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {"device": "cpu", "pin_memory": True},
"offload_param": {"device": "cpu", "pin_memory": True},
"overlap_comm": True, # overlap communication with computation
"contiguous_gradients": True,
"stage3_gather_16bit_weights_on_model_save": True,
},
}
# Usage: deepspeed.initialize(model=model, config=ds_config, ...)
# Each GPU now stores only 1/N of params, gradients, and optimizer states,
# gathering shards from peers on demand during forward/backward passes.ZeRO vs. model parallelism: a different axis entirely
It's easy to conflate ZeRO with tensor or pipeline parallelism since both reduce per-device memory, but they operate on different axes. Tensor and pipeline parallelism change *how computation is distributed* — different devices do genuinely different parts of the math. ZeRO keeps the data-parallel computation pattern (every device still computes forward/backward on its own data batch using the full model logically) but changes *how memory is stored*, gathering and releasing shards transparently so that, from a correctness standpoint, training proceeds identically to standard data parallelism — just with far less redundant storage. In practice, ZeRO and model parallelism are frequently combined.
PyTorch's FSDP is ZeRO's ideas, built in
Sharding trades memory for communication and complexity
What's next?
Memory sharding tackles one dimension of the training-at-scale problem; another, orthogonal lever — useful even on a single device — is adjusting the effective batch size without changing available memory at all, which is exactly what gradient accumulation does.
I build these systems professionally.
Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.