ZeRO: Memory optimizations Toward Training Trillion Parameter Models

By Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He. 2019

1,913 citations recorded by the source corpus.

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 Turing-NLG, the world’s largest language model at the time (17B parameters) with record breaking accuracy.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Efficiency & Compression

Citations in the corpus

Training Compute-Optimal Large Language Models (2022, neutral): “• Whilst the forward and backward pass are computed in bfloat16 , we store a float32 copy of the weights in the distributed optimiser state (Rajbhandari et al., 2020).”
Reducing Activation Recomputation in Large Transformer Models (2022, methodology): “As we detail in Section 2, there are several other techniques to reduce the memory requirements of training large models, such as partitioning various data across the data parallel ranks or oﬄoading data to CPU memory [14, 17].”
Efficiently Scaling Transformer Inference (2022, positive): “For the reasons outlined in Rajbhandari et al. (2020), we typically split all-reduce into two phases: a reduction phase and a broadcast phase.”
Llemma: An Open Language Model For Mathematics (2023, methodology): “, 2019) with a world size of 2 for LLEMMA-7B , and a world size of 8 for LLEMMA-34B, alongside ZeRO Stage 1 sharded optimizer states (Rajbhandari et al., 2020) across Data Parallel (Goyal et al.”
FLAVA: A Foundational Language And Vision Alignment Model (2021, methodology): “We use Fully-Sharded Data Parallel (FSDP) [86,87] and train in full FP16 precision except the layer norm [3] to reduce GPU memory consumption.”
CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X (2023, methodology): “To increase training efficiency, we adopt an 8-way model parallel training together with 192-way data parallel training, with ZeRO2 [27] optimizer enabled to further reduce thememory consumption of optimizer states.”
The Future of Sensitivity Analysis: An essential discipline for systems modeling and policy support (2020, neutral): “In the machine learning context, the number of model parameters can reach millions (e.g., BERT; Houlsby et al., 2019), even trillions of parameters (e.g., ZeRO; Rajbhandari et al., 2020).”
This paper is included in the Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation. (positive): ““ZeRO-2” [44] is a memory-efﬁcient version of data parallelism which partitions gradients and optimizer states.”
A Survey of Large Language Models for Healthcare: from Data, Technology, and Applications to Accountability and Ethics (2023, neutral): “However, the ZeRO [247] is essentially a parallelism solution.”
A Comprehensive Overview of Large Language Models (2023, neutral): “[55]: Gated Linear Unit [56] is a neural network layer that is an element-wise product ( ⊗ ) of a linear transformation and a sigmoid transformed ( σ ) linear projection of the input given as”
SPHINX-X: Scaling Data and Parameters for a Family of Multi-modal Large Language Models (2024, positive): “To accommodate the large model volume, a combination of ZeRO2-style (Rajbhandari et al., 2020) data parallel and Megatron-style (Shoeybi et al., 2019) model parallel is used.”
A Survey on Mixture of Experts in Large Language Models (2024, neutral): “Thisapproachcanbeviewedas anaugmentationofdataparallelism[160],[161],[162],where eachexpertinanMoElayerisassignedtoadistinctdevice,while allnon-expertlayersareduplicatedacrossdevices.”
Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models (2024, methodology): “In deep learning, it has been observed that training models using data samples arranged in a strategically meaningful order can lead to improved performance compared to training on randomly shuffled data.”
DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving (2024, background): “DistServe is orthogonal to the large body of work on model parallelism in training [28,36,40,46,59].”
mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval (2024, methodology): “We train the model by transformers with deepspeed ZeRO (Rajbhandari et al., 2020) stage 1 in FP16 precision for roughly 0.4 epoch (240k steps, took 154 hours on 16 A100 80G GPUs) of our data (3.93B pairs on sampled data by Eq.1).”
Pre-Trained Language Models for Text Generation: A Survey (2022, methodology): “More-over, to train PLMs under limited computational resource, many efficient and scalable techniques are proposed, exemplified by 3D parallelism [97, 216], ZeRO [195], mixed precision training [166], and LoRA [92].”
CodeS: Towards Building Open-source Language Models for Text-to-SQL (2024, positive): “We leverage the DeepSpeed Zero-3 framework [54], employing BF16 mixed precision, to optimize GPU memory consumption during pre-training.”