ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Samyam Rajbhandari; Jeff Rasley; Olatunji Ruwase; Yuxiong He

Source-linked AI summary

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He

arXiv:1910.02054v3Oct 4, 2019cs.LGcs.DCstat.ML

TL;DR

Training models from tens of billions to trillions of parameters exceeds single-device memory and is not solved simply by adding devices. ZeRO reduces redundant memory in data- and model-parallel training, enabling substantially larger models and faster training, including trillion-parameter scale in aggregate memory.

Problem
Models are growing because larger size offers significant accuracy gains, but their parameters, optimizer states, gradients, activations, and buffers exceed available device memory.
Method
ZeRO optimizes memory by eliminating redundancies in model states and partitioning parameters, gradients, and activations across devices while retaining communication and computation efficiency.
Results
ZeRO-100B efficiently runs 170B-parameter models, delivers over 15 Petaflops and more than 10x training-speed improvement over SOTA, and supports models up to 13B without model parallelism.
Takeaways & Limitations
With all three stages enabled, ZeRO can train a trillion-parameter model on 1024 NVIDIA GPUs and has potential to scale beyond 1 trillion parameters using today’s hardware.
Takeaways & Limitations
Training a 1T-parameter model could still take over a year in practice because data samples and sequence lengths are likely to increase with model size.

Abstract

from arXiv · show

Large deep learning models offer significant accuracy gains, but training billions to trillions of parameters is challenging. Existing solutions such as data and model parallelisms exhibit fundamental limitations to fit these models into limited device memory, while obtaining computation, communication and development efficiency. We develop a novel solution, Zero Redundancy Optimizer (ZeRO), to optimize memory, vastly improving training speed while increasing the model size that can be efficiently trained. ZeRO eliminates memory redundancies in data- and model-parallel training while retaining low communication volume and high computational granularity, allowing us to scale the model size proportional to the number of devices with sustained high efficiency. Our analysis on memory requirements and communication volume demonstrates: ZeRO has the potential to scale beyond 1 Trillion parameters using today's hardware. We implement and evaluate ZeRO: it trains large models of over 100B parameter with super-linear speedup on 400 GPUs, achieving throughput of 15 Petaflops. This represents an 8x increase in model size and 10x increase in achievable performance over state-of-the-art. In terms of usability, ZeRO can train large models of up to 13B parameters (e.g., larger than Megatron GPT 8.3B and T5 11B) without requiring model parallelism which is harder for scientists to apply. Last but not the least, researchers have used the system breakthroughs of ZeRO to create the world's largest language model (Turing-NLG, 17B parameters) with record breaking accuracy.

1 Extended Introduction

ZeRO addresses the memory and scaling limits of large-model training by partitioning redundant model states and optimizing residual memory. It enables substantially larger models, higher throughput, and broader usability than existing approaches while retaining efficient training.

Motivation: Large models exceed single-device memory, and simply adding devices does not solve the training-scale challenge.Existing approaches also face trade-offs among memory, computation, communication, and usability.
ZeRO-DP: ZeRO partitions optimizer states, gradients, and parameters across data-parallel processes to eliminate replicated model-state memory.Its three cumulative stages provide 4x, 8x, and up to 64x memory reductions, with parameter partitioning adding a modest 50% communication increase.
ZeRO-R: ZeRO-R reduces residual memory by partitioning or offloading activations, sizing temporary buffers, and preventing fragmentation.These optimizations target memory consumed by activations, temporary buffers, and unusable fragments after model-state memory is optimized.
System design: ZeRO combines ZeRO-DP and ZeRO-R, and can also combine with model parallelism for a theoretical Nd × Nm per-device memory reduction.This combination could fit a trillion-parameter model on 1024 GPUs using 16-way model parallelism and 64-way data parallelism.
Results: ZeRO exhibits super-linear speedup from 64 to 400 GPUs and supports training models up to 13B parameters without model or pipeline parallelism.The system also powered a 17B-parameter Turing-NLG model with record-breaking accuracy.

2 Related Work

