ZeRO (Zero Redundancy Optimizer)

Optimizer states (e.g., momentum, variance for Adam) are partitioned across data-parallel processes. Each process stores and updates only its assigned shard, reducing memory footprint by ~4x.

• Mechanism: The optimizer is sharded. During the backward pass, each process gathers the necessary optimizer states for its local parameter update.
• Impact: Enables training of models ~4x larger than standard data parallelism, as optimizer states are a major memory consumer for adaptive optimizers like Adam.

Gradients are partitioned across processes in addition to optimizer states. After the backward pass, each process only retains the gradients for its assigned parameter partition.

• Mechanism: Uses a reduce-scatter operation instead of an all-reduce. Each GPU becomes responsible for a unique subset of gradients.
• Impact: Provides an additional ~2x memory savings over Stage 1. Total memory reduction is ~8x compared to standard data parallelism, enabling even larger models.

The model parameters themselves are partitioned across GPUs. Each process only stores the parameters for its assigned partition, fetching others on-demand during forward and backward passes.

• Mechanism: Parameters are gathered just-in-time for computation and discarded afterward. This introduces communication overhead but maximizes memory savings.
• Impact: Enables memory reduction linear with the degree of data parallelism (N). In theory, a model of size M can be trained across N GPUs with ~M/N memory per GPU.

ZeRO trades increased communication for reduced memory. The overhead is managed through optimized collective operations.

• Stage 1/2: Use reduce-scatter and all-gather operations. Overhead is moderate and often offset by the ability to use larger batch sizes.
• Stage 3: Requires frequent all-gather and reduce-scatter for every layer, increasing communication volume. This is mitigated by overlapping communication with computation (communication hiding).

ZeRO is implemented in progressive stages, each eliminating a different type of memory redundancy.

• ZeRO-1 (Optimizer State Partitioning): Partitions the optimizer states (e.g., momentum, variance) across processes, reducing memory by the number of data-parallel workers (D).
• ZeRO-2 (Add Gradient Partitioning): Additionally partitions gradients across processes, with each process only updating the parameters for which it holds the gradient.
• ZeRO-3 (Full Parameter Partitioning): Partitions the model parameters themselves. Parameters are gathered on-demand for the forward and backward passes, enabling the training of models larger than the memory of a single GPU.

Each stage builds upon the last, with ZeRO-3 offering the highest memory savings at the cost of increased communication overhead.

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

• ZeRO-Offload: Enables training of large models on a single GPU by offloading optimizer states and gradients to CPU memory. It strategically keeps parameters on GPU for compute, using the CPU as a massive memory buffer.
• ZeRO-Infinity: Goes further by offloading to CPU and NVMe storage. It uses innovative techniques like infinity offload engine and bandwidth-centric partitioning to efficiently leverage terabytes of CPU/NVMe memory, enabling training of trillion-parameter models.
• Use Case: Critical for research and organizations without massive GPU clusters, democratizing access to large-scale model training.

Choosing a ZeRO implementation involves trade-offs between memory efficiency, communication cost, and implementation complexity.

• Memory vs. Communication: ZeRO-1 saves memory with low overhead. ZeRO-3 saves the most memory but introduces significant communication for parameter gathering (all-gather operations).
• Framework Choice:
  • DeepSpeed: Most feature-rich and battle-tested for extreme scale. Higher configuration complexity.
  • FSDP: PyTorch-native, easier debugging, and integrates with PyTorch's future roadmap. May have different performance characteristics.
  • Hugging Face: Best for rapid prototyping and using pre-existing model code.
• Hybrid Parallelism: In practice, ZeRO (data parallelism) is often combined with Tensor Parallelism (intra-layer) and Pipeline Parallelism (inter-layer) to train the largest models.

Mixed precision training uses lower numerical precision (e.g., BF16 or FP16) for most operations to accelerate computation and reduce memory usage, while maintaining higher precision (FP32) for a master copy of weights and critical operations to ensure numerical stability.

• Key Components:
  • FP16/BF16 for Forward/Backward Pass: Faster computation and halved memory for activations and gradients.
  • FP32 Master Weights & Optimizer States: Maintains a full-precision copy of weights to avoid underflow/overflow during gradient updates.
  • Loss Scaling: Scales the loss to prevent gradient values from vanishing in the lower-precision range.
• Integration with ZeRO: ZeRO's optimizer state partitioning is especially powerful here, as the FP32 master states are a major memory bottleneck. Combining mixed precision with ZeRO-2 or ZeRO-3 is standard for modern large model training.

Model parallelism is a distributed training strategy that partitions a single model's layers or components across multiple devices. This contrasts with data parallelism (which ZeRO enhances) where the model is replicated and data is split.

• Tensor Parallelism: Splits individual layer operations (e.g., matrix multiplications within an MLP or attention head) across GPUs. Used in models like Megatron-LM.
• Pipeline Parallelism: Places different groups of model layers on different GPUs. Micro-batching is used to keep the pipeline full and avoid idle devices.
• Relationship to ZeRO: ZeRO is a form of optimized data parallelism. In practice, the largest models use 3D parallelism: a combination of ZeRO/FSDP (Data Parallelism), Tensor Parallelism, and Pipeline Parallelism to distribute both the model and the data.

ZeRO operates in progressive optimization stages, each eliminating a different type of memory redundancy.

• ZeRO-1 (Optimizer State Partitioning): Partitions only the optimizer states (e.g., momentum, variance) across processes. Reduces memory proportional to the data parallelism degree.
• ZeRO-2 (Add Gradient Partitioning): Adds partitioning of gradients in addition to optimizer states. Further reduces memory and enables larger batch sizes.
• ZeRO-3 (Add Parameter Partitioning): Partitions the model parameters themselves across processes. This is the most memory-efficient stage, enabling model sizes that exceed the memory of any single GPU. Parameters are gathered on-demand for computation.
• ZeRO-Offload & ZeRO-Infinity: Advanced variants that offload partitioned states to CPU RAM (Offload) or NVMe storage (Infinity), pushing the boundaries of trainable model size.

Distributed Data Parallel is the classic, non-sharded data parallelism framework in PyTorch. It serves as the baseline that ZeRO optimizes.

• Baseline Mechanism: Each GPU holds a full replica of the model, optimizer, and its states. Each processes a different shard of the data batch.
  • After the backward pass, gradients are averaged across all processes via an all-reduce operation.
  • Each GPU then performs an identical weight update.
• The Redundancy Problem: This full replication is the source of memory redundancy. Training a 10B parameter model requires 10B parameters * (1 + 2 + 12) bytes ≈ 40GB per GPU just for parameters, gradients, and Adam optimizer states.
• ZeRO as an Optimization: ZeRO can be viewed as a memory-optimized version of DDP. It removes the replication of optimizer states, gradients, and parameters, transforming DDP from a memory-redundant to a memory-efficient paradigm.