ZeRO (Zero Redundancy Optimizer)

ZeRO Stage 1 partitions only the optimizer states (e.g., momentum, variance for Adam) across the data parallel processes. Each GPU stores and updates only its assigned partition of the optimizer states, reducing the memory footprint of the optimizer by a factor of the data parallel degree (N). This is the most basic form of ZeRO, offering significant memory savings with minimal communication overhead.

• Example: With 8 GPUs (N=8), the memory for Adam optimizer states is reduced by 8x.
• Benefit: Enables training of models 8x larger with the same GPU memory, or allows the use of larger batch sizes.

ZeRO Stage 2 builds on Stage 1 by also partitioning the gradients across processes. After the backward pass, each GPU only retains the gradients corresponding to its partition of the parameters. A global all-gather operation is required to collect the full gradients before the optimizer step, but they are immediately discarded after use.

• Mechanism: Eliminates the memory redundancy of storing a full copy of gradients on every GPU.
• Benefit: Further increases model size capacity. Combined with Stage 1, memory reduction scales with N for both optimizer states and gradients.

ZeRO Stage 3 is the most memory-efficient stage, partitioning the model parameters themselves across GPUs. Each GPU stores only a fraction of the parameters. During the forward and backward passes, parameters are gathered on-demand for each layer via communication collectives and then released, a process known as the ZeRO-3 forward/backward pass.

• Key Concept: Parameter offloading can be combined with Stage 3, where inactive parameter partitions are moved to CPU RAM, enabling training of models larger than the aggregate GPU memory.
• Trade-off: Maximizes model size capability but introduces the highest communication overhead, requiring careful tuning for performance.

ZeRO's memory savings come with increased inter-GPU communication, which is managed through optimized collective operations.

• All-Gather: Used in Stage 3 to collect parameters for a layer and in Stage 2 to collect gradients. This is a bandwidth-intensive operation.
• Reduce-Scatter: Used to sum and partition gradients across GPUs during the backward pass.
• Overhead Management: The communication overhead is often hidden by overlapping it with computation. Frameworks like DeepSpeed implement sophisticated communication scheduling to minimize the performance impact.

ZeRO is fundamentally a data parallelism technique. It is most powerful when combined with other forms of parallelism to train extremely large models.

• ZeRO + Tensor/Pipeline Parallelism: A common 3D parallelism strategy. ZeRO handles data parallelism across model replicas, while tensor and pipeline parallelism split the model itself across GPUs. This combination allows scaling to thousands of GPUs.
• Role in Training: ZeRO efficiently manages the memory footprint of the data parallel dimension, while model parallelism techniques manage the memory and compute of the model's layers.

The core ZeRO technique is implemented in three progressive stages, each partitioning more model state to save more memory:

• ZeRO Stage 1 (Optimizer State Partitioning): Partitions the optimizer states (e.g., momentum, variance for Adam) across processes. Reduces memory by ~4x.
• ZeRO Stage 2 (Gradient Partitioning): Adds partitioning of gradients after the backward pass. Reduces memory by ~8x.
• ZeRO Stage 3 (Parameter Partitioning): Partitions the model parameters across processes, fetching them via all-gather only when needed for computation. Enables training of models with trillions of parameters.

These extensions push memory savings beyond GPU RAM limits by leveraging CPU and storage.

• ZeRO-Offload: Efficiently offloads optimizer states and gradients to CPU memory while keeping parameters on GPU. Enables training models 10x larger on a single GPU.
• ZeRO-Infinity: The most advanced stage, offloading all ZeRO components—optimizer states, gradients, parameters—to both CPU and NVMe solid-state storage. This allows for training models with tens of trillions of parameters by leveraging terabytes of CPU/NVMe memory.

ZeRO is accessible to a broad audience through tight integration with the Hugging Face Transformers and Accelerate libraries.

• The Trainer class in Transformers natively supports DeepSpeed via a JSON configuration file.
• Hugging Face Accelerate provides a unified interface for distributed training, supporting ZeRO configurations through its DeepSpeedPlugin.
• This integration allows users to apply ZeRO optimizations with minimal code changes to existing training scripts.

FSDP is PyTorch's native implementation of the ZeRO-3 concept, integrated into the torch.distributed module. It shards model parameters, gradients, and optimizer states across data-parallel processes.

Key comparisons with DeepSpeed ZeRO-3:

• Tighter PyTorch Integration: Uses PyTorch's native communication primitives.
• Flexible Sharding: Can shard parameters at the layer level for more efficient communication.
• Simplified Configuration: Managed via wrapping model layers with FullyShardedDataParallel. Often preferred for new PyTorch-centric projects.

A family of techniques for splitting a neural network across multiple devices (GPUs) when it is too large to fit on a single device. It is the foundational concept upon which ZeRO builds. There are two primary strategies:

• Tensor Parallelism: Splits individual layers (e.g., the weights of a linear layer) across devices, requiring frequent communication (all-reduce) during each forward/backward pass.
• Pipeline Parallelism: Places different groups of layers (stages) on different devices, passing activations sequentially like an assembly line, often using micro-batching to keep all devices busy. ZeRO is often categorized as a form of optimized data parallelism that incorporates model parallelism ideas for optimizer states, gradients, and parameters.

The standard distributed training approach where the entire model is replicated on each GPU. A large batch of training data is split into smaller mini-batches, each processed independently on a different GPU. Gradients are then synchronized across all GPUs by averaging them (using an all-reduce operation) before the optimizer step updates the weights. While simple, its major limitation is memory redundancy—each GPU holds a full copy of the model, optimizer states, gradients, and activations. ZeRO directly addresses this redundancy by partitioning these components across the data parallel processes.

A memory optimization technique that trades compute for memory. During the forward pass, only a subset of layer activations are stored (the 'checkpoints'). During the backward pass, the missing activations are recomputed on-the-fly from the nearest checkpoint. This can reduce activation memory by up to 5-10x but increases computation time by ~20-30%. It is complementary to ZeRO; they are often used together. While ZeRO optimizes memory for parameters, gradients, and optimizer states, gradient checkpointing targets the memory consumed by activations, enabling the training of models with even deeper architectures.

The use of lower-precision numerical formats (like FP16 or BF16) for parts of the training process to speed up computation and reduce memory usage. A typical scheme uses FP16 for forward/backward passes and FP32 for the master weight copy and optimizer states to maintain stability. ZeRO's memory savings are particularly impactful here, as the optimizer states (e.g., momentum, variance in Adam) are often kept in FP32, becoming a major memory bottleneck. ZeRO-Offload, a stage of ZeRO, can move these FP32 optimizer states to CPU memory, further leveraging mixed precision's characteristics.

PyTorch's native implementation of the ZeRO-3 algorithm, introduced in PyTorch 1.11. It shards model parameters, gradients, and optimizer states across data parallel workers. Key features include:

• Automatic wrapping: Simplifies integration by allowing users to specify which sub-modules to shard.
• Flexible precision policies: Native support for mixed precision training.
• Integration with PyTorch ecosystem: Works seamlessly with torch.compile, DistributedDataParallel (DDP) APIs, and other PyTorch components. FSDP has become the standard method for training large models within the PyTorch framework, providing a user-friendly API for ZeRO's capabilities.