ZeRO

The Problem

Before ZeRO, the standard way to train large models across multiple GPUs was Data Parallelism (DP). In DP, every single GPU holds a complete, identical copy of the entire model, the gradients, and the optimizer states (like Adam's momentum and variance).

As models grew beyond a billion parameters, this became impossible. A 100B parameter model requires hundreds of gigabytes just to store the optimizer states in 32-bit float, far exceeding the memory of any single GPU. Researchers were forced to use complex Model Parallelism, splitting the model vertically, which required massive code rewrites and suffered from severe communication bottlenecks.

The Idea

The Microsoft DeepSpeed team realized that Data Parallelism was wildly inefficient because it replicated the exact same data across every GPU.

They introduced ZeRO (Zero Redundancy Optimizer). Instead of replicating the model states, ZeRO partitions them across all the GPUs in the cluster. When a specific GPU needs a piece of the model state to do its math, it requests that specific piece from whichever GPU holds it over the network.

How It Works

ZeRO is implemented in three distinct stages, offering a spectrum of memory savings vs. communication overhead:

• ZeRO Stage 1: Partitions only the Optimizer States (which take up the vast majority of memory). Achieves 4x memory reduction.
• ZeRO Stage 2: Partitions Optimizer States and Gradients. Achieves 8x memory reduction.
• ZeRO Stage 3: Partitions Optimizer States, Gradients, and Model Parameters.

In Stage 3, no GPU holds the full model. During the forward pass, a GPU broadcasts a request for the parameters of Layer 1. The GPUs holding those partitions send them over. The GPU computes the math, discards the parameters to free memory, and then requests the parameters for Layer 2.

Why It Mattered

ZeRO democratized large-scale model training. It allowed researchers to train models that were 10x to 100x larger than what could fit on a single GPU without requiring them to rewrite their model code for complex model parallelism.

Because ZeRO Stage 1 and 2 introduce virtually zero communication overhead compared to standard Data Parallelism, they became the absolute default standard for distributed training. It is the core technology powering Microsoft's DeepSpeed library.

What Came After

ZeRO became foundational to modern AI infrastructure. It powers the training of models from Megatron-Turing NLG to LLaMA. PyTorch later natively implemented the exact same concept as FSDP (Fully Sharded Data Parallel), which is now the standard built-in tool for training massive models in the PyTorch ecosystem.