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).