ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

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.

References in corpus (9)

• Mixed Precision Training
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Training Deep Nets with Sublinear Memory Cost
• Adafactor: Adaptive Learning Rates with Sublinear Memory Cost
• An Empirical Model of Large-Batch Training
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• Mesh-TensorFlow: Deep Learning for Supercomputers
• Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
• Training Large Neural Networks with Constant Memory using a New Execution Algorithm

Cited by in corpus (25)

• Zero-Shot Text-to-Image Generation
• Beyond English-Centric Multilingual Machine Translation
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• PyTorch Distributed: Experiences on Accelerating Data Parallel Training
• ProtTrans: Towards Cracking the Language of Life's Code Through Self-Supervised Deep Learning and High Performance Computing
• Memory-Efficient Pipeline-Parallel DNN Training
• Adaptivity without Compromise: A Momentumized, Adaptive, Dual Averaged Gradient Method for Stochastic Optimization
• GSPMD: General and Scalable Parallelization for ML Computation Graphs
• TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models
• Training Large Neural Networks with Constant Memory using a New Execution Algorithm
• Parameter-Efficient Transfer Learning with Diff Pruning
• PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers
• TensorTEE: Unifying Heterogeneous TEE Granularity for Efficient Secure Collaborative Tensor Computing
• Revisiting BFloat16 Training
• Breaking the Memory Wall: A Study of I/O Patterns and GPU Memory Utilization for Hybrid CPU-GPU Offloaded Optimizers
• LAMP: Large Deep Nets with Automated Model Parallelism for Image Segmentation
• NxMTransformer: Semi-Structured Sparsification for Natural Language Understanding via ADMM
• Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloads
• Varuna: Scalable, Low-cost Training of Massive Deep Learning Models
• DistIR: An Intermediate Representation and Simulator for Efficient Neural Network Distribution
• IndoNLG: Benchmark and Resources for Evaluating Indonesian Natural Language Generation
• HetSeq: Distributed GPU Training on Heterogeneous Infrastructure
• Cloud Collectives: Towards Cloud-aware Collectives forML Workloads with Rank Reordering
• Improving Layer-wise Adaptive Rate Methods using Trust Ratio Clipping