Elamur Research
☰
home / papers / 2022 / Reducing Activation Recomputation in Large Transformer Models

Reducing Activation Recomputation in Large Transformer Models

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "arXiv — non-exclusive license to distribute" does not allow redistributing the text itself. Original — arxiv.org/pdf/2205.05198

Summary

Training transformer models with hundreds of billions of parameters is limited by GPU memory. Activations saved for the backward pass are a major cost. Pipeline parallelism does not reduce this cost, and tensor parallelism reduces it only partly. The usual workaround, full activation recomputation, discards most activations and recomputes them during backpropagation, which adds about 30–40% to execution time. The NVIDIA authors first derive an approximate formula for activation memory per transformer layer. They then add two simple techniques on top of Megatron-style tensor parallelism:

  • Sequence parallelism splits layer-norm and dropout regions along the sequence dimension, so these activations are no longer duplicated on every tensor-parallel rank.
  • Selective activation recomputation recomputes only the parts of the layer that use a lot of memory but little compute.

They test GPT-style models ranging from 22B to 1T parameters on NVIDIA A100 GPUs. Activation memory falls about fivefold, and over 90% of the recomputation overhead disappears. End-to-end throughput rises by roughly 30%. This matters because it shortens and cheapens large-model training without extra communication or memory costs. The code is released in Megatron-LM and NeMo-Megatron.

Key points

  • Estimated activation memory per layer: attention needs 11sbh + 5as²b bytes, the MLP 19sbh, and the two layer norms 4sbh.
  • Recomputation overhead per layer falls to about 2% for the 530B and 1T models, compared with 36% under full recompute.
  • End-to-end throughput improves by 29.0–32.1% across the 22B, 175B, 530B and 1T configurations. The 1T model reaches 56.3% MFU and 57.0% HFU.
  • A 530B model trained on 2240 A100 GPUs with 8-way data parallelism achieves 54.2% Model FLOPs Utilization, compared with 42.1% using full recompute (about 29% faster).
  • The method complements data-parallel sharding and CPU offloading, which the authors did not compare. Future work targets memory fragmentation and memory pressure on the first pipeline stage.

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

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

Cite

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

@misc{korthikanti2022reducing,
  title = {{Reducing Activation Recomputation in Large Transformer Models}},
  author = {Korthikanti, Vijay and Casper, Jared and Lym, Sangkug and McAfee, Lawrence and Andersch, Michael and Shoeybi, Mohammad and Catanzaro, Bryan},
  year = {2022},
  doi = {10.48550/arxiv.2205.05198},
  eprint = {2205.05198},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2205.05198}
}

References

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

  1. NVIDIA A100 Tensor Core GPU. https://www.nvidia.com/en-us/data-center/a100/. ↗
  2. NVLink and nVSwitch. https://www.nvidia.com/en-us/data-center/nvlink/. ↗
  3. Selene. https://www.top500.org/system/179842/. ↗
  4. 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 Winter,… find
  5. Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR, abs/1604.06174, 2016. find
  6. Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emi… find
  7. Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d’… find
  8. Can Karakus, Rahul Huilgol, Fei Wu, Anirudh Subramanian, Cade Daniel, Derya Cavdar, Teng Xu, Haohan Chen, Arash Rahnama, and Luis Quintela. Amazon sagemaker model parallelism: A general and flexible framework for large model training, 2021. find
  9. Shenggui Li, Fuzhao Xue, Yongbin Li, and Yang You. Sequence parallelism: Long sequence training from system perspective, 2021. find
  10. Zhuohan Li, Siyuan Zhuang, Shiyuan Guo, Danyang Zhuo, Hao Zhang, Dawn Song, and Ion Stoica. Terapipe: Token-level pipeline parallelism for training large-scale language models. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 1… find
  11. Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gregory R. Ganger, Phillip B. Gibbons, and Matei Zaharia. Pipedream: generalized pipeline parallelism for DNN training. In Tim Brecht and Carey Williamson, editors, Proceedings of the 27th ACM Symposium on Operating… find
  12. Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory-Efficient Pipeline-Parallel DNN Training. arXiv preprint arXiv:2006.09503, 2020. find
  13. Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catan-zaroand Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on gpu clusters using megatron-… find
  14. Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations toward training trillion parameter models. In Christine Cuicchi, Irene Qualters, and William T. Kramer, editors, Proceedings of the International Conference for High Performance Computing, Networking, Storag… find
  15. Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero-infinity: breaking the GPU memory wall for extreme scale deep learning. In Bronis R. de Supinski, Mary W. Hall, and Todd Gamblin, editors, SC ’21: The International Conference for High Performance Computing, Network… find
  16. Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Rajesh Gupta, Yan Liu, Jiliang Tang, and B. Aditya Prakash, editors, KDD ’20: The 26th ACM SIGKDD Conference on Knowledge Discov… find
  17. Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scale model training. In Irina Calciu and Geoff Kuenning, editors, 2021 USENIX Annual Technical Conference, USENIX ATC 2021, July 14-16, 20… find our card
  18. 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, Kristen… find our card
  19. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. CoRR, abs/1909.08053, 2019. find
  20. Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, Elton Zhang, Rewon Child, Reza Yazdani Aminabadi, Julie Bernauer, Xia Song, Mohammad Shoeybi, Yuxiong He, Michael Houston, Saurabh Tiwa… find our card
  21. Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake A. Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, Ruoming Pang, Noam Shazeer, Shibo Wang, Tao Wang, Yonghui Wu, and Zhifeng Chen. GSPMD: general and scalable parallelization for ML computation graphs. C… — [arXiv:2105.04663](https://arxiv.org/abs/2105.04663) find