Understanding ZeRO Stages 1, 2, and 3 in AI Clusters

Training massive deep learning models requires immense computational memory, often exceeding the capacity of a single graphics processing unit (GPU). The Zero Redundancy Optimizer (ZeRO) solves this bottleneck by eliminating memory redundancy across data-parallel artificial intelligence clusters. Instead of replicating the entire model state on every GPU, ZeRO partitions the memory footprint across three progressive levels: Stage 1 divides optimizer states, Stage 2 splits gradients, and Stage 3 shards the model parameters themselves. This article breaks down the mechanics of how each stage partitions these specific components to maximize memory efficiency and enable the training of massive AI models.

The Memory Bottleneck in Standard Data Parallelism

In traditional data-parallel training, every GPU in a cluster holds an identical copy of the entire model. This includes the model parameters (weights), the gradients calculated during the backward pass, and the optimizer states (such as momentum and variance in the Adam optimizer). While this setup allows GPUs to process different batches of data simultaneously, it is highly inefficient regarding memory. Optimizer states alone can consume up to four times the memory of the actual model parameters when using mixed-precision training. ZeRO targets this exact redundancy by distributing these three elements across the available GPUs rather than duplicating them.

ZeRO Stage 1: Optimizer State Partitioning

ZeRO Stage 1 focuses exclusively on the largest memory consumer in mixed-precision training: the optimizer states.

  • What is partitioned: The optimizer states (e.g., 32-bit master weights, momentum, and variance) are divided equally across all GPUs in the cluster. If there are \(N\) GPUs, each GPU stores and updates only \(\frac{1}{N}\) of the total optimizer states.
  • What remains replicated: Both the model parameters and the gradients remain fully duplicated on every GPU.
  • How it works: During the backward pass, each GPU computes the gradients for its data batch. An All-Reduce communication operation averages these gradients across all GPUs. Next, each GPU updates only its specific partition of the optimizer states and the corresponding subset of model parameters. Finally, an All-Gather operation broadcasts the updated parameter subsets to all GPUs, ensuring every device has the complete, updated model for the next forward pass.
  • Memory Impact: Stage 1 reduces the memory footprint by up to a factor of 4 without incurring any additional communication overhead compared to standard data parallelism.

ZeRO Stage 2: Gradient Partitioning

ZeRO Stage 2 builds directly upon Stage 1 by targeting the gradients, which are only needed when updating their corresponding optimizer states.

  • What is partitioned: Both the optimizer states and the gradients are divided equally across all GPUs. Each GPU holds only \(\frac{1}{N}\) of the optimizer states and \(\frac{1}{N}\) of the gradients.
  • What remains replicated: Only the model parameters remain fully duplicated across the cluster.
  • How it works: As gradients are computed layer by layer during the backward pass, they are immediately averaged and sent only to the specific GPU responsible for updating those corresponding parameters (using a Reduce-Scatter operation). Once a GPU sends away the gradients it is not responsible for, it immediately deletes them from its local memory. Each GPU then uses its retained \(\frac{1}{N}\) gradients to update its \(\frac{1}{N}\) optimizer states and parameters. Finally, an All-Gather operation synchronizes the updated parameters across all GPUs.
  • Memory Impact: Stage 2 achieves up to an 8x memory reduction while still maintaining the same communication volume as standard data-parallel training.

ZeRO Stage 3: Parameter Partitioning

ZeRO Stage 3 represents the complete elimination of redundant model states across the AI cluster by partitioning the final component: the model weights.

  • What is partitioned: Optimizer states, gradients, and the model parameters are all divided equally across the cluster. Each GPU stores only \(\frac{1}{N}\) of the entire model state.
  • What remains replicated: Nothing. Every GPU holds a unique shard of the model.
  • How it works: Because no single GPU holds the complete model, GPUs must communicate dynamically during both the forward and backward passes. Before a specific layer is computed in the forward pass, the GPU holding the parameters for that layer broadcasts them to all other GPUs. Once the computation for that layer is finished on a given GPU, the parameters are immediately discarded from memory to free up space for the next layer. This process repeats in reverse during the backward pass, followed by the partitioned gradient and optimizer updates established in Stages 1 and 2.
  • Memory Impact: Stage 3 reduces memory consumption linearly with the number of GPUs. Across a cluster of 64 GPUs, the memory required for model states on a single device is reduced by a factor of 64, trading a modest increase in communication volume for the ability to train models with trillions of parameters.