ZeRO Stage Comparison and Memory Optimization

Introduction

Welcome back. In our previous lesson, we established the fundamental memory challenge in large-scale training: the $16\Psi$ memory cost per GPU for model states (parameters, gradients, and optimizer states) in standard data parallelism. We introduced the Zero Redundancy Optimizer (ZeRO) as a solution that systematically eliminates this memory redundancy by partitioning these states across devices.

This lesson builds directly on that foundation. Our goal is to move from a high-level overview to a precise, quantitative comparison of the three ZeRO stages. You will learn to analyze the specific memory reduction each stage provides and its corresponding impact on communication volume. Understanding these trade-offs is critical for designing efficient training systems and selecting the right optimization strategy for a given model and hardware configuration.

By the end of this lesson, you will be able to compare the different stages of ZeRO and their corresponding reductions in memory consumption, a key skill for optimizing large model training.

A Framework for Comparison

To systematically compare the ZeRO stages, we'll analyze each one through the lens of four key metrics:

1. Partitioned States: Which components of the model state are sharded across devices?
2. Replicated States: Which components remain duplicated on every device?
3. Per-GPU Memory: The resulting memory footprint per device, expressed as a formula.
4. Communication Volume: The total amount of data communicated per step, relative to standard data parallelism.

Let's start by defining our baseline. As we established, for a model with $\Psi$ parameters using mixed-precision training with the Adam optimizer:

• Baseline (Standard Data Parallelism):
  • Partitioned States: None.
  • Replicated States: Parameters ($2\Psi$), Gradients ($2\Psi$), and Optimizer States ($12\Psi$).
  • Per-GPU Memory: $16\Psi$ bytes.
  • Communication Volume: $2\Psi$ bytes. This comes from the AllReduce operation on the gradients, where each GPU effectively sends and receives the entire gradient tensor.

The diagram below, which we saw briefly in the last lesson, provides the essential visual intuition for what follows. It shows how the memory on each GPU (from gpu0 to gpuN-1) is progressively reduced as we move from the baseline to the more aggressive ZeRO stages.

Now, let's dissect each stage using our framework.

ZeRO Stage 1: Optimizer State Partitioning (Pos)

Stage 1 targets the largest single consumer of memory: the optimizer states.

• Partitioned States: Optimizer States ($12\Psi$).

• Replicated States: FP16 Parameters ($2\Psi$) and FP16 Gradients ($2\Psi$).

• Per-GPU Memory: The optimizer state memory is sharded across $N_d$ devices. The total per-GPU memory for model states becomes:

  $${\text{Mem}}_{\text{S1}}=\underset{\text{Params}}{\underbrace{2\Psi }}+\underset{\text{Gradients}}{\underbrace{2\Psi }}+\underset{\text{Optimizer States}}{\underbrace{\frac{12\Psi }{N_d}}}=4\Psi +\frac{12\Psi }{N_d}$$

  As $N_d$ becomes large, the $\frac{12\Psi }{N_d}$ term becomes negligible, and the memory approaches $4\Psi$. This represents approximately a 4x reduction from the baseline $16\Psi$.

• Communication Volume: The communication volume remains the same as standard data parallelism ($2\Psi$). This is a crucial benefit.

You might wonder how this is possible, since each GPU only updates a partition of the parameters. The process ensures consistency without extra communication:

1. Gradients are fully synchronized across all GPUs via a standard AllReduce operation.
2. Each GPU updates its assigned partition of the optimizer states and, consequently, its partition of the master FP32 parameters.
3. An AllGather operation is then performed, where each GPU broadcasts its updated parameter partition to all other GPUs.

The key insight is that the standard AllReduce collective can be implemented as a ReduceScatter followed by an AllGather. ZeRO Stage 1 effectively performs the ReduceScatter and AllGather on different data (gradients vs. parameters), but the total volume of data moved remains the same.

ZeRO Stage 2: Gradient and Optimizer State Partitioning (Pos+g)

Stage 2 builds on Stage 1 by also partitioning the gradients.

• Partitioned States: Optimizer States ($12\Psi$) and Gradients ($2\Psi$).

• Replicated States: FP16 Parameters ($2\Psi$).

• Per-GPU Memory: Now, both gradients and optimizer states are sharded.

  $${\text{Mem}}_{\text{S2}}=\underset{\text{Params}}{\underbrace{2\Psi }}+\underset{\text{Gradients}}{\underbrace{\frac{2\Psi }{N_d}}}+\underset{\text{Optimizer States}}{\underbrace{\frac{12\Psi }{N_d}}}=2\Psi +\frac{14\Psi }{N_d}$$

  For large $N_d$, the memory approaches $2\Psi$, yielding an 8x reduction from the baseline.

