Elamur Research
☰
home / papers / 2023 / Ring Attention with Blockwise Transformers for Near-Infinite Context

Ring Attention with Blockwise Transformers for Near-Infinite Context

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/2310.01889

Summary

Self-attention memory grows quadratically with sequence length. Even memory-efficient Transformers must keep each layer's full output, so a single accelerator's high-bandwidth memory caps the usable context. For example, 100 million tokens at hidden size 1024 would need over 1000GB, while current GPUs and TPUs offer less than 100GB. The authors extend blockwise parallel transformers (BPT) by splitting the sequence across hosts arranged in a ring. Each host keeps its own query block and computes attention and feedforward for it block by block. Meanwhile, key-value blocks are passed to the next host and received from the previous one. Blockwise attention gives the same result regardless of the order in which blocks are processed, so transfers can run in the background during computation. This adds no overhead as long as computing a block takes longer than sending it. Attention stays exact, and the maximum context scales linearly with the number of devices. Experiments on language modeling and reinforcement learning show contexts beyond 100 million tokens. This makes very long inputs, such as books, video, codebases and scientific data, practical without approximating attention.

Key points

  • With Ring Attention, sequences can be more than 500× longer than with earlier memory-efficient methods, exceeding 100 million tokens; maximum context was measured with end-to-end training on TPUv4-1024.
  • The baselines are vanilla Transformers, memory-efficient attention, and BPT. BPT already lowers per-layer activation memory to about 2bsh, with the feedforward reduced from 8bsh; Ring Attention removes the remaining per-device limit.
  • LLaMA-13B was finetuned on ShareGPT with a 512K-token context using 32 A100 GPUs, producing Ring Attention-13B-512K.
  • On the line retrieval test, Ring Attention-13B-512K kept high accuracy at long contexts. GPT3.5-turbo-16K, Vicuna-16B-16K and Claude-2-100K were competitive only at shorter lengths.
  • Unlike earlier ring-based sequence parallelism, this method fully overlaps communication with computation. Code is available at github.com/lhao499/llm_large_context (Jax, TPUs via Google TPU Research Cloud).

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

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

Cite

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

@misc{liu2023ring,
  title = {{Ring Attention with Blockwise Transformers for Near-Infinite Context}},
  author = {Liu, Hao and Zaharia, Matei and Abbeel, Pieter},
  year = {2023},
  doi = {10.48550/arxiv.2310.01889},
  eprint = {2310.01889},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2310.01889}
}

