无烟锅网站规划与建设,做废钢推广网站,武威 网站建设,wordpress help sheet这篇文章的目的是详细的解释Flash Attention#xff0c;为什么要解释FlashAttention呢#xff1f;因为FlashAttention 是一种重新排序注意力计算的算法#xff0c;它无需任何近似即可加速注意力计算并减少内存占用。所以作为目前LLM的模型加速它是一个非常好的解决方案…这篇文章的目的是详细的解释Flash Attention为什么要解释FlashAttention呢因为FlashAttention 是一种重新排序注意力计算的算法它无需任何近似即可加速注意力计算并减少内存占用。所以作为目前LLM的模型加速它是一个非常好的解决方案本文介绍经典的V1版本最新的V2做了其他优化我们这里暂时不介绍。因为V1版的FlashAttention号称可以提速5-10倍所以我们来研究一下它到底是怎么实现的。 介绍
论文的标题是
“FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”
内存的效率与普通注意力相比序列长度是二次的O(N²)FlashAttention是次二次的/线性的N (O(N))。并且它不是注意力机制的近似值(例如稀疏或低秩矩阵近似值方法)-它的输出与“传统”注意力机制相同。与普通的注意力相比FlashAttention的注意力是”有感知“的。
它利用底层硬件的内存层次知识(例如gpu但其他AI加速器也应该工作我这里使用gpu作为示例)。一些[近似]方法在序列长度上将计算要求降低到线性或近线性但其中许多方法专注于减少FLOP而忽略内存访问(IO)的开销。 经过多年的发展gpu的FLOPS的增长速度一直在以比内存吞吐量(TB/s)更快。内存的瓶颈应该引起重视。FLOPS和内存吞吐量需要紧密结合由于硬件上的差距我们就需要软件层面上的工作进行平衡。
根据计算和内存访问之间的比率操作可以分为以下两种:
计算约束 矩阵乘法内存约束:元素操作(激活dropoutmasking)归并操作(softmax layer normsum等)
在当前的AI加速器GPU上是受内存大小限制的。因为它“主要由元素操作组成”或者更准确地说注意力的算术密度不是很高。
我们看看这个图: 可以看到maskingsoftmax和dropout是占用大量时间的操作而不是矩阵乘法(即使大部分FLOPS是在matmul中)。内存不是一个单一的工件它在本质上是分层的一般的规则是:内存越快越昂贵容量越小。 我们在上面说的FlashAttention的注意力是”有感知“的可以归结为利用SRAM比HBM(高带宽内存)快得多来确保减少两者之间的通信。
以A100为例
A100 GPU有40-80GB的高带宽内存(HBM)带宽为1.5-2.0 TB/s而每108个流处理器有192KB的SRAM带宽估计在19TB/s左右。
可以看到大小小了很多但是速度却提升了10倍所以如何高效的利用SRAM是提速的关键让我们看看标准注意力实现背后的计算: 标准实现如何显示对HW操作方式不大尊重。它基本上将HBM加载/存储操作视为0成本(它不是“io感知”)。
我们首先考虑如何使这个实现更有效(时间和内存方面)。最简单的方法是删除冗余的HBM读/写。
如何把S写回HBM只是为了(重新)加载它来计算softmax那么我们可以将其保存在SRAM中执行所有中间步骤然后将最终结果写回HBM。 内核基本上是“GPU操作”的一种奇特的说法参考我们以前发布的CUDA入门往简单了说就是一个函数。融合则可以将多个操作融合在一起。所以只从HBM加载一次执行融合的op然后将结果写回来。这样做可以减少通信开销。
这里还有一个专业名词术语是“materialization”物化/实体化。它指的是在上面的标准注意力实现中已经分配了完整的NxN矩阵(S, P)。下面我们将看到如何直接将内存复杂度从O(N²)降低到O(N)。
Flash attention基本上可以归结为两个主要观点:
Tiling (在向前和向后传递时使用)-基本上将NxN softmax/scores矩阵分块成块。
Recomputation (仅在向后传递中使用)
算法如下: 上面我们提到了很多名词你可能还不了解。没关系下面我们开始逐行解释算法。
FlashAttention算法
让Tiling方法的主要障碍是softmax。因为softmax需要将所有的分数列耦合在一起。 看到分母了吗?这就是问题所在。
要计算输入序列中的特定第i个标记对序列中其他标记的关注程度需要在SRAM中随时可用所有这些分数(这里用z_j表示)。
但是SRAM的容量是有限的。N(序列长度)可以是1000甚至100000个令牌。所以N²爆炸得很快。所以论文使用了一个技巧把softmax的计算分成更小的块最终仍然得到完全相同的结果。 我们可以只获取前一个B分数(x_1到x_B)并为它们计算softmax。然后通过迭代“收敛”到正确的结果。以一种聪明的方式组合这些每块部分softmax的数字这样最终的结果实际上是正确的。方法如下: 基本上为了计算属于前2个块(大小为B)的分数的softmax必须要跟踪每个块的2个统计数据:m(x)(最大分数)和l(x) (exp分数总和)。然后就可以用归一化系数将它们无缝地融合在一起。
这里主要是一些基本的代数运算通过展开f(x)和l(x)项并与e^x相乘一些项会相互抵消这里就不写了。这个逻辑递归地一直持续到最后一个(N/B)块这样就得到了N维正确的softmax输出!
为了详细的介绍这个算法假设有一个大小为1的批处理(即单个序列)和单个注意力头稍后会扩展它(通过简单地跨GPU的并行化-稍后会详细介绍)。我们暂时忽略了dropout和masking因为稍后再添加。 我们开始计算
初始化:HBM的容量以GB为单位测量(例如RTX 3090有24 GB的VRAM/HBM, A100有40-80 GB等)因此分配Q, K和V不是问题。
第1步
计算行/列块大小。为什么ceil(M / 4 d) ?因为查询、键和值向量是d维的所以我们还需要将它们组合成输出的d维向量。所以这个大小基本上允许我们用q k v和0个向量最大化SRAM的容量。
比如说假设M 1000, d 5。那么块大小为(1000/4*5) 50。所以一次加载50个q, k, v, o个向量的块这样可以减少HBM/SRAM之间的读/写次数。 对于B_r我也不太确定他们为什么要用d执行最小运算?如果有人知道请评论指教!
第2步 用全0初始化输出矩阵O。它将作为一个累加器l也类似它的目的是保存softmax的累积分母——exp分数的总和)。M(保存逐行最大分数)初始化为-inf因为我们将对其进行Max运算符因此无论第一个块的Max是什么-它肯定大于-inf 。
第3步 步骤1中的块大小将Q, K和V分成块。
第4步 将O, l, m分割成块(与Q的块大小相同)。
第5步 开始跨列循环即跨键/值向量(上图中的外部循环)。
第6步 将K_j和V_j块从HBM加载到SRAM。在这个时间点上我们仍然有50%的SRAM未被占用(专用于Q和O)。所以SRAM是这样的 第7步 开始跨行内部循环即跨查询向量。
第8步 将Q_i (B_r x d)和O_i (B_r x d)块以及l_i (B_r)和m_i (B_r)加载到SRAM中。
这里需要保证l_i和m_i能够载入SRAM(包括所有中间变量)这块可能是CUDA的知识我不太确定如何计算所以如果你有相关的信息请留言
第9步 计算Q_i (B_r x d)和K_j转置(d x B_c)之间的点积得到分数(B_r x B_c)。并没有将整个nxns(分数)矩阵“物化”。
假设外部循环索引为j (j3)内部循环索引为i (i2) N为25块大小为5下面就是刚刚计算的结果(假设以1为基础的索引): 也就是输入序列中标记11-15的标记6-10的注意力得分。这里的一个要点是这些都是精确的分数它们永远不会改变。
第10步 使用上一步计算的分数计算m_i_j、l*i_j和P~*i_j。M ~_i_j是按行计算的找到上面每一行的最大元素。
然后通过应用元素运算得到P~_i_j:
归一化-取行最大值并从行分数中减去它然后EXP
l~_i_j是矩阵P的逐行和。
第11步 计算m_new_i和l_new_i。同样非常简单可以重复使用上面的图表: M_i包含之前所有块的逐行最大值(j1 j2用绿色表示)。M _i_j包含当前块的逐行最大值(用黄色表示)。为了得到m_new_i我们只需要在m_i_j和m_i之间取一个最大值l_new_i也类似。
第12步(最重要): 这是算法中最难的部分。 它允许我们用矩阵的形式做逐行标量乘法。如果你有一列标量s (N)和一个矩阵a (NxN)如果你做diag(s)* a你基本上是在用这些标量做a行的元素乘法。
公式1(为了方便再次粘贴在这里): 第12步的第一项所做的(用绿色下划线)是:更新了在同一行块中当前块之前的块的当前softmax估计。如果j1(这是这一行的第一个块。
第一项乘以diag(l_i)是为了抵消之前迭代中除以的相同常数(这个常数隐藏在O_i中)。
表达式的第二项(黄色下划线)是不需要消去的因为可以看到我们直接将P~_i_j矩阵与V向量块(V_j)相乘。
e^x项是用来修改矩阵P~_i_j O_i的方法是消去前一次迭代中的m用最新的估计(m_new_i)来更新它该估计包含到目前为止逐行最大值。
以下是我的逐步分析(实际上只需要5分钟希望能有所帮助!) 重点是这些外面的e项和P/O矩阵里面的e项消掉了所以总是得到最新的m_new_1估计! 第三次迭代也是类似的得到了正确的最终结果!
回想一下:这只是对最终O_i的当前估计。只有在我们遍历上图中的所有红色块之后我们才能最终得到确切的结果。
第13步 将最新的累加到统计数据(l_i m_i)写回HBM。注意它们的维数是B_r。
第13、14、15、1步 嵌套的for循环结束O (Nxd)将包含最终结果:每个输入令牌的注意力加权值向量!
简单汇总
算法可以很容易地扩展到“block-sparse FlashAttention”这是一种比FlashAttention快2-4的稀疏注意力算法扩展到64k的序列长度!通过使用一个块形式的掩码矩阵可以跳过上面嵌套的for循环中的某些加载/存储这样我们可以按比例节省稀疏系数比如下图 现在让我们简单地讨论一下复杂性。
复杂度分析
空间:在HBM中分配了Q, K, V, O (Nxd) l和m (N)。等于4Nd 2*N。去掉常量并且知道d也是一个常量并且通常比N小得多(例如d{32,64,128} N{1024…100k})可以得到O(N)的空间这有助于扩展到64k序列长度(再加上一些其他“技巧”比如ALiBi)。
时间:这里不会严格地进行时间复杂度分析但是我们将使用一个好的指标:HBM访问的数量。
论文的解释如下 他们是怎么得到这个数字的?让我们来分析嵌套的for循环:
我们的块大小是M/4d。这意味着向量被分割成N/(M/4d)块。取它的2次方(因为要遍历行/列块)得到O(N²d²/ M²)
我们不能一次获取整个块如果做一个大O分析可能会让我们认为这并不比标准注意力好多少但对于典型的数字这导致访问次数减少了9倍(根据上面的论文截图)。
我们的伪算法集中在一个单头注意力假设批处理大小为1。下面我们就开始进行扩展了
多头注意力
要扩展到batch_size 1和num_heads 1实际上并不难。
算法基本上是由单个线程块(CUDA编程术语)处理的。这个线程块在单个流多处理器(SM)上执行(例如A100上有108个这样的处理器)。为了并行化计算只需要在不同的SMs上并行运行batch_size * num_heads线程块。该数字与系统上可用的SMs数量越接近利用率就越高(理想情况下是多个因为每个SM可以运行多个线程块)。
反向传播
对于GPU内存的占用另外一个大头就是反向传播通过存储输出O (Nxd)和softmax归一化统计数据(N)我们可以直接从SRAM中的Q, K和V (Nxd)块中反向计算注意力矩阵S (NxN)和P (NxN) !从而使内存保持在O (N)。这个比较专业了我们了解以下就可以了所以需要详细的内容请看原论文。
代码实现
最后让我们看看在使用flash attention时可能出现的一些问题。因为涉及到显存的操作所以我们只能深入CUDA但是CUDA又比较复杂。
这就是OpenAI的Triton等项目的优势(参见他们的FlashAttention实现)。Triton基本上是一种DSL(领域特定语言)介于CUDA和其他领域特定语言(例如TVM)之间的抽象级别。可以编写超级优化的Python代码(一旦编译)而不必直接处理CUDA。这样Python代码可以部署在任意的加速器上(这是Triton任务)。
另外一个好消息是Triton最近已经与PyTorch 2.0集成了。
另外对于某些用例比如对于超过1K的序列长度一些近似注意方法(如Linformer)开始变得更快。但是flash attention的块稀疏实现优于所有其他方法。
总结
你有没有想过对于这种底层优化的算法为什么是一个斯坦福大学的学生发布而不是NVIDIA的工程师?
我认为有2种可能的解释:
1、FlashAttention更容易/只能在最新的gpu上实现(原始代码库不支持V100)。
2、通常“局外人”是那些以初学者的眼光看待问题能够看到问题的根源并从基本原则出发解决问题
最后我们还是要进行个总结
FlashAttention能够让BERT-large训练中节省15%将GPT训练速度提高2/3并且是在不需要修改代码的情况下这是一个非常重要的进步它为LLM的研究又提出了一个新的方向。
论文地址
https://avoid.overfit.cn/post/9d812b7a909e49e6ad4fb115cc25cdc1
作者Aleksa Gordić
