Redian新闻
>
别再「浪费」GPU了,FlashAttention重磅升级,实现长文本推理速度8倍提升

别再「浪费」GPU了,FlashAttention重磅升级,实现长文本推理速度8倍提升

公众号新闻

机器之心报道

机器之心编辑部

处理小说、法律文件等长文本是大模型的一个重要应用方向,但也面临速度上的挑战。FlashAttention 作者 Tri Dao 等人提出的「Flash-Decoding」通过充分利用 GPU,可以将大模型的长上下文推理速度提高至 8 倍。

最近,像 ChatGPT 或 Llama 这样的大型语言模型(LLM)引起了前所未有的关注。然而,它们的运行成本仍然极高。虽然生成单个响应可能仅需 0.01 美元(在 AWS 上的 8xA100 实例上运行几秒钟),但当扩大规模以满足数十亿用户的需求时,成本会迅速累积。而且,这些用户可能每天与 LLM 进行多次互动。某些用例的成本更高,例如代码自动生成,因为它会随着每次输入新字符而运行。随着 LLM 应用的不断增加,即使在生成时间方面实现细微的效率提升,也将产生巨大的影响。

LLM 推理(或「解码」)是一个迭代的过程:token 逐个生成。生成包含 N 个 token 的完整句子需要通过模型进行 N 次前向传递。幸运的是,我们可以缓存先前计算的 token:这意味着单个生成步骤不依赖于上下文长度,除了一个单独的操作 —— 注意力。这个操作导致上下文长度不能很好地扩展。

在 LLM 的重要新兴用例中,有一些需要利用更长的上下文。只有拥有了更长的上下文窗口,LLM 才能对更长的文档进行推理,无论是总结文档还是回答其中的问题。此外,它们还可以保持更长的对话历史,甚至在编写代码之前处理整个代码库。举个例子,在 2022 年,大多数 LLM 的上下文长度最多为 2k(例如 GPT-3),但现在,有些开源 LLM 已经可以扩展到 32k(比如 Llama-2-32k),甚至有些模型已经达到了 100k(比如 CodeLlama)。在这些情境中,注意力操作在推理过程中占据了相当大的时间比例。

在扩展 batch size 维度时,即使上下文相对较短,注意力也可能成为一个瓶颈。这是因为随着 batch 维度的增加,需要读取的内存量也会增加,而对于模型的其余部分,内存需求只取决于模型的大小。

为了解决上述问题,FlashAttention 的作者 Tri Dao 等人提出了一项名为「Flash-Decoding」的技术,它显著加速了推理过程中的注意力计算,使长序列的处理生成速度提高到了原来的 8 倍。其主要思想是以最快的速度并行加载键和值,然后分别重新缩放和合并结果,以维持正确的注意力输出。

解码时的多头注意力

在解码期间,生成的每个新 token 都需要关注所有先前的 token,以计算:softmax (queries @ keys.transpose) @ values

这个操作已经在训练阶段通过 FlashAttention 进行了优化(包括最近的 v1 和 v2 版本),瓶颈是读写中间结果的内存带宽(如 Q @ K^T)。然而,这些优化并不直接适用于推理情况,因为瓶颈不同。在训练中,FlashAttention 并行处理 batch size 和查询长度两个维度。而在推理过程中,查询长度通常为 1:这意味着,如果 batch size 小于 GPU 上的流多处理器(streaming multiprocessor,SM)数量(例如 A100 有 108 个),该操作只会利用 GPU 的一小部分!特别是在处理长上下文时,情况尤为明显,因为它需要较小的 batch size 以适应 GPU 内存。当 batch size 为 1 时,FlashAttention 将使用不到 1% 的 GPU!

FlashAttention 只在查询块和 batch size 之间并行,并且在解码期间不会设法占用整个 GPU。