Related work scales large-model training through data, model, and pipeline parallelism, alongside activation, optimizer-state, and CPU-offloading techniques. These approaches trade memory reduction against computation, communication, usability, or convergence considerations.

Parallel training: Data parallelism replicates model parameters across devices, while model and pipeline parallelism split models vertically or horizontally when they do not fit.Pipeline parallelism uses micro-batching to hide pipeline bubbles, but horizontal splitting complicates tied weights and batch normalization.
Memory reduction: Activation compression, checkpointing, and live analysis reduce activation memory and can operate alongside ZeRO, including ZeRO-R activation reduction.These methods are described as complementary rather than mutually exclusive.
Memory reduction: CPU-offloading methods reduce memory through heterogeneous memory but may spend up to 50% of training time on GPU-CPU-GPU transfers.ZeRO instead reduces memory without generally storing model states in bandwidth-constrained CPU memory.
Optimizer memory: ZeRO is orthogonal to coarser adaptive-optimization statistics and reduces optimizer-state and gradient memory without changing the optimization method or convergence.Adaptive optimizers remain important for large-model training despite their fine-grained state memory costs.

3 Where Did All the Memory Go?

Training memory is dominated by model states—optimizer states, gradients, and parameters—but activations, temporary buffers, and fragmentation create substantial residual demands. Mixed-precision Adam and large-model activations illustrate why parameter count alone understates training memory.

Memory components: Most training memory consists of model states, while activations, temporary buffers, and fragmented memory form residual states.The model examines both categories to explain why a 1.5B-parameter model cannot be trained on a single 32GB GPU despite 3GB of 16-bit weights.
Model states: Adam stores momentum and variance in addition to parameters and gradients, making optimizer states a major memory component.The optimizer maintains fine-grained first- and second-order statistics for each parameter and gradient.
Model states: 16Ψ bytes is the mixed-precision Adam memory requirement for a model with Ψ parameters.This comprises fp16 parameters and gradients plus fp32 parameters, momentum, and variance.
Model states: 24 GB is the minimum mixed-precision Adam memory requirement for GPT-2 with 1.5 billion parameters.The estimate follows directly from the 16Ψ-byte requirement.
Residual states: 60 GB is required by 1.5B GPT-2 with sequence length 1K and batch size 32 for activations, while a 100B GPT-like model needs around 60 GB even with checkpointing.Checkpointing reduces activation memory by approximately the square root of total activations but adds 33% recomputation overhead.
Residual states: Temporary fused buffers and memory fragmentation can trigger out-of-memory failures beyond the memory required by model states and activations.Fragmentation can leave insufficient contiguous memory even when total available memory exceeds a request.

4 ZeRO: Insights and Overview

ZeRO combines ZeRO-DP for model-state memory with ZeRO-R for residual memory. Its design partitions redundant states and uses dynamic communication and memory-management strategies to preserve data-parallel efficiency.

Overview: ZeRO has two optimization families: ZeRO-DP reduces model-state memory, while ZeRO-R reduces residual memory from activations, buffers, and fragmentation.The design explicitly retains efficiency rather than relying on unconstrained CPU offloading or arbitrarily high model-parallelism.
ZeRO-DP: ZeRO-DP combines data parallelism’s computational granularity and communication efficiency with model parallelism’s state partitioning.It partitions model states and schedules communication dynamically because different states are needed at different times.
ZeRO-DP: ZeRO-DP reduces per-device model-state memory linearly as the data-parallel degree increases.The approach partitions states instead of replicating them while minimizing communication volume through a dynamic schedule.
ZeRO-R: ZeRO-R partitions activation checkpoints across GPUs and reconstructs them on demand with all-gather, reducing activation memory proportional to model-parallel degree.For very large models, activation partitions may be moved to CPU memory while maintaining efficiency due to high arithmetic intensity.
ZeRO-R: ZeRO-R uses constant-size buffers and on-the-fly defragmentation to control temporary-buffer growth and improve contiguous memory availability.Defragmentation moves activation checkpoints and gradients into pre-allocated contiguous buffers.

