什么是 Flash Attention

2024-09-09 07:36
文章标签 attention flash

本文主要是介绍什么是 Flash Attention,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

Flash Attention 是 由 Tri Dao 和 Dan Fu 等人在2022年的论文 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness 中 提出的,

论文可以从 https://arxiv.org/abs/2205.14135 页面下载,点击 View PDF 就可以下载。

下面我们通过详细解读这篇论文,来说明什么是Flash Attention。

Transformer在处理长序列时速度慢且占用大量内存,因为自注意力的时间和内存复杂度与序列长度的平方成正比。近似注意力方法尝试通过牺牲模型质量来减少计算复杂度来解决这个问题,但往往不能实现实际速度提升。我们认为一个缺失的原则是使注意力算法具有IO感知能力——考虑GPU内存层级之间的读取和写入。我们提出了FlashAttention,这是一种IO感知的精确注意力算法,使用平铺技术来减少GPU高带宽内存(HBM)和GPU片上静态随机存储器(SRAM)之间的内存读写次数。我们分析了FlashAttention的IO复杂度,表明它需要比标准注意力更少的HBM访问,并且在一定范围内的SRAM大小下是最优的。FlashAttention训练Transformer比现有基准更快:与MLPerf 1.1训练速度记录相比,BERT-large(序列长度为512)的端到端墙钟速度提高了15%,GPT-2(序列长度为1K)的速度提高了3倍,长距离竞技场(序列长度为1K-4K)的速度提高了2.4倍。

左图:FlashAttention使用平铺技术,防止在(相对)较慢的GPU HBM上生成大型的𝑁×𝑁注意力矩阵(虚线框)。在外部循环(红色箭头)中,FlashAttention循环遍历K和V矩阵的块,并将它们加载到快速的片上静态随机存储器(SRAM)中。在每个块中,FlashAttention循环遍历Q矩阵的块(蓝色箭头),将它们加载到SRAM,并将注意力计算的输出写回HBM。右图:与GPT-2上注意力的PyTorch实现相比的加速效果。FlashAttention不会将大型的𝑁×𝑁注意力矩阵读取和写入HBM,从而使注意力计算加速了7.6倍。

在现代GPU上,计算速度已经超过了内存读取速度,而Transformer中的大多数操作都受到内存访问的限制。对于类似受内存限制的操作,IO感知算法对于读取和写入数据可能占据运行时较大部分的情况至关重要,比如数据库连接、图像处理、数值线性代数等等。然而,诸如PyTorch和Tensorflow等常见的Python深度学习接口并不允许对内存访问进行精细控制。

我们提出了FlashAttention,这是一种新的注意力机制算法,可以在极少的内存访问次数下计算精确的注意力。我们的主要目标是避免从高带宽存储器(HBM)读取和写入注意力矩阵。这需要 (i) 在没有访问整个输入的情况下计算softmax归一化 (ii) 在反向传播过程中不存储大型的中间注意力矩阵。我们应用了两种成熟的技术来解决这些挑战。 (i) 我们重新构造了注意力计算,将输入分成块,并对输入块进行多次遍历,逐步执行softmax归一化(也称为平铺)。 (ii) 我们在正向传播过程中存储了softmax归一化因子,以便在反向传播过程中快速重新计算注意力,这比从HBM读取中间注意力矩阵的标准方法更快。我们在CUDA中实现了FlashAttention,以实现对内存访问的精细控制,并将所有的注意力操作融合到一个GPU核函数中。即使由于重新计算而增加了浮点运算次数,我们的算法在运行时比标准注意力更快,并且使用的内存量(与序列长度成正比)更少,这得益于大大减少的HBM访问量。

我们分析了FlashAttention的IO复杂度,证明它需要𝑂(𝑁^2𝑑^2/𝑀)次HBM访问,其中𝑑是头部维度,𝑀是SRAM的大小,而标准注意力的HBM访问次数为Ω(𝑁𝑑 + 𝑁^2)。对于𝑑和𝑀的典型值,我们证明FlashAttention相比标准注意力需要更少的HBM访问次数。此外,我们提供了一个下界,表明没有精确的注意力算法能够在所有SRAM大小上渐近地改进HBM访问次数。

GPU内存层次结构。GPU内存层次结构包括多种不同大小和速度的内存,较小的内存速度更快。以A100 GPU为例,其具有40-80GB的高带宽存储器(HBM),带宽为1.5-2.0TB/s,每个108个流多处理器具有192KB的片上静态随机存储器(SRAM),其带宽估计约为19TB/s。片上SRAM的速度比HBM快一个数量级,但在尺寸上小几个数量级。随着计算相对内存速度变得更快,操作越来越受到内存(HBM)访问的限制。因此,利用快速的SRAM变得更加重要。

执行模型。GPU拥有大量线程来执行操作(称为核函数)。每个核函数从HBM加载输入到寄存器和SRAM,进行计算,然后将输出写入HBM。

性能特征。根据计算和内存访问的平衡,操作可以被归类为计算受限或内存受限。这通常由算术密集度来衡量,即每个内存访问的算术操作数。

