Zero Redundancy Optimizer (ZeRO)

This is the foundational stage of ZeRO. It partitions the optimizer states (e.g., momentum, variance for Adam) across all data-parallel processes. Each GPU stores and updates only its assigned portion of the optimizer states, reducing memory usage proportional to the data-parallel degree.

• Memory Reduction: The memory footprint for optimizer states is divided by the number of GPUs (D_pp).
• Communication Overhead: Minimal. Only requires a single all-gather operation when the full optimizer state is needed for a parameter update, which is infrequent.
• Impact: Enables training models with ~3x larger batch sizes or model sizes compared to standard data parallelism.

Builds on Stage 1 by also partitioning the gradients across devices. After the backward pass, each GPU only retains the gradients corresponding to its assigned portion of the parameters. A reduction operation is performed to aggregate gradients across devices only for the locally owned slices.

• Memory Reduction: Eliminates the memory redundancy for gradients, in addition to optimizer states.
• Communication Pattern: Uses a reduce-scatter operation during gradient synchronization instead of an all-reduce. This is more communication-efficient for this purpose.
• Impact: Enables training models roughly 8x larger than possible with standard data parallelism, as both major training memory components (optimizer states and gradients) are partitioned.

The most memory-aggressive stage. It partitions the model parameters themselves across GPUs. Each GPU stores only the slice of parameters it is responsible for. Parameters are gathered on-demand from all devices just before computation in the forward or backward pass and are discarded immediately after.

• Memory Reduction: The model's parameter memory is divided by the number of GPUs. This enables training models with trillions of parameters.
• Communication Overhead: Significant. Requires all-gather operations before every layer's computation in both forward and backward passes, and reduce-scatter after backward passes. This trades increased communication for massive memory savings.
• Use Case: Essential for training the largest foundational models (e.g., models with hundreds of billions to trillions of parameters).

ZeRO's performance depends critically on optimizing communication. It replaces standard all-reduce with more efficient collective operations and overlaps them with computation.

• Reduce-Scatter: Used in Stage 2/3 for gradient synchronization. Each GPU ends up with a unique slice of the summed gradients.
• All-Gather: Used in Stage 3 to collect parameters from all GPUs before computation.
• Overlap Strategy: Modern implementations (e.g., in DeepSpeed) carefully schedule these communications to overlap with independent computation (e.g., the forward pass of one layer with the all-gather for the next layer), hiding much of the communication latency.

ZeRO is primarily a data parallelism optimizer. For maximum scalability, it is combined with other forms of parallelism in 3D Parallelism.

• ZeRO + Model Parallelism (Tensor Parallelism): ZeRO partitions data-parallel groups, while tensor parallelism splits individual layers (e.g., attention heads, MLP blocks) across devices within a group. This combination reduces memory per device and communication volume for very large layers.
• ZeRO + Pipeline Parallelism: Pipeline parallelism splits the model vertically by layers. ZeRO can be applied within each pipeline stage to further reduce memory usage for the layers resident on that stage's devices.
• Orchestration: Frameworks like DeepSpeed and PyTorch Fully Sharded Data Parallel (FSDP) manage this complex orchestration, allowing the training of the world's largest AI models.

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

• ZeRO Stage 1 (P_os): Partitions optimizer states (e.g., momentum, variance) across data-parallel processes. Reduces memory by ~4x.
• ZeRO Stage 2 (P_os+g): Adds partitioning of gradients. Each process only stores gradients for its assigned parameters, reducing memory by ~8x.
• ZeRO Stage 3 (P_os+g+p): Partitions the model parameters themselves. Parameters are gathered on-demand for the forward/backward pass and discarded afterward, enabling near-linear memory scaling with the number of GPUs.

ZeRO enhances, rather than replaces, standard Data Parallel (DP) training.

• Classic DP: Replicates the entire model, optimizer states, and gradients on every GPU. Communication is a single all-reduce of gradients. Memory usage is identical on all devices.
• ZeRO-powered DP: Maintains the data parallelism paradigm but partitions the model states. It uses a reduce-scatter operation to distribute gradient aggregation and an all-gather to collect parameters, trading increased communication for massive memory savings. The computational workload per device remains the same.

