ZeRO Optimization

ZeRO-1 eliminates memory redundancy for the optimizer states. In standard data parallelism, each GPU holds a full copy of the optimizer states (e.g., momentum and variance for Adam), which can consume ~12 bytes per parameter.

• Mechanism: Partitions the optimizer states across the data-parallel processes. Each process is responsible for updating only its assigned partition.
• Memory Reduction: Reduces optimizer state memory by a factor equal to the data-parallel degree (e.g., 8x reduction with 8 GPUs).
• Communication: Requires a single all-gather operation before the weight update to collect the relevant optimizer partitions, but no extra communication during the backward pass.
• Use Case: The first step, enabling larger models than standard data parallelism but not the full memory savings of later stages.

ZeRO-2 builds on ZeRO-1 by also partitioning the gradients. In standard training, each GPU stores the full gradients for all parameters after the backward pass.

• Mechanism: Gradients are partitioned across processes immediately after they are computed during the backward pass. Each process only stores and updates the gradients for its assigned parameter partition.
• Memory Reduction: Reduces gradient memory by a factor equal to the data-parallel degree, on top of the optimizer state savings from ZeRO-1.
• Communication: Uses a reduce-scatter operation during the backward pass to aggregate gradients across partitions, ensuring correct averaging. This replaces the standard all-reduce used in data parallelism.
• Impact: Often provides the most significant memory savings relative to its communication overhead, making it a highly efficient stage for scaling.

ZeRO-3 is the most comprehensive stage, partitioning the model parameters themselves across GPUs. This eliminates all major memory redundancies.

• Mechanism: The full model parameters are partitioned across the data-parallel processes. Each GPU only stores the parameters for its assigned partition.
• Memory Reduction: Reduces model parameter memory by a factor equal to the data-parallel degree. Combined with ZeRO-1 and ZeRO-2, it enables linear memory scaling with the number of GPUs.
• Communication Overhead: Requires sophisticated communication scheduling:
  • Forward Pass: Parameters needed for a given layer are all-gathered just before computation and discarded afterward.
  • Backward Pass: Parameters are again gathered for gradient calculation, followed by a reduce-scatter for the gradients.
• Use Case: Essential for training models with trillions of parameters, as it allows the aggregate GPU memory of a cluster to define the trainable model size.

Each ZeRO stage introduces distinct communication operations that replace those of standard data parallelism.

• Baseline (DP): Uses a single all-reduce operation after backward pass to average gradients.
• ZeRO-1: Adds an all-gather before the optimizer step (for optimizer states).
• ZeRO-2: Replaces the all-reduce with a reduce-scatter during backward pass (for gradients).
• ZeRO-3: Adds multiple all-gather operations (before each layer in forward/backward) and a reduce-scatter (backward).
• Overhead Management: The communication volume remains similar to baseline data parallelism, but the number of operations increases. This overhead is managed in DeepSpeed via optimized kernels and overlapping communication with computation.

When adapting a large pre-trained model (e.g., Llama 2 70B) to a new domain via full fine-tuning, all model parameters are updated. This process requires storing the model weights, gradients, optimizer states (like momentum and variance for Adam), and activations—a massive memory footprint. ZeRO stages 2 and 3 enable this by distributing these components. For example, ZeRO-2 partitions gradients and optimizer states, while ZeRO-3 also partitions the model parameters themselves. This makes full fine-tuning feasible on a modest cluster of GPUs, allowing for deeper specialization than parameter-efficient fine-tuning (PEFT) methods like LoRA.

Training on long sequences (e.g., 32k+ tokens) is critical for document understanding, code generation, and long-context reasoning. The memory for activations (intermediate layer outputs) scales linearly with sequence length and batch size, quickly becoming the dominant memory consumer. While ZeRO primarily targets optimizer and parameter memory, its memory savings free up GPU resources to store these larger activations. Furthermore, ZeRO can be combined with activation checkpointing (or gradient checkpointing), which recomputes activations during the backward pass instead of storing them, enabling training with sequences that would be otherwise impossible.

