[Paper Review] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Samyam Rajbhandari, Jeff Rasley+2 more|arXiv (Cornell University)|Oct 4, 2019

Advanced Neural Network Applications21 references72 citations

TL;DR

ZeRO introduces memory-efficient partitioning of optimizer states, gradients, and parameters (plus residual memory optimizations) to enable training models up to trillion parameters with data-parallelism, achieving large-speedups and enabling democratized large-model training.

ABSTRACT

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.

Motivation & Objective

• Motivate the need to train trillion-parameter models and identify memory bottlenecks in existing data/model/pipeline parallelism.
• Propose ZeRO to eliminate memory redundancies in model states and residuals while preserving efficiency.
• Develop and analyze three partitioning stages (optimizer states, gradients, parameters) and residual-memory optimizations to scale model size with available hardware.
• Demonstrate practical training improvements, show integration with MP, and provide open-source tooling to democratize large-model training.

Proposed method

• Partition optimizer states across data-parallel processes (P_os) so each rank holds 1/N_d of states.
• Partition gradients corresponding to each parameter partition (P_g) and perform bucketized reduce-scatter at partition boundaries.
• Partition parameters across data-parallel processes (P_p) with on-demand broadcast of non-local parameters to forward/backward passes.
• Combine ZeRO-DP with ZeRO-R to also partition activations and manage residual memory (checkpointing, buffers, and fragmentation) including optional CPU offloading (P_a, P_a+cpu).
• Use a dynamic communication schedule to minimize communication while decoupling memory savings from compute; analyze communication overhead (max 1.5x with P_p).
• Support combination with model parallelism (MP) when beneficial, but show ZeRO-DP can match or exceed MP in many cases and enable training without extensive model refactors.

Experimental results

Research questions

• RQ1 Can ZeRO partition model states to reduce per-device memory by large factors without sacrificing DP efficiency?
• RQ2 How much memory can be saved for optimizer states, gradients, and parameters as a function of data-parallel degree N_d?
• RQ3 How can residual memory (activations, buffers, fragmentation) be managed to further enable very large models?
• RQ4 What is the resulting impact on communication volume and scalability when applying ZeRO-DP and ZeRO-R?
• RQ5 How does ZeRO enable training of models up to hundreds of billions of parameters, and what are practical speedups on real hardware?

Key findings

• ZeRO enables data-parallel training to fit models much larger than traditional DP/MP limits, with memory reductions up to factors roughly proportional to N_d for parameters and gradients and up to 4x–8x for optimizer state depending on stages.
• With three stages (P_os, P_g, P_p) the memory footprint per device can be reduced significantly, enabling trillion-parameter models on large GPU clusters.
• ZeRO-DP with 64-way DP can train a 7.5B model with much lower memory, and with P_os+g+p the memory scales down to roughly 1.88 GB for the 7.5B model, enabling much larger scales.
• ZeRO-R further reduces residual activation memory, uses constant-size buffers for efficiency, and performs on-the-fly memory defragmentation to mitigate fragmentation-related OOMs.
• Implementation results show 100B-parameter models trained on 400 Nvidia V100 GPUs with over 38 TFlops per GPU and aggregate performance exceeding 15 Petaflops; claimed data show up to 8x larger max model size and about 10x higher achievable performance than state-of-the-art at the time.
• ZeRO was released as part of the DeepSpeed library to democratize large-model training and aims to scale toward 1 trillion parameters with further partitioning (P_os+g+p).