使用矩阵乘法基元也能执行注意力计算,这样就不需要使用 FlashAttention 了。在这种情况下,该操作会占用整个 GPU,但会启动许多写入和读取中间结果的内核,因此并不是最优的做法。

更快的注意力解码:Flash-Decoding

新方法 Flash-Decoding 基于 FlashAttention,同时引入了一个新的并行维度:键值序列的长度。它综合了上述两种方法的优点。与 FlashAttention 类似,它在全局内存中存储的额外数据很少。然而,只要上下文足够长,即使 batch size 较小,它也能充分利用 GPU。

Flash-Decoding 也在键和值之间并行化,代价是一个小的最终归约(reduction 步骤。

Flash-Decoding 主要有三个工作步骤:

  1. 首先,将键 / 值分成更小的块;
  2. 使用 FlashAttention 并行计算查询与每个这些分块的注意力,为每行和每个分块额外写入一个标量值:注意力值的 log-sum-exp
  3. 最后,通过对所有分块进行归约来计算实际输出,使用 log-sum-exp 来调整每个分块的贡献。

这一切之所以可行,都是因为注意力 /softmax 可以进行迭代计算。在 Flash-Decoding 中,它在两个级别上被使用:在分块内部(类似 FlashAttention),以及跨分块进行最终的归约计算。

实际操作中,步骤(1)不涉及任何 GPU 操作,因为键 / 值块是完整键 / 值张量的视图。然后,有两个独立的核函数,分别用于执行步骤(2)和(3)。

在 CodeLlama 34B 上进行的基准测试

为了验证上述新方法,研究者对 CodeLLaMa-34b 的解码吞吐量进行了基准测试。该模型与 Llama 2 具有相同的架构,一般来说,结果应该适用于许多大型语言模型。研究者在不同序列长度下(从 512 到 64k),以 tok/s 为单位来测量解码速度,并比较了多种计算注意力的方式:

  • Pytorch:使用纯粹的 PyTorch 基元来运行注意力计算(不使用 FlashAttention);
  • FlashAttention v2;
  • FasterTransformer:使用 FasterTransformer 的注意力内核;
  • Flash-Decoding;
  • 以及一个上限值,该值计算了从内存中读取整个模型和 KV-cache 所需的时间

对于非常大的序列,Flash-Decoding 可以将解码速度提高至 8 倍,并且比其他方法的扩展性要好得多。

在 prompt 比较小时,所有方法表现接近。但是当序列长度从 512 增加到 64k 时,除了 Flash-Decoding,其他方法的可扩展性都很差。在 Flash-Decoding 的这种模式下(batch size 为 1),扩展序列长度对生成速度的影响很小。

组件级微基准测试

研究者还在 A100 上对多头注意力进行了微基准测试,输入为 f16,考虑了不同的序列长度和 batch size。他们将 batch size 设置为 1,并且使用 16 个 128 维的查询头,以及 2 个键 / 值头(分组查询注意力),这与在 4 个 GPU 上运行的 CodeLLaMa-34b 使用的维度相匹配。

上述微基准测试展示了多头注意力的运行时间,单位为微秒。Flash-Decoding 在序列长度扩展到高达 64k 时,几乎实现了恒定的运行时间。

之前测量的高达 8 倍的端到端加速是可能的,因为注意力本身的速度比 FlashAttention 快高达 50 倍。在序列长度达到 32k 之前,注意力的时间大致是恒定的,因为 Flash-Decoding 能够完全利用 GPU。

使用 Flash-Decoding

Flash-decoding 可以在以下链接中找到:

  • FlashAttention 包,从 v2.2 开始:https://github.com/Dao-AILab/flash-attention/tree/main
  • xFormers 包(搜索 xformers.ops.memory_efficient_attention),从 0.0.22 开始:调度程序将根据问题的大小自动使用 Flash-Decoding 或 FlashAttention 方法。当这些方法不受支持时,它可以调度到一个高效的 triton 内核,该内核实现了 Flash-Decoding 算法。

一个完整的使用 LLaMa v2 / CodeLLaMa 的解码示例可以在 FlashAttention  repo 和 xFormers  repo 中找到。此外,作者还提供了一个简单的 LLaMa v1/v2 模型的高效解码代码示例,旨在快速、易读、有教育意义和易于修改。

参考链接:https://princeton-nlp.github.io/flash-decoding/



© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

微信扫码关注该文公众号作者

戳这里提交新闻线索和高质量文章给我们。
相关阅读
一周国际财经|股价一夜狂飙40倍!肝癌三联治疗法究竟有多厉害?ChatGPT下月或迎重磅升级;星链推出直连手机业务老黄给H100“打鸡血”:英伟达推出大模型加速包,Llama2推理速度翻倍微信公众号迎来重磅升级,越来越像小红书了?老黄深夜炸场,世界最强新品震撼发布:性能飙升90%,Llama 2推理速度翻倍 |【经纬低调分享】H100推理飙升8倍!英伟达官宣开源TensorRT-LLM,支持10+模型「专题速递」GPU算力思考、内容生产技术实践、AI动漫、多模态可控的作曲框架谁是真正的不完美受害人?LM-Infinite: 一种简单有效的大模型即时长度泛化,解决更长文本推理问题GPT-4一天顶6个月,人类审核要失业?OpenAI凌晨发布重磅升级,或大量取代人类审核员加拿大使馆官宣:签证审理速度慢,因为需要背调!如何申请加拿大签证?魯秋警官的故事2老黄深夜炸场,世界最强AI芯片H200震撼发布!性能飙升90%,Llama 2推理速度翻倍,大批超算中心来袭PyTorch官方认可!斯坦福博士新作:长上下文LLM推理速度提8倍ChatGPT之父首次「剧透」GPT-5/6 , OpenAI还想改变一切NewsFlash 3.0:大升级,焕然一新 | Linux 中国2024招聘季 | 美加学生均可申!Point72(US)开放Winter InternshipGPT-5 秘密训练曝光!ChatGPT 将再次迎来重磅升级?|Hunt Good 周报Erklärung zur Zusammenarbeit0.3%参数推理,实现78倍加速!ETH团队提出UltraFastBERT,构筑语言模型巨人孩子被人夸,别再「过度谦虚」了数学启蒙别再「硬教」了!这项底层思维能力,在家就能跟孩子玩出来GPT又迎重磅升级!万物皆可大模型(赠大模型系列课程+书籍)Ubuntu 23.10 将首次推出基于 Flutter 的新 Ubuntu 商店 | Linux 中国新MacBook Air或明年春季推出/Google发布Gemini大模型,多领域超越GPT-4/微软Copilot重磅升级谁统治下次大战的天空房企都进ICU了,还不「允许降房价」?别再使用 RestTemplate了,来了解一下官方推荐的 WebClient !骁龙8 Gen 3处理器发布:CPU性能提升30%、GPU提升25%女人的八字星雲與李玟:學會放下不容易颠覆Transformer霸权!CMU普林斯顿推Mamba新架构,解决致命bug推理速度暴增5倍【最新】《活力中国》第五集《城市“侦探”》,跟着沉浸式剧本推理游戏“重新发现”上海nǚ hóng?nǚ gōngMoonshot AI 初亮相,开启大模型长文本时代|Z News早财经丨北京一酒店竟现“楼梯间房”?查封;阳澄湖大闸蟹正式开捕;ChatGPT重磅升级:可以看图、说话
logo
联系我们隐私协议©2024 redian.news
Redian新闻
Redian.news刊载任何文章,不代表同意其说法或描述,仅为提供更多信息,也不构成任何建议。文章信息的合法性及真实性由其作者负责,与Redian.news及其运营公司无关。欢迎投稿,如发现稿件侵权,或作者不愿在本网发表文章,请版权拥有者通知本网处理。