Elamur Research
☰
home / papers / 2021 / ZeRO-Offload: Democratizing Billion-Scale Model Training

ZeRO-Offload: Democratizing Billion-Scale Model Training

Written for the Research library (claude-opus-5-5) from the paper text. Original — arxiv.org/pdf/2101.06840

Summary

Training models with billions of parameters usually needs many GPUs, because existing parallel methods only work when the combined GPU memory can hold the parameters, gradients and optimizer states. ZeRO-Offload is a heterogeneous GPU+CPU training system that lets a single GPU, or a small cluster, handle much larger models. Using first-principles analysis, the authors argue for one optimal split. Parameters and the forward/backward passes stay on the GPU. Gradients, optimizer states and the optimizer update move to CPU memory and CPU compute. This keeps CPU–GPU traffic and CPU work low while saving the most GPU memory. A vectorized CPU-Adam and a one-step delayed parameter update (DPU) stop the CPU from slowing training. For multi-GPU setups, the system builds on ZeRO-powered data parallelism and also works with model parallelism. On one V100 it trains 13B-parameter models at about 40 TFLOPS, roughly ten times what plain PyTorch can fit. It ships in DeepSpeed and needs only a few code changes, so large-model training becomes possible for users with modest hardware.

Key points

  • On a single NVIDIA V100, a 10B model reaches over 40 TFLOPS. Plain PyTorch peaks at about 30 TFLOPS with roughly 1.2–1.4B parameters, its memory limit (the paper gives both figures).
  • Combined with model parallelism, it trains models above 70B parameters on one DGX-2 node, about 4–4.5x more than model parallelism alone. A 10B GPT-2 scales almost linearly up to 128 GPUs at over 30 TFLOPS per GPU.
  • CPU-Adam is over 5x faster than the PyTorch CPU Adam across 1–10B parameters. For a 1B model it takes 0.22 s versus 1.39 s, a 6.4x speedup.
  • At micro batch size 8, DPU raises GPT-2 throughput by 1.12–1.59x without hurting convergence. BERT-large fine-tuned on SQuAD still reaches the baseline F1 of 92.8.
  • The code is open source in the DeepSpeed PyTorch library and does not require model refactoring.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Computing & GPUs

Cited by

246 papers of our library cite it (of the 140673 whose bibliographies we read); 9 of them have a page here, 237 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2101.06840 · arXiv 2101.06840

@misc{ren2021zero,
  title = {{ZeRO-Offload: Democratizing Billion-Scale Model Training}},
  author = {Ren, Jie and Rajbhandari, Samyam and Aminabadi, Reza Yazdani and Ruwase, Olatunji and Yang, Shuangyan and Zhang, Minjia and Li, Dong and He, Yuxiong},
  year = {2021},
  doi = {10.48550/arxiv.2101.06840},
  eprint = {2101.06840},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2101.06840}
}

References