Large batch sizes are often necessary for stable training, especially in distributed settings, and to maximize GPU utilization. However, per-GPU memory limits batch size. By drastically reducing the memory footprint of the model states, ZeRO allows for significantly larger effective batch sizes per GPU. Users can either increase the micro-batch size per GPU for better hardware efficiency or scale out to more GPUs with data parallelism while maintaining a manageable memory profile. This is crucial for data-intensive tasks like contrastive learning or when using massive datasets.

Instruction tuning on large, diverse datasets (e.g., mixtures of Alpaca, ShareGPT, and proprietary data) involves processing many distinct task formats. Efficiently handling this requires flexible data loading and often larger model capacities. ZeRO's memory efficiency allows practitioners to use larger base models for instruction tuning, leading to more capable and generalist models. It also supports the full fine-tuning phase often used after initial supervised fine-tuning (SFT) in alignment pipelines like RLHF, where both the policy model and the reward model can be large.

Gradient checkpointing is a memory optimization technique that trades compute for memory by selectively saving only a subset of the intermediate activations during the forward pass of neural network training. During the backward pass, the unsaved activations are recalculated on-demand. This enables the training of models with significantly larger layers or longer sequence lengths than would fit in GPU memory otherwise.

• Core Mechanism: Instead of storing all activations (O(n) memory), it stores only checkpoints (O(√n) memory).
• Trade-off: Reduces memory consumption by up to 60-70% at the cost of approximately a 20-30% increase in computation time due to recomputation.
• Use Case: Essential for training very deep networks, such as large vision transformers or language models with extremely long context windows, and is often used in conjunction with ZeRO.

Mixed precision training is a computational technique that uses lower-precision numerical formats (like 16-bit floating-point - FP16 or Brain Floating Point - BF16) for most operations to accelerate computation and reduce memory usage, while maintaining higher precision (32-bit FP) for critical operations to preserve model stability and accuracy.

• Performance Gains: Using FP16/BF16 halves the memory footprint of tensors and can provide up to 3x speedup on modern NVIDIA Tensor Core GPUs.
• Precision Management: A master copy of weights is kept in FP32. During the forward and backward passes, weights are cast to FP16/BF16. Gradients are computed in lower precision, then upcast to FP32 for the weight update.
• Loss Scaling: To prevent underflow of small gradient values, gradients are scaled up before the backward pass and scaled down before the weight update.

Parameter-Efficient Fine-Tuning (PEFT) refers to a family of techniques designed to adapt large pre-trained models to downstream tasks by updating only a small, targeted subset of the model's parameters. This stands in contrast to full fine-tuning, which updates all billions of parameters.

• Primary Goal: Drastically reduce the computational, storage, and memory overhead of fine-tuning.
• Common Methods: Includes LoRA (Low-Rank Adaptation), Adapter layers, and prefix tuning.
• Key Benefit: Enables task adaptation with minimal risk of catastrophic forgetting of the model's original knowledge. PEFT is orthogonal to ZeRO; ZeRO optimizes memory during the training process, while PEFT reduces the number of parameters that need to be trained.

Model parallelism is a distributed training strategy where the computational graph of a single neural network model is partitioned across multiple devices (GPUs). This is necessary when a model is too large to fit on a single device's memory.

• Types: Includes tensor parallelism (splitting individual layers, e.g., the attention heads in a transformer, across devices) and pipeline parallelism (splitting the model into sequential stages of layers).
• Relationship to ZeRO: ZeRO is a form of data parallelism that eliminates memory redundancy. In practice, ZeRO is often combined with model parallelism (e.g., in 3D parallelism) to train the world's largest models. ZeRO handles the data-parallel dimension efficiently, while tensor/pipeline parallelism handle the model's size.

Full fine-tuning is the conventional process of adapting a pre-trained model by updating all of its parameters during supervised training on a new dataset. This is the baseline approach that optimization techniques like ZeRO and PEFT aim to improve upon or enable.

• Process: The entire pre-trained model is loaded, and every weight receives gradient updates based on a task-specific loss function (e.g., cross-entropy loss).
• Drawbacks: It is computationally expensive, requires significant GPU memory (addressed by ZeRO), and carries a high risk of catastrophic forgetting.
• Contrast: Unlike Parameter-Efficient Fine-Tuning (PEFT) methods that freeze most weights, full fine-tuning offers maximum flexibility for task adaptation but at the highest cost, making memory optimization critical.