斯坦福大学的研究人员提出了一种名为FlashAttention的新颖注意力算法，该算法针对GPU进行了优化，能够实现Transformer模型训练速度提升2-4倍。传统自注意力机制在序列长度上的计算复杂度和内存需求较高，阻碍了长上下文的处理。尽管有稀疏和低秩近似方法试图解决这一问题，但它们并未显著提高实际运行时速度。研究强调了IO感知的重要性，即考虑内存访问速度对计算效率的影响。FlashAttention通过IO感知设计实现了更快的wall-clock速度，有望改善Transformer模型的训练效率。
#今日论文推荐# 斯坦福提出新型Attention算法,提速2-4倍，BERT单节点训练最快
Transformer 已然成为自然语言处理和图像分类等应用中最广泛使用的架构。随着研究的不断前进，Transformer 尺寸变得越来越大、层数也越来越深，但是给 Transformer 配备更长的上下文仍然很困难，因为 Transformer 核心自注意力模块的时间复杂度以及内存复杂度在序列长度上是二次方的。
有研究者提出一些近似注意力的方法，旨在减少注意力计算和内存需求。这些方法包括稀疏近似、低秩近似以及它们的组合。从序列长度来看，尽管这些方法可以将计算降低到线性或接近线性，但它们并没有显示出针对标准注意力的 wall-clock 加速，因而没有被广泛使用。这其中一个主要原因是这些研究专注于减少 FLOP（这可能与 wall-clock 速度无关）并且倾向于忽略来自内存访问 (IO) 的开销。
在本文中，该研究认为应该让注意力算法具有 IO 感知——即考虑显存级间的读写。现代 GPU 计算速度超过了内存速度，transformer 中的大多数操作都被内存访问所阻塞。IO 感知算法对于类似的内存绑定操作至关重要，这种重要性体现在当读写数据占据很大运行时——例如数据库连接、图像处理、数值线性代数等。然而，用于深度学习的常见 Python 接口，如 PyTorch 和 Tensorflow，不允许对内存访问进行细粒度控制。
论文题目：FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
详细解读：https://www.aminer.cn/research_report/62aff0767cb68b460fd6d24b?download=false
https://www.aminer.cn/research_report/62aff0767cb68b460fd6d24b?download=false
AMiner链接：https://www.aminer.cn/?f=cs
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