34 references, 3 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. for Programming Languages and Operating Systems, ASPLOS ’20, page 1341–1355, New York, NY, USA, 2020. Association for Computing Machinery. find
  2. [1] The Stanford Question Answering Dataset (SQuAD) leaderboard. https://rajpurkar.github.io/ SQuAD-explorer/. [10] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of gi… ↗
  3. [2] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Wint… find
  4. [11] Hai Jin, Bo Liu, Wenbin Jiang, Yang Ma, Xuanhua Shi, Bingsheng He, and Shaofeng Zhao. Layer-centric memory reuse and data migration for extreme-scale deep learning on many-core architectures. ACM Trans. Archit. Code Optim., 15(3), September 2018. find
  5. [12] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020. find
  6. [3] Yu Cao, Wei Bi, Meng Fang, and Dacheng Tao. Pre-trained language models for dialogue generation with multiple input sources, 2020. find
  7. [13] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2014. find
  8. [4] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv: Learning, 2016. find
  9. [14] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, and Soumith Chintala. Pytorch distributed: Experiences on accelerating data parallel training. Proc. VLDB Endow., 13(12):3005– 3018, 2020. find our card
  10. [5] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc’aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, Quoc V. Le, and Andrew Y. Ng. Large scale distributed deep networks. In F. Pereira, C. J. C. Burges, L. Bottou, and K. Q. Weinberger, editors, Advances in Neura… find
  11. [15] Gaurav Mitra, Beau Johnston, Alistair Rendell, Eric McCreath, and Jun Zhou. Use of simd vector operations to accelerate application code performance on low-powered arm and intel platforms. pages 1107–1116, 05 2013. find
  12. [6] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018. find
  13. [16] Nvidia. Automatic Mixed Precision for Deep Learning. https://developer.nvidia.com/ automatic-mixed-precision, 2019. ↗
  14. [17] Xuan Peng, Xuanhua Shi, Hulin Dai, Hai Jin, Weiliang Ma, Qian Xiong, Fan Yang, and Xuehai Qian. Capuchin: Tensor-based gpu memory management for deep learning. In Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems, A… find
  15. [7] Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, and Phil Gibbons. Pipedream: Fast and efficient pipeline parallel dnn training, 2018. find
  16. [8] Mark Hildebrand, Jawad Khan, Sanjeev Trika, Jason Lowe-Power, and Venkatesh Akella. Autotm: Automatic tensor movement in heterogeneous memory systems using integer linear programming. In Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages a… find
  17. [18] Bharadwaj Pudipeddi, Maral Mesmakhosroshahi, Jinwen Xi, and Sujeeth Bharadwaj. Training large neural networks with constant memory using a new execution algorithm. June 2020. find
  18. [9] Chien-Chin Huang, Gu Jin, and Jinyang Li. Swapad-visor: Pushing deep learning beyond the gpu memory limit via smart swapping. In Proceedings of the Twenty-Fifth International Conference on Architectural Support [19] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutske… find
  19. [20] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2020. models using model parallelism. CoRR, abs/1909.08053, 2019. find
  20. [29] Roman Tezikov. PyTorch implementation of L2L execution algorithm. https://github.com/TezRomacH/ layer-to-layer-pytorch, 2020. ↗
  21. [21] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. In International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2020. find
  22. [30] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances… find
  23. [22] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100, 000+ questions for machine comprehension of text. In Jian Su, Xavier Carreras, and Kevin Duh, editors, Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas… find our card
  24. [31] G. Velkoski, M. Gusev, and S. Ristov. The performance impact analysis of loop unrolling. In 2014 37th International Convention on Information and Communication Technology, Electronics and Microelectronics (MIPRO), pages 307–312, 2014. find
  25. [23] Jie Ren, Jiaolin Luo, Kai Wu, Minjia Zhang, Hyeran Jeon, and Dong Li. Sentinel: Efficient Tensor Migration and Allocation on Heterogeneous Memory Systems for Deep Learning. In International Symposium on High Performance Computer Architecture (HPCA), 2020. find
  26. [32] Oreste Villa, Mark Stephenson, David Nellans, and Stephen Keckler. Buddy Compression: Enabling Larger Memory for Deep Learning and HPC Workloads on GPUs. In International Symposium on Computer Architecture, 2020. find
  27. [24] Minsoo Rhu, Natalia Gimelshein, Jason Clemons, Arslan Zulfiqar, and Stephen W. Keckler. vdnn: Virtualized deep neural networks for scalable, memory-efficient neural network design. In The 49th Annual IEEE/ACM International Symposium on Microarchitecture, MICRO-49, 2016. find
  28. [33] Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song, Zenglin Xu, and Tim Kraska. Superneurons: Dynamic gpu memory management for training deep neural networks. In Proceedings of the 23rd ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, PPoPP ’18, pa… find
  29. [25] Corby Rosset. Turing-nlg: A 17-billion-parameter language model by microsoft, 2020. find
  30. [26] Christopher J. Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E. Dahl. Measuring the Effects of Data Parallelism on Neural Network Training. Journal of Machine Learning Research, 20, 2019. find
  31. [34] Junzhe Zhang, Sai-Ho Yeung, Yao Shu, Bingsheng He, and Wei Wang. Efficient memory management for gpu-based deep learning systems. CoRR, abs/1903.06631, 2019. find
  32. [27] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, and Blake A. Hechtman. Mesh-tensorflow: Deep learning for supercomputers. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kri… find our card
  33. [35] M. A. Zinkevich, M.Weimer, A. Smola, and L. Li. Parallelized Stochastic Gradient Descent. In International Conference on Neural Information Processing Systems, 2010. find
  34. [28] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language find