5 Deep Dive into ZeRO-DP

ZeRO-DP progressively partitions optimizer states, gradients, and parameters across data-parallel processes. The stages reduce per-device memory while adding limited communication, enabling substantially larger models as the data-parallel degree grows.

Partitioning stages: ZeRO-DP removes replicated model-state memory by partitioning optimizer states, gradients, and parameters across data-parallel processes.The three cumulative stages are Pos, Pos+g, and Pos+g+p.
Optimizer-state partitioning: Each process updates 1/Nd of optimizer states and parameters, then all-gathers fully updated parameters after each training step.The optimizer-state partition determines which process performs each corresponding update.
Optimizer-state partitioning: 31.4GB is required by a 7.5B-parameter model with Pos and 64-way DP, versus 120GB with standard DP.Optimizer-state partitioning reduces model-state memory from 4Ψ + KΨ to approximately (4Ψ + KΨ)/Nd.
Parameter partitioning: 1.9GB of model-state memory is required by a 7.5B-parameter model with Pos+g+p and 64-way DP, versus 120GB with standard DP.Parameter partitioning reduces model-state memory from 16Ψ to 16Ψ/Nd and increases baseline communication volume to 1.5x.
Implications for model size: With Nd = 64, ZeRO fits models up to 7.5B, 14B, and 128B parameters using Pos, Pos+g, and Pos+g+p, respectively.With Nd = 1024 and all optimizations enabled, the paper reports capacity for 1 Trillion parameters.

6 Deep Dive into ZeRO-R

ZeRO-R reduces activation and fragmentation memory in model-parallel training by partitioning activations and managing contiguous buffers. These optimizations substantially lower activation requirements while balancing memory use against computational efficiency.

Activation Partitioning: ZeRO-R partitions activations across model-parallel GPUs and materializes them in replicated form only when needed for computation.This removes replicated activation copies while preserving the activation data required by each layer.
Activation Partitioning: 33 GB of activation checkpoints per GPU falls to about 2 GB for a 100B model with batch size 32, sequence length 1024, and MP degree 16.The reduction comes from partitioned activation checkpointing with Pa.
Activation Partitioning: Partitioned activation checkpointing reduces activation footprint by a factor proportional to the model-parallel degree.For the cited 100B-model configuration, the MP degree is 16.
Buffer Management: ZeRO selects temporal-data buffer sizes to balance memory efficiency with operation efficiency, since larger inputs can improve communication-library bandwidth.The design accounts for operations such as fused parameter communication and large all-reduce operations.
Memory Defragmentation: ZeRO defragments memory online by pre-allocating contiguous chunks for activation checkpoints and gradients before copying produced data into them.This supports larger models and batch sizes under memory constraints while improving efficiency.

7 Communication Analysis of ZeRO-DP

ZeRO-DP removes replicated optimizer states, gradients, and parameters by partitioning them across data-parallel processes. It preserves communication efficiency through optimizer-state and gradient partitioning, while full partitioning trades additional communication for further memory reduction.

Partitioning Strategy: ZeRO-DP partitions optimizer states, gradients, and parameters across data-parallel processes instead of replicating them on every device.These are the three cumulative optimization phases of ZeRO-DP.
Communication Trade-offs: No additional communication is incurred with optimizer-state and gradient partitioning, while memory reduction reaches up to 8x.The communication comparison is against baseline data parallelism.
Communication Trade-offs: Gradient partitioning replaces all-reduce with scatter-reduce followed by parameter all-gather after each process updates its assigned parameter partition.Parameter broadcasts during forward propagation can be pipelined to limit memory overhead.
Communication Trade-offs: 3Ψ total communication volume results from gradient reduce-scatter plus parameter all-gathers, or 1.5x the baseline volume.The parameter all-gathers are scheduled across forward and backward propagation.

8 Communication Analysis of ZeRO-R

ZeRO-R adds activation-checkpoint partitioning to model-parallel training and analyzes when its communication cost is worthwhile. The method can substantially reduce data-parallel communication, while CPU offloading nearly eliminates activation memory at additional data-movement cost.