References

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

  1. Anthropic. Introducing claude, 2023. introducing-claude. find
  2. Christian Bischof. Parallel computing: Architectures, algorithms, and applications, volume 15. IOS Press, 2008. find
  3. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. find
  4. Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084–15097, 2021. find
  5. Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016. URL https://www.anthropic.com/index/ ↗
  6. Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E Gonzalez, et al. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality. See https://vicuna.lmsys.org, 2023. ↗
  7. Anthony Danalis, Ki-Yong Kim, Lori Pollock, and Martin Swany. Transformations to parallel codes for communication-computation overlap. In SC’05: Proceedings of the 2005 ACM/IEEE conference on Supercomputing, pages 58–58. IEEE, 2005. find
  8. Anthony Danalis, Lori Pollock, Martin Swany, and John Cavazos. Mpi-aware compiler optimizations for improving communication-computation overlap. In Proceedings of the 23rd international conference on Supercomputing, pages 316–325, 2009. find
  9. Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344–16359, 2022. find
  10. Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc’aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. Advances in neural information processing systems, 25, 2012. find
  11. Facebook. Fully Sharded Data Parallel: faster AI training with fewer GPUs — engineering.fb.com. https://engineering.fb.com/2021/07/15/open-source/fsdp/, 2023. ↗
  12. Xinyang Geng and Hao Liu. Openllama: An open reproduction of llama, may 2023. URL https://github. com/openlm-research/open\_llama, 2023. find
  13. Xinyang Geng, Arnav Gudibande, Hao Liu, Eric Wallace, Pieter Abbeel, Sergey Levine, and Dawn Song. Koala: A dialogue model for academic research. Blog post, April, 1, 2023. find
  14. Andrew Gibiansky. Bringing hpc techniques to deep learning. Baidu Research, Tech. Rep., 2017. find
  15. Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019. find
  16. Joshua Hursey and Richard L Graham. Building a fault tolerant mpi application: A ring communication example. In 2011 IEEE International Symposium on Parallel and Distributed Processing Workshops and Phd Forum, pages 1549–1556. IEEE, 2011. find
  17. Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509, 2023. find our card
  18. Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. arXiv preprint arXiv:2205.05198, 2022. find our card
  19. Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel. Urlb: Unsupervised reinforcement learning benchmark. arXiv preprint arXiv:2110.15191, 2021. — [arXiv:2110.15191](https://arxiv.org/abs/2110.15191) find
  20. Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph E. Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang. How long can open-source llms truly promise on context length?, June 2023. URL https://lmsys.org/blog/2023-06-29-longchat. ↗
  21. Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence parallelism: Long sequence training from system perspective. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volum… find
  22. Hao Liu and Pieter Abbeel. Emergent agentic transformer from chain of hindsight experience. International Conference on Machine Learning, 2023. find
  23. Hao Liu and Pieter Abbeel. Blockwise parallel transformer for large context models. Advances in neural information processing systems, 2023. find
  24. Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.02867, 2018. find
  25. MosaicML. Introducing mpt-7b: A new standard for open-source, commercially usable llms, 2023. URL https://www.mosaicml.com/blog/mpt-7b. ↗
  26. Sharan Narang, Hyung Won Chung, Yi Tay, William Fedus, Thibault Fevry, Michael Matena, Karishma Malkan, Noah Fiedel, Noam Shazeer, Zhenzhong Lan, et al. Do transformer modifications transfer across implementations and applications? arXiv preprint arXiv:2102.11972, 2021. find
  27. 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 Proceedings of the 27th ACM Symposium on Operating Systems Principles, pages 1–15, 2019. find
  28. Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory-efficient pipeline-parallel dnn training. In International Conference on Machine Learning, pages 7937–7947. PMLR, 2021. find
  29. OpenAI. Gpt-4 technical report, 2023. find our card
  30. Markus N Rabe and Charles Staats. Self-attention does not need o(n2) memory. arXiv preprint arXiv:2112.05682, 2021. find
  31. Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–16. IEEE, 2020. find
  32. J. Schulman, B. Zoph, C. Kim, J. Hilton, J. Menick, J. Weng, J. F. C. Uribe, L. Fedus, L. Metz, M. Pokorny, R. G. Lopes, S. Zhao, A. Vijayvergiya, E. Sigler, A. Perelman, C. Voss, M. Heaton, J. Parish, D. Cummings, R. Nayak, V. Balcom, D. Schnurr, T. Kaftan, C. Hallacy, N. Turley, N. Deutsch, and V.… find
  33. Alexander Sergeev and Mike Del Balso. Horovod: fast and easy distributed deep learning in tensorflow. arXiv preprint arXiv:1802.05799, 2018. find
  34. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019. find
  35. Yi Tay, Mostafa Dehghani, Samira Abnar, Hyung Won Chung, William Fedus, Jinfeng Rao, Sharan Narang, Vinh Q Tran, Dani Yogatama, and Donald Metzler. Scaling laws vs model architectures: How does inductive bias influence scaling? arXiv preprint arXiv:2207.10551, 2022. find
  36. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023. find
  37. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. find
  38. Shibo Wang, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, et al. Overlap communication with dependent computation via decomposition in large deep learning models. In Proceedings of the 28th ACM International… find
  39. Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto. Don’t change the algorithm, change the data: Exploratory data for offline reinforcement learning. arXiv preprint arXiv:2201.13425, 2022. — [arXiv:2201.13425](https://arxiv.org/abs/2201.13425) find