ZeRO: Deep Memory Optimization for Training Trillion-Parameter Models

In 2020, Microsoft researchers introduced ZeRO (Zero Redundancy Optimizer) via their paper "ZeRO: Memory Optimization Towards Training Trillion Parameter Models" (arXiv:1910.02054). ZeRO is a memory optimization technique that eliminates redundancy in distributed training, enabling efficient scaling to trillion-parameter models. This provides an in-depth technical breakdown of ZeRO's partitioning strategies, memory usage analysis, and integration with DeepSpeed.

1. What is ZeRO?

ZeRO eliminates redundant memory copies of model states across GPUs. Instead of replicating parameters, gradients, and optimizer states across each GPU, ZeRO partitions them across all devices. This results in near-linear memory savings as the number of GPUs increases.

2. Limitations of Traditional Data Parallelism

In standard data-parallel training, every GPU maintains:

• Model Parameters $\theta$
• Gradients $\nabla \theta$
• Optimizer States $O(\theta)$

This causes memory usage to scale linearly with model size, regardless of GPU count.

3. ZeRO’s Three Optimization Stages

 Comparing the per-device memory consumption of model states, with three stages of ZeRO-DP optimizations. Ψ denotes model size (number of parameters), K denotes the memory multiplier of optimizer states, and Nd denotes DP degree. In the example, we assume a model size of Ψ = 7.5B and DP of Nd = 64 with K = 12 based on mixed-precision training with Adam optimizer.

3.1 ZeRO-1: Partitioning Optimizer States

Each GPU stores only a portion of the optimizer states (e.g., Adam's first and second moments), reducing memory from $2 \times |\theta|$ to:

$$ \text{Memory}_{optimizer} = \frac{2 \times |\theta|}{N} $$

3.2 ZeRO-2: Partitioning Gradients

Gradients are no longer reduced to all GPUs. Instead, each GPU owns and updates a subset of $\nabla \theta$, reducing redundant storage.

3.3 ZeRO-3: Partitioning Model Parameters

ZeRO-3 splits the model parameters themselves, allowing a model to scale linearly with the number of GPUs. Forward and backward passes are executed using temporary broadcasts or gathers as needed.

4. Memory Efficiency

The combined memory usage in ZeRO is:

$$ \text{Total Memory}_{ZeRO} = \frac{|\theta| + |\nabla \theta| + |O(\theta)|}{N} $$

Compared to the traditional:

$$ \text{Total Memory}_{DP} = |\theta| + |\nabla \theta| + |O(\theta)| $$

This achieves up to $N$-fold improvement in memory efficiency.

5. DeepSpeed Integration

ZeRO is fully integrated into the DeepSpeed library. DeepSpeed supports ZeRO stages, CPU/NVMe offloading, mixed precision, and efficient activation checkpointing, enabling the training of trillion-scale models with manageable hardware.

6. Experimental Results

ZeRO showed up to 10× larger model capacity and improved throughput on models like BERT-Large and GPT-2. The performance remained stable even as GPU count increased.

7. Limitations and Future Work

While ZeRO increases communication overhead, ZeRO-Infinity (2021) introduced CPU and NVMe offloading to alleviate these costs. Future directions include asynchronous data movement and hierarchical memory systems.

References

• Rajbhandari et al., "ZeRO: Memory Optimization Towards Training Trillion Parameter Models", arXiv:1910.02054
• https://www.deepspeed.ai/
• ZeRO-Infinity: https://arxiv.org/abs/2104.07857