Communication Trade-off: Pa trades activation-memory reduction for additional model-parallel communication, with the trade-off depending on model size, checkpointing strategy, and MP strategy.The analysis is conducted for transformer models implemented with Megatron-LM.
Communication Overhead: Partitioned activation checkpointing generally adds less than 10% of baseline model-parallel communication.For transformer blocks, the additional all-gather communication is one sequence-length-by-hidden-dimension message per checkpointed block.
Communication Trade-off: An order-of-magnitude increase in batch size from Pa can reduce data-parallel communication by an order of magnitude when that communication bottlenecks performance.The activation-memory reduction scales with the model-parallel degree, allowing larger batches.
CPU Offloading: Pa+cpu reduces activation memory requirements to nearly zero but adds 2x the CPU data movement of Pa.It is useful when data-parallel communication dominates and CPU transfer overhead remains lower than that communication overhead.

9 Step Towards 1 Trillion Parameters

ZeRO targets the memory and communication barriers that limit trillion-parameter training on current hardware. Its measured scaling results show much larger runnable models, while end-to-end training time at trillion scale remains constrained by available compute power.

Existing Scaling Limits: Megatron reaches acceptable throughput at 16–20B parameters on one DGX-2, but scaling model parallelism across DGX nodes causes significant efficiency loss.The stated cause is limited internode bandwidth.
Memory Scaling: More than 1T parameters can fit on 1024 GPUs with ZeRO using data parallelism alone or 16-way model parallelism within nodes and 64-way data parallelism across nodes.The combined configuration is presented as an alternative to data parallelism alone.
Compute Power Gap: A trillion-parameter model could require 140 days under Bert-Large-like assumptions and over a year when data and sequence length also increase.The paper states that an exa-flop system would be required to train such a model in a reasonable time.

10 Implementation and Evaluation

ZeRO-100B demonstrates that memory optimizations can support substantially larger models and high throughput on GPU clusters, while preserving usability without model changes. Its evaluations cover scalability, model size, memory, performance, and a 17B-parameter language model.

Speed and Model Size: 170B parameters run on 400 GPUs, more than 8x bigger than Megatron-LM, with up to 10x speedup over baseline.ZeRO-100B sustains 15 PetaFlops on average for 8B–100B models.
Scalability: ZeRO-100B achieves super-linear scalability for a 60B-parameter model as the data-parallel degree increases.Parameter and gradient partitioning reduce per-GPU memory, enabling larger batch sizes and higher arithmetic intensity.
Usability: 13B parameters can be trained without model parallelism on 128 GPUs, exceeding 40 TFlops per GPU on average.ZeRO requires no model changes and can be used similarly to baseline data parallelism.
Memory and Performance Analysis: 170B-parameter training requires Pa+cpu to execute without running out of memory, while CPU activation movement can reduce performance.The implementation enables Pa+cpu only when it is beneficial, particularly when models otherwise cannot run or batch sizes are very small.
Turing-NLG: Turing-NLG exceeds 17B parameters, reaches 10.21 Webtext-103 perplexity, and sustains 41.4 TFlops/GPU when trained with ZeRO-100B.The reported perplexity is compared with the previous state-of-the-art Megatron-LM 8.3B model.

11 Concluding Remarks

The authors position ZeRO as a scalable approach for large-model training and emphasize its usability for data scientists. They report substantial gains while identifying broader potential beyond the evaluated ZeRO-100B implementation.

Concluding Remarks: ZeRO-100B enables an 8x increase in model size, over 10x throughput improvement, and super-linear speedups on modern GPU clusters.The authors describe these results as only a portion of ZeRO’s broader potential.
Concluding Remarks: The complete ZeRO approach is projected to increase trainable model size by another order of magnitude beyond ZeRO-100B.This is presented as potential rather than as an evaluated result of the implementation described.
Concluding Remarks: ZeRO requires no model refactoring and is as easy to use as standard data parallelism, unlike model and pipeline parallelism.The authors identify this usability as a reason for future large-model-training investigations.