Glossary

ZeRO Optimization

A memory optimization technology that partitions model states, gradients, and optimizer states across data-parallel processes to eliminate memory redundancy.

Get in touch Learn more

Data scientist building training data pipeline on laptop, data preprocessing visible, technical workspace.

DISTRIBUTED TRAINING EFFICIENCY

What is ZeRO Optimization?

ZeRO (Zero Redundancy Optimizer) is a memory optimization technology that partitions model states, gradients, and optimizer states across data-parallel processes to eliminate memory redundancy, enabling the training of models with trillions of parameters.

ZeRO Optimization, introduced by Microsoft as part of the DeepSpeed library, fundamentally addresses the memory bottleneck in distributed training. Standard data parallelism replicates the entire model state—including optimizer states, gradients, and parameters—on every GPU, creating massive redundancy. ZeRO eliminates this by partitioning these states across the data-parallel dimension, allowing each device to store only a fraction of the total state, dramatically reducing per-GPU memory consumption.

ZeRO operates in three progressive stages. Stage 1 partitions only the optimizer states, reducing memory by a factor equal to the number of data-parallel processes. Stage 2 additionally partitions gradients, and Stage 3 partitions model parameters themselves, distributing the entire model state. This enables near-linear scaling of model size with the number of GPUs, making it possible to train models with hundreds of billions of parameters without model parallelism.

Memory Optimization

Core Characteristics of ZeRO

ZeRO (Zero Redundancy Optimizer) eliminates memory redundancy in data-parallel training by partitioning model states across GPUs, enabling the training of models with trillions of parameters.

Three-Stage Partitioning Strategy

ZeRO operates in three sequential optimization stages, each partitioning a different component of the model state to progressively reduce memory consumption:

Stage 1 (ZeRO-1): Partitions optimizer states across data-parallel processes, reducing memory by 4x.
Stage 2 (ZeRO-2): Additionally partitions gradients, reducing memory by 8x cumulatively.
Stage 3 (ZeRO-3): Partitions model parameters as well, achieving memory reduction linear with the number of GPUs (Nd x).

Each stage communicates only the necessary shards during forward and backward passes, maintaining computational efficiency while drastically lowering per-GPU memory requirements.

Memory Reduction (Stage 2)

Nd x

Linear Scaling (Stage 3)

Residual State Management

Unlike traditional data parallelism where each GPU holds a full replica of the model, ZeRO treats model states as residual—only the actively required shards are materialized on a given device.

Activation Partitioning: ZeRO can also partition activation memory, which is often the largest consumer during training of massive models.
Offloading: ZeRO-Offload extends the strategy by moving optimizer states and gradients to CPU memory or NVMe storage, enabling training of models larger than aggregate GPU memory.
Communication Overlap: ZeRO overlaps gradient reduction and parameter broadcast with computation to hide communication latency.

This residual approach ensures that memory usage per GPU is a fraction of the total model size, not a multiple of it.

EXPLORE

Communication Primitives and Efficiency

ZeRO leverages high-performance collective communication primitives to manage sharded states without becoming communication-bound:

All-Gather: Used to collect parameter shards before a forward or backward pass.
Reduce-Scatter: Used to aggregate and redistribute gradients efficiently after the backward pass.
Bandwidth Optimization: ZeRO dynamically selects between ring-based and tree-based algorithms depending on message size and topology.

By carefully orchestrating these primitives, ZeRO achieves near-linear scaling efficiency even as the number of GPUs increases, making it practical for clusters with thousands of accelerators.

ZeRO-Infinity for Trillion-Parameter Models

ZeRO-Infinity extends the core ZeRO paradigm to leverage heterogeneous memory hierarchies—combining GPU memory, CPU DRAM, and NVMe storage as a unified memory pool.

Bandwidth-Centric Partitioning: Data placement decisions are driven by the bandwidth characteristics of each memory tier.
Overlap Centric Design: Data movement between tiers is aggressively overlapped with computation to prevent stalls.
Memory-Centric Tiling: Large operators are tiled to process data in chunks that fit within the fastest available memory.

This architecture enabled the first training of models exceeding 1 trillion parameters on a single GPU cluster without requiring model parallelism code changes.

1T+

Parameters Enabled

Integration with DeepSpeed Runtime

ZeRO is implemented as a core component of the DeepSpeed optimization library, providing a drop-in replacement for standard PyTorch optimizers with minimal code changes.

Configuration-Driven: Users specify ZeRO stages and offloading targets via a JSON configuration file.
Automatic Tensor Partitioning: The runtime handles all sharding logic, communication scheduling, and memory management transparently.
Compatibility: Works alongside other DeepSpeed features like mixed-precision training (fp16/bf16) and gradient accumulation.

This tight integration allows engineers to scale from a single GPU to thousands without rewriting training logic, significantly reducing the engineering overhead of large model training.

ZERO REDUNDANCY OPTIMIZATION

Frequently Asked Questions

Explore the core mechanisms behind ZeRO, the memory optimization technology that makes training massive multi-billion parameter models feasible by eliminating data-parallel redundancy.

ZeRO (Zero Redundancy Optimizer) is a memory optimization technology that partitions model states—parameters, gradients, and optimizer states—across data-parallel processes to eliminate memory redundancy, enabling the training of models with trillions of parameters. Unlike classic data parallelism where each GPU holds a full replica of the model, ZeRO shards these states across devices. During the forward and backward passes, it uses all-gather communication collectives to reconstruct the required state partitions on the fly, then discards them. This trades a slight increase in communication overhead for a dramatic reduction in per-device memory consumption, allowing models that would otherwise exceed GPU memory to be trained efficiently.

