Transformers have emerged as the architecture of choice for many state-of-the-art AI models, showcasing exceptional performance across a wide range of AI applications. However, the memory demands imposed by Transformers limit their ability to handle long sequences, thereby creating challenges for tasks involving extended sequences or long-term dependencies. We present a distinct approach, Ring Attention, which leverages blockwise computation of self-attention to distribute long sequences across multiple devices while concurrently overlapping the communication of key-value blocks with the computation of blockwise attention. By processing longer input sequences while maintaining memory efficiency, Ring Attention enables training and inference of sequences that are device count times longer than those of prior memory-efficient Transformers, effectively eliminating the memory constraints imposed by individual devices. Extensive experiments on language modeling tasks demonstrate the effectiveness of Ring Attention in allowing large sequence input size and improving performance.
翻译:变换器已成为众多先进人工智能模型的架构首选,在广泛的人工智能应用中展现出卓越性能。然而,变换器带来的内存需求限制了其处理长序列的能力,从而对涉及长序列或长期依赖关系的任务构成挑战。我们提出一种独特的方法——环状注意力,该方法利用自注意力的分块计算,将长序列分布到多个设备上,同时使关键值块的通信与分块注意力的计算重叠。通过在处理更长输入序列的同时保持内存效率,环状注意力能够训练和推理比此前内存高效变换器长设备数量倍的序列,有效消除了单个设备带来的内存约束。在语言建模任务上的大量实验证明了环状注意力在允许大序列输入大小和提升性能方面的有效性。