计算受限:操作所需时间取决于算术操作的数量,而在HBM上的访问时间要小得多。典型的例子包括具有大内部维度的矩阵乘法和具有大通道数量的卷积。
内存受限:操作所需时间取决于内存访问次数,而计算所需时间要小得多。例如,大多数其他操作:逐元素(例如,激活,丢弃)和缩减(例如,求和,softmax,批标准化,层标准化)。
核函数融合。加速内存受限操作的最常见方法是核函数融合:如果有多个操作应用于相同的输入,那么可以一次从HBM加载输入,而不是每个操作都进行多次加载。编译器可以自动融合许多逐元素操作。

然而,在模型训练的情境中,中间值仍然需要被写入HBM以供反向传播保存,降低了简单核函数融合的效果。

标准注意力实现

给定输入序列 Q、K、V ∈ R^(𝑁 ×𝑑) ,其中 𝑁 是序列长度,𝑑 是头维度,我们希望计算注意力输出 O ∈ R^(𝑁 ×𝑑) : S = QK^T∈ R^(𝑁 ×𝑁) ,P = softmax(S) ∈ R^(𝑁 ×𝑁) ,O = PV ∈ R^(𝑁 ×𝑑) , 其中 softmax 按行应用。 标准的注意力实现会将矩阵 S 和 P 实现在 HBM,这将占用 𝑂(𝑁^2 ) 的内存。 通常 𝑁 远大于 𝑑(例如,对于 GPT2,𝑁 = 1024,𝑑 = 64)。我们在算法 0 中描述了标准的注意力实现。由于一些或大部分操作是内存绑定的(例如 softmax),大量的内存访问会导致较慢的墙钟时间。 这个问题会因为其他应用于注意力矩阵的逐元素操作而变得更加严重,比如应用于 S 的掩码或应用于 P 的 dropout。因此,已经有许多尝试将多个逐元素操作融合在一起,比如将掩码与 softmax 融合在一起。标准注意力实现在序列长度 𝑁 的 HBM 访问方面呈二次增长。

算法 0

标准注意力实现 要求:矩阵 Q、K、V ∈ R^(𝑁 ×𝑑) 在 HBM 中。

 1: 从 HBM 中按块加载 Q、K,计算 S = QK>,将 S 写入 HBM。

 2: 从 HBM 中读取 S,计算 P = softmax(S),将 P 写入 HBM。

 3: 从 HBM 中按块加载 P 和 V,计算 O = PV,将 O 写入 HBM。

 4: 返回 O。

下面我们展示如何以更少的 HBM 读/写操作并且在不存储大型中间矩阵的情况下计算精确的注意力。这产生了一种既内存高效又在墙钟时间上更快的注意力算法。

一种具有平铺和重计算的高效注意力算法
给定输入Q、K、V ∈ R 𝑁 ×𝑑 在HBM中,我们的目标是计算注意力输出O ∈ R^(𝑁 ×𝑑) 并将其写入HBM。我们的目标是减少HBM访问的数量(至少是𝑁的平方级别)。我们应用了两种已建立的技术(平铺、重计算)来克服在子二次HBM访问中计算精确注意力的技术挑战。我们在算法1中描述了这一点。其主要思想是我们将输入Q、K、V分成块,从较慢的HBM加载到较快的SRAM,然后针对这些块计算注意力输出。通过在将每个块的输出按正确的归一化因子进行缩放后相加,我们最终得到了正确的结果。
平铺:我们通过块来计算注意力。Softmax将K的列耦合起来,因此我们分解了大型softmax并进行了缩放。为了数值稳定性,向量𝑥 ∈ R^𝐵的softmax计算如下:

对于向量 𝑥(1), 𝑥(2) ∈ R^𝐵,我们可以将拼接向量 𝑥 = [𝑥(1)𝑥(2)] ∈ R^(2𝐵) 的 softmax 分解为:

因此,如果我们跟踪一些额外的统计数据(𝑚(𝑥), ℓ(𝑥)),我们可以逐块计算softmax。因此,我们将输入 Q、K、V 分割成块(算法1第3行),计算softmax值以及额外的统计数据(算法1第10行),然后将结果合并(算法1第12行)。

重新计算。我们的一个目标是不要为反向传播存储𝑂(𝑁^2)个中间值。反向传播通常需要矩阵 S、P ∈ R^𝑁×𝑁 来计算相对于 Q、K、V 的梯度。然而,通过存储输出 O 和 softmax 归一化统计数据 (𝑚, ℓ),我们可以在反向传播中轻松地从 Q、K、V 的块中重新计算注意力矩阵 S 和 P。这可以看作是一种选择性梯度检查点技术。虽然梯度检查点技术被建议用于减少最大所需内存,但所有已知的实现都必须以速度为代价来换取内存。相比之下,即使有更多的浮点操作数,我们的重新计算也能加速反向传播,因为减少了高带宽内存访问。