MEMORY REDUNDANCY ELIMINATION

ZeRO Stages vs. Standard Data Parallelism

Comparison of memory consumption per GPU and communication overhead across ZeRO optimization stages versus standard data parallelism for training large models.

Feature	Standard DP	ZeRO Stage 1	ZeRO Stage 2	ZeRO Stage 3
Optimizer State Partitioning
Gradient Partitioning
Parameter Partitioning
Memory Reduction vs. Standard DP	1x (baseline)	4x	8x	Linear with N GPUs
Communication Volume Overhead	1x (baseline)	1x	1x	1.5x
Full Model Replication per GPU
Offload to CPU/NVMe Supported
Max Model Size (Adam, 32GB GPU)	~1.4B params	~6B params	~12B params	~100B+ params

Enabling Efficiency, Speed & Accuracy

Intelligent Analysis, Decision & Execution

We build AI systems for teams that need search across company data, workflow automation across tools, or AI features inside products and internal software.

Talk to Us

Search across company data

Give teams answers from docs, tickets, runbooks, and product data with sources and permissions.

Useful when people spend too long searching or get different answers from different systems.

Enterprise searchRAGPermissions

Automate internal workflows

Use AI to route work, draft outputs, trigger actions, and keep approvals and logs in place.

Useful when repetitive work moves across multiple tools and teams.

AI agentsWorkflow automationGovernance

Add AI to products and internal tools

Build assistants, guided actions, or decision support into the software your team or customers already use.

Useful when AI needs to be part of the product, not a separate tool.

AI integrationDecision supportModel routing

DISTRIBUTED TRAINING ECOSYSTEM

Related Terms

ZeRO optimization is a foundational technique within a broader ecosystem of distributed training and inference technologies. The following concepts are essential for understanding how model state partitioning interacts with other memory and compute strategies.

DeepSpeed

The deep learning optimization library developed by Microsoft that implements ZeRO redundancy elimination. It provides a complete runtime system for training models with hundreds of billions of parameters across thousands of GPUs.

Implements all three stages of ZeRO (optimizer state, gradient, and parameter partitioning)
Includes ZeRO-Offload to leverage CPU memory and compute
Integrates with PyTorch and popular model frameworks
Enables training of models like Megatron-Turing NLG 530B

EXPLORE

Tensor Parallelism

A complementary strategy to ZeRO that splits individual weight matrices across multiple accelerators within a single layer. While ZeRO partitions model states across data-parallel processes, tensor parallelism shards the computation itself.

Essential when a single layer's parameters exceed one GPU's memory
Often combined with ZeRO in 3D parallelism (data + tensor + pipeline)
Requires high-bandwidth interconnects like NVLink due to frequent all-reduce operations
Used in frameworks like Megatron-LM

Pipeline Parallelism

A technique that partitions model layers sequentially across multiple devices, with each device processing a subset of layers. Micro-batches flow through the pipeline to keep all devices busy.

Reduces memory by distributing layers, but introduces pipeline bubbles (idle time)
GPipe and PipeDream are foundational scheduling algorithms
Combined with ZeRO and tensor parallelism for largest-scale training runs
Effective when model depth exceeds single-device capacity

Gradient Checkpointing

A memory-saving technique that trades compute for memory by discarding intermediate activations during the forward pass and recomputing them during backpropagation.

Reduces activation memory by up to O(√n) with selective checkpointing
Often used alongside ZeRO to further reduce peak memory
Minimal implementation: mark specific layers as checkpoint boundaries
Critical for training deep transformers with long sequences

Mixed Precision Training

The practice of using lower-precision formats (FP16, BF16) for forward and backward passes while maintaining a master copy of weights in FP32 for numerical stability.

Reduces memory consumption by up to 50% compared to full FP32
Loss scaling prevents gradient underflow in FP16
BF16 (bfloat16) preserves dynamic range, eliminating the need for loss scaling
ZeRO partitions the FP32 master weights and FP16 gradients across data-parallel ranks

All-Reduce Communication

The collective communication primitive at the heart of data-parallel training. Each GPU computes local gradients, then all-reduce aggregates and distributes the sum across all participants.

Ring all-reduce optimizes bandwidth by circulating data in a logical ring
ZeRO reduces communication volume by partitioning states, but still requires gradient averaging
NCCL (NVIDIA Collective Communications Library) provides optimized implementations
Communication overhead becomes the primary bottleneck at scale

About the author

Prasad Kumkar

CEO & MD, Inference Systems

Prasad Kumkar is the CEO & MD of Inference Systems and writes about AI systems architecture, LLM infrastructure, model serving, evaluation, and production deployment. Over 5+ years, he has worked across computer vision models, L5 autonomous vehicle systems, and LLM research, with a focus on taking complex AI ideas into real-world engineering systems.

His work and writing cover AI systems, large language models, AI agents, multimodal systems, autonomous systems, inference optimization, RAG, evaluation, and production AI engineering.

Limited slotsGet a Free AI Consultation

How We Work

Custom AI workflows for your Business

One-fit-all AI don't work for modern businesses. At Inferensys, we aim to understand your business & custom requirements; which we use to define most efficient agentic workflows, the data, and the tools for your business.