BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

有效的注意力模块在基于Transformer的大型语言模型（LLMs）的成功中发挥了关键作用，但这些注意力模块的二次时间和内存复杂度在处理长序列时也带来了挑战。处理长序列的一个潜在解决方案是利用分布式集群，在多个设备（例如GPU）上并行计算注意力模块。然而，采用分布式方法不可避免地会引入额外的内存开销来存储本地注意力结果，并产生额外的通信成本来将本地结果聚合到全局结果中。本文提出了一个名为“BurstAttention”的分布式注意力框架，以优化全局集群和本地设备级别的内存访问和通信操作。在实验中，我们将BurstAttention与其他竞争性的分布式注意力解决方案进行比较，以处理长序列。在不同长度设置下的实验结果表明，与这些竞争基线相比，BurstAttention在处理长序列时提供了显着的优势，减少了40%的通信开销，并在32 X A100上训练128K序列长度时实现了1.37倍的加速。

论文旨在解决Transformer-based大型语言模型处理长序列时的二次时间和内存复杂度问题，提出了一种分布式注意力框架BurstAttention来优化内存访问和通信操作。

BurstAttention在全局集群和本地设备级别上优化了内存访问和通信操作，相比于其他竞争性的分布式注意力解决方案，在处理长序列时具有显著的优势。

论文使用了128K序列长度的32 X A100进行实验，与其他竞争性的分布式注意力解决方案进行比较，结果表明BurstAttention可以减少40%的通信开销，并在训练时实现1.37倍的加速。此外，论文还提供了开源代码。

最近的相关研究包括：《Scaling Local Self-Attention for Parameter Efficient Visual Backbones》、《Longformer: The Long-Document Transformer》等。