ZeRO and Sharded Optimizers
ZeRO (Zero Redundancy Optimizer) eliminates the wasteful memory duplication of data parallelism by sharding optimizer state, gradients, and weights across GPUs.
Overview
It lets you train enormous models with the simplicity of data parallelism but a fraction of the per-GPU memory.
Deep Dive
In ordinary data parallelism, every GPU stores a redundant full copy of the optimizer state, gradients, and parameters, which is hugely wasteful, especially for Adam, where optimizer state can be several times the size of the model itself. ZeRO, introduced by Microsoft in DeepSpeed, removes this redundancy by partitioning these tensors across GPUs so each device owns only a slice. ZeRO comes in three progressive stages: Stage 1 shards optimizer state, Stage 2 adds gradient sharding, and Stage 3 shards the parameters themselves. As needed, GPUs gather the missing slices via communication, compute, then release them. The result is dramatically lower memory per GPU, enabling billion- to trillion-parameter training, while keeping the easy programming model of data parallelism.
Technical Insight
ZeRO trades extra communication for memory savings. In Stage 3, before a layer's forward pass, an all-gather collects that layer's full parameters onto each GPU; afterward the non-owned slices are discarded to reclaim memory. Gradients are reduce-scattered so each GPU keeps only the gradient slice matching the parameters it owns. PyTorch's FSDP (Fully Sharded Data Parallel) implements the same idea natively, wrapping modules to shard and reshard on the fly.
Strategic Impact
Cost and budget
Architecture decisions drive performance and operating cost for years.
Clearer decisions
Technical education helps teams choose the right stack, not just the newest one.
Quality control
Better engineering choices reduce reliability incidents in production.
The Future of ZeRO and Sharded Optimizers
Sharding is becoming the default for large-scale training rather than an exotic option. Expect deeper integration with offloading (pushing slices to CPU or NVMe via ZeRO-Infinity), better overlap of all-gather and reduce-scatter with computation to hide their cost, and combinations with tensor and pipeline parallelism. As models keep growing, memory-efficient sharded optimizers are central to fitting them onto realistic hardware budgets.
Real-World Implementation
Using DeepSpeed ZeRO Stage 2 to fine-tune a multi-billion-parameter language model that would otherwise overflow GPU memory.
Training with PyTorch FSDP, which shards parameters, gradients, and optimizer state across GPUs and gathers them per layer on demand.
Applying ZeRO-Offload to push optimizer state to CPU memory, letting a single GPU train a model many times larger than its VRAM.
Scaling a trillion-parameter model with ZeRO-Infinity by streaming parameter shards from NVMe storage when GPU and CPU memory run out.
Risks & Guardrails
Optimizing one benchmark can hide broader system weaknesses.
Infrastructure and maintenance costs are often underestimated.
Security and observability gaps can grow as systems become more complex.
Implementation Roadmap
Define latency, quality, and cost targets before implementation.
Benchmark under realistic load and data conditions.
Instrument monitoring for errors, drift, and user impact.
Prepare rollback and incident response paths before scaling.
Keep Exploring
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the ZeRO and Sharded Optimizers quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
Next guide
Lookahead and Lion Optimizers
Frequently asked questions
What is ZeRO and Sharded Optimizers?
ZeRO (Zero Redundancy Optimizer) eliminates the wasteful memory duplication of data parallelism by sharding optimizer state, gradients, and weights across GPUs. It lets you train enormous models with the simplicity of data parallelism but a fraction of the per-GPU memory.
What redundancy does ZeRO eliminate compared to plain data parallelism?
Standard data parallelism stores a full copy of optimizer state, gradients, and weights on every GPU; ZeRO shards these so each GPU holds only a slice.
Why is optimizer state often the biggest memory hog with Adam?
Adam maintains running estimates such as first and second moments per parameter, which combined with fp32 master weights can dwarf the model's own size.
What does ZeRO Stage 3 shard that Stages 1 and 2 do not?
Stage 1 shards optimizer state, Stage 2 adds gradients, and Stage 3 goes further by sharding the model parameters across GPUs as well.
In ZeRO Stage 3, how does a GPU get the full parameters it needs for a layer's forward pass?
Before computing a layer, an all-gather assembles its full parameters on each GPU; once done, the non-owned slices are freed to reclaim memory.
Which PyTorch feature natively implements ZeRO-style sharding?
PyTorch's Fully Sharded Data Parallel (FSDP) shards parameters, gradients, and optimizer state, gathering and resharding them on the fly, mirroring ZeRO.