网站开发模板代码,网站开发实战视频,西安搜索引擎优化,wordpress本地做好了怎么备份Flash Attention 并没有减少 Attention 的计算量#xff0c;也不影响精度#xff0c;但是却比标准的Attention运算快 2~4 倍的运行速度#xff0c;减少了 5~20 倍的内存使用量。究竟是怎么实现的呢#xff1f;
Attention 为什么慢#xff1f;
此处的“快慢”是相对而言的…
Flash Attention 并没有减少 Attention 的计算量也不影响精度但是却比标准的Attention运算快 2~4 倍的运行速度减少了 5~20 倍的内存使用量。究竟是怎么实现的呢
Attention 为什么慢
此处的“快慢”是相对而言的。严格意义上来说相比于传统的 RNNTransformer中的Attention可以并行地处理序列所有位置的信息RNN 只能串行处理因此计算效率并不低但是仍然有可以进一步改进的空间。
众所周知科学计算通常分为计算密集型 (compute-bound) 和内存密集型 (memory-bound) 两类。其中计算密集型运算的时间瓶颈主要在于算数计算比如大型矩阵的相乘等而内存密集型运算的时间瓶颈主要在于内存的读写时间比如批归一化、层归一化等等。
时间复杂度Attention 需要对矩阵 Q 和矩阵 K 的转置做乘法来得到注意力权重矩阵。不考虑 batch 维度假设矩阵QK 的尺寸都为 ( n , d i m ) (n,dim) (n,dim)那么两个维度为 ( n , d i m ) (n,dim) (n,dim)的矩阵相乘的时间复杂度是序列长度n的平方级 O ( n 2 ) O(n^2) O(n2)在计算完注意力权重矩阵后还需要对其进行softmax操作这个算法需要分成三次迭代来执行 O ( n 3 ) O(n^3) O(n3)空间复杂度Attention的计算过程需要存储 S Q K T SQK^T SQKT和 P s o f t m a x ( S ) Psoftmax(S) Psoftmax(S)这两个尺寸均为 ( n , n ) (n,n) (n,n)的矩阵
为了对 Attention 的内存读取时间有更清晰的感知这里简单介绍 GPU 的内存层级。 GPU 的内存可以分为 HBM 和 SRAM 两部分。例如A100 GPU具有40-80 GB的高带宽内存上图中的 HBM即我们平时说的“显存”带宽为 1.5TB/s并且108个流式多核处理器都有 192 KB 的片上 SRAM带宽约为 19 TB/s。片上 SRAM 比 HBM 快一个数量级但容量要小很多个数量级。
在 GPU 运算之前数据和模型先从 CPU 的内存上图中的DRAM移动到 GPU 的 HBM然后再从 HBM 移动到 GPU 的 SRAMCUDA kernel 在 SRAM 中对这些数据进行运算运算完毕后将运算结果再从 SRAM 移动到 HBM。
所以提高Attention运算效率需要从降低attention的时间和空间复杂度入手。
时间复杂度
在 S Q K T SQK^T SQKT的计算过程中理论上尝试的方法主要可以分为稀疏 (sparse) 估计和低秩 (low-rank) 估计。但是在实际应用中仍然存在一些缺陷
性能比不上原始 attention。不论是稀疏估计、低秩估计还是其他这些方法都采用了某种近似算法来估算注意力权重矩阵难免会丢失信息。目前主流的还是原始的attention无法减少内存读取的时间消耗。这些方法只能降低 attention 的计算复杂度但是无法对 attention 运算过程中的空间复杂度等进行控制无法减少内存读写带来的时间损耗
所以在时间复杂度方向的优化主要在softmax的计算过程中 softmax ( x i ) e x i ∑ k 1 N e x k \operatorname{softmax}\left(x_{i} \right)\frac{e^{x_{i}}}{\sum_{k1}^{N} e^{x_{k}}} softmax(xi)∑k1Nexkexi
softmax 有个问题那就是很容易溢出。比如float16的最大值为65504所以只要 x ≥ 11 x\geq11 x≥11 的话softmax就溢出了。好在 exp 有这么一个性质那就是 e x − y e x e y e^{x-y} \frac{e^x}{e^y} ex−yeyex根据这个性质可以在分子分母上同时除以一个数这样可以将 x x x的范围都缩放到范围内保证计算 softmax 时的数值稳定性。这个算法可以分成三次迭代来执行
遍历所有数求 x 中的最大值m for i ← 1 , N do m i max ( m i , x i ) \begin{array}{l}\text { for } i \leftarrow 1, N \text { do } \\ \quad m_{i}\max \left(m_{i}, x_{i}\right)\end{array} for i←1,N do mimax(mi,xi) 2. 计算 softmax 分母并根据m对其进行缩放 for i ← 1 , N do d i d i − 1 e x i − m N \begin{aligned} \text { for } i \leftarrow 1, N \text { do } \\ d_{i} d_{i-1}e^{x_{i}-m_{N}}\end{aligned} for idi←1,N do di−1exi−mN 3. 求对应位置的 softmax for i ← 1 , N d o a i e x i − m N d N \begin{aligned} \text { for } i \leftarrow 1, N d o \\ a_{i} \frac{e^{x_{i}-m_{N}}}{d_{N}}\end{aligned} for iai←1,NdodNexi−mN
分析以上步骤可以发现如果是不做任何优化的话至少要进行和 GPU 进行6次通信3次写入3次写出如果对每一步的for循环进行一些并行切分的的话还要加上 reduce_sum 和 reduce_max 之类的通信成本。所以2018年 Nvidia 提出了《Online normalizer calculation for softmax》核心改进是去掉第二步 d i d i − 1 e x i − m N d_{i} d_{i-1}e^{x_{i}-m_{N}} didi−1exi−mN中对$m_N 的依赖设 的依赖设 的依赖设d_{i}{\prime}\sum_{j}{i} e^{x_{j}-m_{i}}$这里的全局最大值变成了当前最大值这个式子有如下的性质 d i ′ ∑ j i e x j − m i ∑ j i − 1 e x j − m i e x i − m i ∑ j i − 1 e x j − m i − 1 m i − 1 − m i e x i − m i ( ∑ j i − 1 e x j − m i − 1 ) e m i − 1 − m i e x i − m i d i − 1 ′ e m i − 1 − m i e x i − m i \begin{aligned} d_{i}^{\prime} \sum_{j}^{i} e^{x_{j}-m_{i}} \\ \sum_{j}^{i-1} e^{x_{j}-m_{i}}e^{x_{i}-m_{i}} \\ \sum_{j}^{i-1} e^{x_{j}-m_{i-1}m_{i-1}-m_{i}}e^{x_{i}-m_{i}} \\ \left(\sum_{j}^{i-1} e^{x_{j}-m_{i-1}}\right) e^{m_{i-1}-m_{i}}e^{x_{i}-m_{i}} \\ d_{i-1}^{\prime} e^{m_{i-1}-m_{i}}e^{x_{i}-m_{i}}\end{aligned} di′j∑iexj−mij∑i−1exj−miexi−mij∑i−1exj−mi−1mi−1−miexi−mi(j∑i−1exj−mi−1)emi−1−miexi−midi−1′emi−1−miexi−mi
这个式子依赖于 d i − 1 ′ m i m i − 1 d_{i-1}^{\prime}m_i m_{i-1} di−1′mimi−1。那么就可以将softmax前两步合并到一起
求 x 的最大值 m, 计算 softmax 的分母 for i ← 1 , N do m i max ( m i , x i ) d i ′ d i − 1 ′ e m i − 1 − m i e x i − m i \begin{array}{l}\text { for } i \leftarrow 1, N \text { do } \\ \qquad m_{i}\max \left(m_{i}, x_{i}\right) \\ \qquad d_{i}^{\prime}d_{i-1}^{\prime} e^{m_{i-1}-m_{i}}e^{x_{i}-m_{i}}\end{array} for i←1,N do mimax(mi,xi)di′di−1′emi−1−miexi−mi 2. 求对应位置的 softmax for i ← 1 , N d o a i e x i − m N d N \begin{aligned} \text { for } i \leftarrow 1, N d o \\ a_{i} \frac{e^{x_{i}-m_{N}}}{d_{N}}\end{aligned} for iai←1,NdodNexi−mN
以上的算法优化可以将3步合并变成2步将softmax的时间复杂度降为 O ( n 2 ) O(n^2) O(n2)。
空间复杂度
在将3步合成2步的同时
借助GPU的share memory来存储中间结果将上面的两步只用一个 kernel 实现这样就只需要与 global memory 通信两次一次写入数据一次读取结果还可以减少 Reduce_max 和 Reduce_sum 之类的通信成本
空间复杂度方面优化的基本思路是降低Attention对于显存的需求减少HBM和SRAM之间的换入换出充分利用 GPU 的并行优势进而减少Attention运算的时间消耗。
总结
Flash Attention的动机是尽可能避免大尺寸的注意力权重矩阵在 HBM 和 SRAM 之间的换入换出。论文中具体方法包含两个部分tiling 和 recomputation。
tiling 的基本思路不直接对整个输入序列计算注意力而是将其分为多个较小的块逐个对这些块进行计算增量式地进行 softmax 的规约。规约过程中只需要更新某些中间变量不需要计算整个注意力权重矩阵就是以上介绍的将三部合并成两步的过程。
recomputation 的基本思路基于 tiling 技巧在反向传播过程中不保留整个注意力权重矩阵而是只保留前向过程中 tiling 的某些中间变量然后在反向传播过程中重新计算注意力权重矩阵。recomputation 可以看作是一种基于 tiling 的特殊的 gradient checkpointing想进一步了解 recomputation 的读者可以翻阅Flash Attention原文。
得益于上述技巧Flash Attention 可以同时做到又快运算速度快又省节省显存。
