ZeRO (Zero Redundancy Optimizer)

Partitions the optimizer states (e.g., momentum, variance for Adam) across all data-parallel processes. Each GPU stores and updates only its assigned slice of the optimizer states for the entire model.

• Memory Reduction: Reduces memory footprint proportional to the data-parallel degree (1/N_dp).
• Communication Overhead: Requires an all-gather operation before the optimizer step and a reduce-scatter after the backward pass to synchronize the partitioned states.
• Impact: Enables training of models roughly 4x larger than standard data parallelism with the same number of GPUs.

Extends Stage 1 by also partitioning the model gradients across data-parallel processes. Gradients are reduced to their respective owning GPUs during the backward pass, eliminating their redundancy.

• Memory Reduction: Further reduces memory compared to Stage 1, as full gradients are not stored on any single device.
• Communication Pattern: Uses a reduce-scatter operation during backward pass instead of an all-reduce, which is more communication efficient for this purpose.
• Benefit: Allows for even larger model sizes or larger batch sizes per GPU by freeing gradient memory.

The most memory-efficient stage. Partitions the model parameters themselves across GPUs. Each GPU stores only the slice of parameters it owns.

• Memory Reduction: Memory consumption per GPU is reduced almost linearly with the number of data-parallel processes.
• On-Demand Activation: Parameters are gathered just-in-time for the forward or backward pass and are released immediately after use.
• Trade-off: Introduces additional communication overhead (all-gather/reduce-scatter) for each layer's computation, which can impact throughput. Often used with ZeRO-Infinity for offloading to CPU/NVMe.

ZeRO employs several strategies to mitigate the communication overhead introduced by partitioning.

• Overlapped Communication: The all-gather for parameters (Stage 3) is overlapped with the computation of the previous layer, hiding latency.
• Reduced Precision Communication: Gradients and optimizer states can be communicated in lower precision (e.g., FP16) while maintaining master weights in FP32 for stability.
• Contiguous Memory Buffers: Aggregates small communication operations into larger, contiguous buffers to maximize network bandwidth utilization.

ZeRO is designed to complement, not replace, other forms of model parallelism. It is most powerful when combined into 3D Parallelism:

• Data Parallelism (DP): ZeRO operates within this dimension.
• Tensor Model Parallelism (TP): Splits individual layers across GPUs (e.g., Megatron-LM).
• Pipeline Model Parallelism (PP): Splits the model into sequential stages.

ZeRO (specifically ZeRO-3) can be combined with TP and PP. In this hybrid strategy, parameters are partitioned within each data-parallel group, while TP and PP handle intra-layer and inter-layer partitioning, respectively. This allows for training models at an unprecedented scale.

ZeRO operates in three progressive stages, each eliminating a different type of memory redundancy:

• ZeRO Stage 1 (P_{os}): Partitions the optimizer states (e.g., momentum, variance) across GPUs. Each GPU stores and updates only its assigned partition.
• ZeRO Stage 2 (P_{os+g}): Adds partitioning of gradients to Stage 1. Gradients are reduced across devices only for the partition each GPU is responsible for, slashing memory further.
• ZeRO Stage 3 (P_{os+g+p}): The most memory-efficient stage. Partitions the model parameters themselves. Each GPU stores only the parameters for its partition, with parameters gathered via communication as needed during the forward and backward passes.

These are advanced extensions that push memory savings beyond GPU RAM limits.

• ZeRO-Offload: Intelligently offloads optimizer states and gradients to CPU memory, while keeping parameters on GPU. It uses optimized asynchronous data movement to minimize performance impact, enabling training of models 10x larger on a single GPU.
• ZeRO-Infinity: Extends offloading to include model parameters, leveraging both CPU RAM and NVMe storage. It uses innovative techniques like infinity_offload and block-level management to enable training of models with trillions of parameters (e.g., 1-100T scale) on limited GPU clusters.

ZeRO's memory savings come with a trade-off in inter-GPU communication, which is carefully managed.

• All-Gather: In Stage 3, parameters are gathered from all GPUs before a layer's forward/backward pass, then discarded afterward.
• Reduce-Scatter: Gradients are reduced and scattered to their respective partition owners after the backward pass.
• Overhead Mitigation: Overlap communication with computation to hide latency. For example, while one layer is computing, parameters for the next layer can be gathered asynchronously. The overhead is often minimal compared to the memory gains, especially on high-bandwidth interconnects like NVLink.

ZeRO solves the fundamental memory inefficiency of Data Parallelism (DP).

• DP Redundancy: In standard DP, the entire model (parameters, gradients, optimizer states) is replicated on every GPU. This limits model size to what fits on a single GPU.
• ZeRO Solution: ZeRO removes this replication by partitioning. It maintains the same computational granularity as DP (each GPU processes a micro-batch) but with linear reduction in memory consumption per GPU as more devices are added.
• Result: Enables model sizes that scale almost linearly with the total aggregate GPU memory of the cluster, not the memory of a single device.

ZeRO is particularly synergistic with Sparse Mixture of Experts models, leading to the ZeRO-MoE system.

• Challenge: MoE models have a massive number of parameters (experts), but only a small subset is active per token. Standard parallelism is inefficient.
• Solution: ZeRO Stage 3 partitions the enormous set of experts across devices. The MoE gating network routes tokens to the devices holding their required experts. ZeRO handles the communication for the shared non-expert layers (e.g., attention).
• Outcome: This combination allows for training and inference of models with effectively hundreds of billions to trillions of parameters while maintaining manageable computational cost per token.