Zero Redundancy Optimizer (ZeRO)

Technical Overview of ZeRO

The Zero Redundancy Optimizer (ZeRO) is a memory optimization protocol developed by Microsoft to scale distributed deep learning training.

Standard data parallelism replicates the entire model, gradients, and optimizer states on every GPU, which creates memory bottlenecks for larger models. ZeRO addresses this by sharding these states across the available data-parallel ranks (GPUs), eliminating redundancy while maintaining the training speed of data parallelism.

Key Architecture & Implementation

ZeRO is implemented in three progressive stages, each sharding a different state:

• ZeRO-1 (Optimizer State Sharding):
  • Shards the optimizer states (e.g., FP32 momentum and variance in AdamW) across GPUs. Each GPU updates only its shard of parameters.
  • Reduces memory usage by up to 4x with zero communication overhead.
• ZeRO-2 (Gradient Sharding):
  • Shards gradients across GPUs as they are calculated during the backward pass.
  • Further reduces the memory footprint.
• ZeRO-3 (Parameter Sharding):
  • Shards model parameters across all GPUs. GPUs dynamically gather parameters for specific layers during the forward and backward passes, discarding them immediately after execution.

ZeRO-1: Shard Optimizer States
ZeRO-2: Shard Optimizer States + Gradients
ZeRO-3: Shard Optimizer States + Gradients + Parameters

Core Parameters

• Offloading (ZeRO-Offload): Allows offloading optimizer states and parameters to host CPU memory or NVMe storage, enabling training of larger models on limited GPU resources.

Real-world Applications

• Integrated into DeepSpeed, which has been used to train models like Megatron-Turing NLG (530B).
• Native support in PyTorch via Fully Sharded Data Parallel (FSDP).

The Three Stages of ZeRO

ZeRO is implemented in progressive stages, each sharding more of the training state across GPUs. Stage 1 shards only the optimizer states (which for AdamW-style optimizers can be as large as the model’s own parameters, since it tracks moving averages per parameter). Stage 2 additionally shards gradients. Stage 3 goes further and shards the model parameters themselves, similar in spirit to PyTorch’s FSDP, which was directly inspired by ZeRO Stage 3. Each additional stage reduces per-GPU memory requirements further but introduces more communication overhead, since more state has to be gathered from other GPUs when it’s actually needed for computation.

Why It Matters

Before ZeRO, training very large models required either accepting the memory ceiling imposed by full replication on every GPU, or resorting to model parallelism techniques that split individual layers across GPUs and require more invasive changes to model code. ZeRO’s insight — that most of a model’s memory footprint during training comes from optimizer states and gradients, not the parameters themselves — let large models be trained using data-parallel-style training code with dramatically reduced per-GPU memory needs, making it one of the foundational techniques (via the DeepSpeed library) behind training many large open-weight models.