实施细节:内核融合。分块使我们能够在一个CUDA内核中实现我们的算法,从高带宽内存加载输入,执行所有的计算步骤(矩阵乘法,softmax,可选的掩码和丢弃,矩阵乘法),然后将结果写回到高带宽内存。这避免了重复地从高带宽内存读取和写入输入和输出。

这篇关于什么是 Flash Attention的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/1150544

相关文章

STM32内部闪存FLASH(内部ROM)、IAP

1 FLASH简介  1 利用程序存储器的剩余空间来保存掉电不丢失的用户数据 2 通过在程序中编程(IAP)实现程序的自我更新 (OTA) 3在线编程(ICP把整个程序都更新掉) 1 系统的Bootloader写死了,只能用串口下载到指定的位置,启动方式也不方便需要配置BOOT引脚触发启动  4 IAP(自己写的Bootloader,实现程序升级) 1 比如蓝牙转串口,

STM32 ADC+DMA导致写FLASH失败

最近用STM32G070系列的ADC+DMA采样时,遇到了一些小坑记录一下; 一、ADC+DMA采样时进入死循环; 解决方法:ADC-dma死循环问题_stm32 adc dma死机-CSDN博客 将ADC的DMA中断调整为最高,且增大ADCHAL_ADC_Start_DMA(&hadc1, (uint32_t*)adc_buffer, ADC_Buffer_Size); 的ADC_Bu

图神经网络框架DGL实现Graph Attention Network (GAT)笔记

参考列表: [1]深入理解图注意力机制 [2]DGL官方学习教程一 ——基础操作&消息传递 [3]Cora数据集介绍+python读取 一、DGL实现GAT分类机器学习论文 程序摘自[1],该程序实现了利用图神经网络框架——DGL,实现图注意网络(GAT)。应用demo为对机器学习论文数据集——Cora,对论文所属类别进行分类。(下图摘自[3]) 1. 程序 Ubuntu:18.04

时序预测|变分模态分解-双向时域卷积-双向门控单元-注意力机制多变量时间序列预测VMD-BiTCN-BiGRU-Attention

时序预测|变分模态分解-双向时域卷积-双向门控单元-注意力机制多变量时间序列预测VMD-BiTCN-BiGRU-Attention 文章目录 一、基本原理1. 变分模态分解(VMD)2. 双向时域卷积(BiTCN)3. 双向门控单元(BiGRU)4. 注意力机制(Attention)总结流程 二、实验结果三、核心代码四、代码获取五、总结 时序预测|变分模态分解-双向时域卷积

阅读笔记--Guiding Attention in End-to-End Driving Models

作者:Diego Porres1, Yi Xiao1, Gabriel Villalonga1, Alexandre Levy1, Antonio M. L ́ opez1,2 出版时间:arXiv:2405.00242v1 [cs.CV] 30 Apr 2024 这篇论文研究了如何引导基于视觉的端到端自动驾驶模型的注意力,以提高它们的驾驶质量和获得更直观的激活图。 摘 要   介绍

通过 Flash 让所有浏览器支持 WebP 格式图像解码

http://www.guao.hk/tag/webp#userconsent# http://www.etherdream.com/WebP/

提高Flash builder编译速度 (转)

提高Flash builder编译速度   2013-04-08 03:14:42|  分类: flash |  标签: |举报 |字号大中小 订阅 我们在开发过程中随着项目的不断壮大,经常会碰到编译速度过慢,IDE崩溃等令人头痛的问题,这里我总结了一下网上别人的经验,对FB进行了3点优化,效果明显  1、把 Build Auto 改成 手动 Build  2

基于 BiLSTM+Attention 实现降雨预测多变量时序分类——明日是否降雨

前言 系列专栏:【深度学习:算法项目实战】✨︎ 涉及医疗健康、财经金融、商业零售、食品饮料、运动健身、交通运输、环境科学、社交媒体以及文本和图像处理等诸多领域,讨论了各种复杂的深度神经网络思想,如卷积神经网络、循环神经网络、生成对抗网络、门控循环单元、长短期记忆、自然语言处理、深度强化学习、大型语言模型和迁移学习。 降雨预测作为气象学和水文学领域的重要研究课题,‌对于农业、‌城市规划、

Show,Attend and Tell: Neural Image Caption Generation with Visual Attention

简单的翻译阅读了一下 Abstract 受机器翻译和对象检测领域最新工作的启发,我们引入了一种基于注意力的模型,该模型可以自动学习描述图像的内容。我们描述了如何使用标准的反向传播技术,以确定性的方式训练模型,并通过最大化变分下界随机地训练模型。我们还通过可视化展示了模型如何能够自动学习将注视固定在显着对象上,同时在输出序列中生成相应的单词。我们通过三个基准数据集(Flickr9k,Flickr

深入理解推荐系统:推荐系统中的attention机制

什么是attention机制、在推荐模型中的应用(会介绍相关模型,AFM/DIN/DIEN/DST)和参考文献  什么是attention机制  Attention函数的本质可以被描述为一个查询(query)到一系列(键key-值value)对的映射,在计算attention时主要分为三步 第一步是将query和每个key进行相似度计算得到权重,常用的相似度函数有点积,拼接,感知机等;