Flash-Attention

这是一篇硬核的优化Transformer的工作。众所周知,Transformer模型的计算量和储存复杂度是 O ( N 2 ) O(N^2) O(N2) 。尽管先前有了大量的优化工作,比如LongFormer、Sparse Transformer、Reformer等等,一定程度上减轻了Transformer的资源消耗,但对Transformer的性能有所折损,且扩展性不强,不能泛化到其它领域、以及复杂结构的叠加。

这篇工作从底层对Transformer的计算和读写进行了优化,主要有三个贡献:

  1. 加速了模型计算:现在GPU的计算速度已经远远超过了内存读写速度,当GPU完成计算后,内存确还在读取数据,造成GPU闲置而内存繁忙读(消费者早就消费完了,生产者还在缓慢生产)的现象,也就是内存墙问题。FlashAttention通过tiling和算子融合计算,将复杂操作放到SRAM中计算,并减少从HBM读取次数,加快了模型计算速度。而之前的工作虽然减少了Transformer的计算复杂度,却并没有减少模型计算时间。
  2. 节省了显存:FlashAttention通过引入全局统计量,避免实例化大注意力矩阵,减少了显存占用。
  3. 精确注意力:FlashAttention从底层优化了Transformer的计算,但是任务指标上没有任何折损,与普通的Transformer结果是完全等价。

现代GPU内存分级


参考

相关推荐
Larcher2 小时前
新手也能学会,100行代码玩AI LOGO
前端·llm·html
架构师日志3 小时前
使用大模型+LangExtract从复杂文本提取结构化数据(三)——提取表格列表类型数据
llm
智泊AI3 小时前
AI圈炸锅了!大模型的下一片蓝海,彻底爆发了!
llm
常先森5 小时前
【解密源码】 RAGFlow 切分最佳实践- naive parser 语义切块(excel & csv & txt 篇)
架构·llm·agent
大模型教程5 小时前
RAG核心基础 Embedding 概念与技术详解
程序员·llm·agent
大模型教程5 小时前
一文搞懂大模型:何为深入理解Agent?
程序员·llm·agent
AI大模型6 小时前
一文搞懂RAG:阿里70K算法岗为什么都在用它?
程序员·llm·agent
AI大模型6 小时前
面试官狂问的 28 个 RAG 问题全解析:一次讲透,从基础到架构优化
程序员·llm·agent
一见8 小时前
多头注意力论文的作用
大模型·llm
在未来等你11 小时前
AI Agent设计模式 Day 1:ReAct模式:推理与行动的完美结合
设计模式·llm·react·ai agent·plan-and-execute