ZeRO Stage 3 introduces significant communication for gathering parameters. Key optimizations mitigate this cost:

• Overlapping Communication and Computation: Parameters for the next layer are gathered while the current layer computes.
• Parameter Persistence: Keeping a small subset of critical parameters (e.g., embedding layers) resident on each GPU.
• Hierarchical Partitioning: In multi-node setups, parameters are partitioned first across nodes, then within nodes, to optimize for slower inter-node links. The communication volume is approximately 1.5x that of classic DP, but the enabling of larger batch sizes often results in better aggregate throughput.

ZeRO is particularly effective for training Mixture of Experts (MoE) models like Google's Switch Transformers.

• In MoE models, only a subset of 'expert' layers are activated per token, creating a naturally sparse, massive model.
• Applying ZeRO Stage 1 to the shared optimizer states and ZeRO Stage 2 to the sparse gradients is highly efficient.
• The combination allows for models with over a trillion parameters while keeping the computational cost per token similar to a dense model 1/10th the size.

ZeRO is often combined with other forms of parallelism for training extreme-scale models:

• ZeRO + Tensor Parallelism: ZeRO handles data parallelism across nodes, while tensor parallelism splits individual layers (e.g., attention heads, MLP blocks) across GPUs within a node. This is used in Megatron-DeepSpeed.
• ZeRO + Pipeline Parallelism: ZeRO optimizes memory within a pipeline stage, while pipeline parallelism splits the model layers sequentially across devices. This hybrid approach is standard for models exceeding hundreds of billions of parameters.
• The choice of hybrid strategy depends on the cluster topology, model architecture, and communication bandwidth.

A distributed training strategy where the parameters and computational graph of a single neural network are partitioned across multiple devices. Unlike data parallelism, where each device holds a full model copy, model parallelism splits the model itself.

• Tensor Parallelism: A specific form where individual tensor operations (e.g., matrix multiplications within a linear layer) are split across devices.
• Pipeline Parallelism: Partitions the model by layers into stages, forming a processing pipeline where micro-batches flow sequentially.
• ZeRO Integration: ZeRO's stages 2 and 3 can be viewed as a form of model parallelism for optimizer states, gradients, and parameters, often combined with data parallelism in frameworks like DeepSpeed.

A memory optimization technique that trades compute for memory by selectively storing only a subset of layer activations during the forward pass. Non-checkpointed activations are recomputed during the backward pass as needed.

• Core Trade-off: Can reduce activation memory by up to 80-90% at the cost of approximately 30% more compute time.
• Synergy with ZeRO: Often used in conjunction with ZeRO. While ZeRO partitions optimizer components, checkpointing reduces the memory footprint of activations, together enabling the training of models with trillions of parameters.
• Implementation: Available in major frameworks (e.g., torch.utils.checkpoint). Critical for training very deep networks like those in large language models.

A model parallelism strategy that partitions a network's layers across multiple devices. The model is divided into sequential stages, each placed on a different device. Micro-batches of data are pumped through this pipeline to keep all devices utilized.

• Bubble Overhead: A key challenge is the 'pipeline bubble'—idle time at the start and end of a batch as the pipeline fills and drains. Techniques like 1F1B (One Forward pass followed by One Backward pass) scheduling minimize this.
• Communication Pattern: Requires point-to-point communication of activations and gradients between adjacent stages in the pipeline.
• DeepSpeed & PyTorch FSDP: Modern implementations (e.g., DeepSpeed's pipeline, FairScale's Pipe) combine pipeline parallelism with ZeRO data parallelism for extreme model scale.

A collective communication operation fundamental to data-parallel training. It sums a target array (like gradients) across all processes in a group and distributes the result back to all processes.

• ZeRO's Impact: Standard data parallelism requires a costly all-reduce on gradients after every backward pass. ZeRO stages 1 and 2 optimize this:
  • ZeRO-1: Uses reduce-scatter for gradients (more efficient than all-reduce).
  • ZeRO-2: Also partitions gradients, further reducing communication volume.
• Bandwidth Cost: All-reduce communication is often the performance bottleneck for data parallelism at scale. ZeRO's optimizations directly address this by reducing the size of the communicated data.