FlashAttention-2:更快的注意力计算,更好的并行和工作分割
提出了 FlashAttention,一种 IO 感知的确切关注算法。FlashAttention 使用平铺减少 GPU 高带宽内存(HBM)和 GPU 片上 SRAM 之间的内存读取 / 写入次数,并可扩展为块状稀疏关注。FlashAttention 使 Transformers 速度提高,使其具有更长的上下文并获得更高质量的模型,以及实现了 Path-X 挑战的首个 Transformers。
May, 2022
我们扩展了 FlashAttention 以容纳一大类稀疏性注意力模式,其中包括关键 / 查询删除和基于哈希的注意力。即使具有相对较低的稀疏度,我们的方法在序列长度增加时可见地改进 FlashAttention 的性能。不牺牲困惑度的情况下,我们将 transformer 语言模型的训练速度分别提高了 2.0 倍和 3.3 倍,对于长度为 8k 和 16k 的序列。
Jun, 2023
通过引入数据依赖的门控机制,我们开发了一种硬件高效的并行形式,使得门控线性注意力(GLA)Transformers 在适度规模的语言建模中表现竞争力,并在训练速度上与基于 CUDA 优化的 FlashAttention-2 相媲美。
Dec, 2023
通过引入可分解的关注机制,我们将注意力机制的计算和内存复杂度从 O (N^2) 降低到 O (N),并保持了注意力矩阵的完整表示,具有稀疏性和全连接关系。经过各种标准设置的测试,结果表明我们的注意力机制具有稳定的性能,并在自我注意力被使用的多种应用领域中具有重要潜力。
Feb, 2024
LeanAttention 是一种可扩展的自注意力计算技术,通过重新设计解码阶段的执行流程,将自注意力机制的实现扩展到具有挑战性的长上下文长度情况,以并行计算的方式提供 2.6 倍的平均注意力执行加速和最多 8.33 倍的速度提升。
May, 2024
我们提出了一种名为 “BurstAttention” 的分布式注意力框架,通过在全局集群和本地设备级别上优化内存访问和通信操作,相比于竞争的基准线,在处理长序列时减少 40% 的通信开销,训练 32K 序列长度时实现 2 倍加速。
Mar, 2024
我们提供了 FlashAttention-2 前向传递的优化实现,使用了自定义融合的 CUDA 内核,以适应 NVIDIA Hopper 架构,并使用开源的 CUTLASS 库编写。在此过程中,我们解释了将在线 softmax 与连续的 GEMM 内核融合在一起的挑战和技术,利用 Hopper 特定的 Tensor Memory Accelerator(TMA)和 Warpgroup Matrix-Multiply-Accumulate(WGMMA)指令,定义和转换 CUTLASS 布局和张量,重叠复制和 GEMM 操作,并选择 Q、K 和 V 注意力矩阵的最优瓦片大小,同时平衡寄存器压力和共享内存利用率。在单个 H100 PCIe GPU 上进行的对比性测试中,针对某些常见的超参数选择,我们观察到与针对上一代 NVIDIA Ampere 架构进行优化的 FlashAttention-2 版本相比,FLOPs/s 高出 20-50%。
Dec, 2023
通过对大型语言模型的多种训练配置进行全面的消融研究,我们总结出几个关键建议,以实现最高效的训练,例如,我们发现使用微批量大小为 1 通常能实现最高效的训练布局。较大的微批量大小需要激活检查点或更高程度的模型并行性,并且还会导致更大的管道延迟。我们最高效的配置使得在一系列模型规模上实现了最先进的训练效率结果,尤其是在训练 13B 模型时,模型 FLOPs 利用率达到了 70.5%。
Nov, 2023
Transformer-based models have achieved state-of-the-art performance, but the quadratic complexity of self-attention limits their applicability to long sequences; Fast Multipole Attention addresses this issue by reducing time and memory complexity, while maintaining a global receptive field with a hierarchical approach.
Oct, 2023
我们引入了三种新的注意力机制,比标准的多头注意力在效率和学习能力方面表现更好,从而提高了 Transformer 模型的性能和广泛部署能力。我们的第一个贡献是优化的注意力,它在头部数量、参数数量和矩阵乘法数量上与标准注意力相近,但参数数量少了 3/4,每个头部少了一次矩阵乘法。接下来,我们介绍了高效的注意力,它在参数数量上只有标准注意力的一半,每个头部少了两次矩阵乘法,并且速度是标准注意力的两倍。最后,我们介绍了超级注意力,在视觉和自然语言处理任务中显著超过标准注意力,同时具有更少的参数和矩阵乘法。除了提供严谨的数学比较,我们还在 MNIST、CIFAR100、IMDB 电影评论和 Amazon 评论数据集上评估了所提出的注意力机制。
Mar, 2024