• Communication Volume: Remarkably, the communication volume remains the same as standard data parallelism ($2\Psi$). Stage 2 replaces the full AllReduce on gradients with a ReduceScatter. During the backward pass, this operation averages the gradients and distributes to each GPU only the partition it needs for its subsequent parameter update. This avoids storing the full gradient tensor on every device. The communication for the final AllGather of updated parameters remains, and the combined volume of ReduceScatter + AllGather is equivalent to a single AllReduce.

ZeRO Stage 3: Full State Partitioning (Pos+g+p)

Stage 3 is the most aggressive optimization, partitioning everything.

• Partitioned States: Optimizer States ($12\Psi$), Gradients ($2\Psi$), and Parameters ($2\Psi$).

• Replicated States: None (beyond transient buffers).

• Per-GPU Memory: All model states are now sharded.

  $${\text{Mem}}_{\text{S3}}=\underset{\text{Params}}{\underbrace{\frac{2\Psi }{N_d}}}+\underset{\text{Gradients}}{\underbrace{\frac{2\Psi }{N_d}}}+\underset{\text{Optimizer States}}{\underbrace{\frac{12\Psi }{N_d}}}=\frac{16\Psi }{N_d}$$

  This provides a linear memory reduction of $N_d$-fold. It effectively makes the aggregate memory of all GPUs in the cluster available, enabling the training of models of practically any size.

• Communication Volume: This is the first stage that increases communication. The volume becomes 1.5x that of standard data parallelism ($3\Psi$). The increase comes from the need to gather the full parameters for each layer during the forward and backward passes. The communication flow for each layer is:

  1. AllGather the necessary parameters for the current layer's forward computation ($\approx \Psi$ total volume over a full pass).
  2. After the backward pass, ReduceScatter the gradients as in Stage 2 ($\Psi$ total volume).
  3. After the local parameter update, AllGather the updated parameters as in S1/S2 is not needed because they are already partitioned. The first AllGather of parameters for the next training step handles this. However, the paper's 1.5x figure implies a total of $3\Psi$. Let's follow the clearer logic from the blog post: AllGather for forward pass ($\Psi$) + ReduceScatter for backward pass ($\Psi$) + an additional communication step to sync updated weights, totaling $3\Psi$. This modest 50% increase in communication is often a worthwhile price for the massive memory savings.

Summary and Quantitative Comparison

Let's consolidate our findings into a summary table.

To see these formulas in action, let's look at the concrete numbers for different model sizes from the original ZeRO paper.

As you can see from the table, for a 7.5B parameter model on 64 GPUs:

• Baseline memory is 120 GB.
• ZeRO-1 (Pos) reduces it to 31.4 GB.
• ZeRO-2 (Pos+g) reduces it further to 16.6 GB.
• ZeRO-3 (Pos+g+p) brings it down to an incredible 1.88 GB.

This clearly demonstrates the power of each successive stage. Finally, let's watch a short summary from one of the creators of ZeRO, which recaps these benefits concisely.

Conclusion

In this lesson, we conducted a detailed, quantitative comparison of the three stages of the Zero Redundancy Optimizer. By systematically analyzing what each stage partitions, we derived the specific memory savings and communication costs, moving from abstract concepts to concrete formulas and numbers.

Key Takeaways:

• ZeRO-1 partitions optimizer states, offering up to a 4x memory reduction with no communication overhead compared to standard data parallelism.
• ZeRO-2 partitions both optimizer states and gradients, yielding up to an 8x memory reduction, also with no communication overhead.
• ZeRO-3 partitions all model states (parameters, gradients, and optimizer states), achieving a linear $N_d$-fold memory reduction. This immense memory saving comes at the cost of a modest 1.5x increase in communication volume.
• The choice of ZeRO stage presents a clear design trade-off: Stages 1 and 2 provide significant "free" memory savings, while Stage 3 offers the maximum possible memory reduction at the cost of increased communication and implementation complexity.

Preview of the Next Lesson:

Now that you can precisely compare the ZeRO stages, the next logical step is to understand how these partitioning strategies are realized in practice. In the next lesson, we will explore how memory sharding is implemented in leading frameworks like DeepSpeed (which houses the canonical ZeRO implementation) and PyTorch's native Fully Sharded Data Parallel (FSDP), connecting this theory to the tools used to train today's largest models.