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内存分级


参考

相关推荐
胡攀峰7 小时前
第12章 微调生成模型
人工智能·大模型·llm·sft·强化学习·rlhf·指令微调
RUNTIME15 小时前
大模型微调实操记录
llm
数据智能老司机15 小时前
构建具备自主性的人工智能系统——探索协调者、工作者和委托者方法
深度学习·llm·aigc
数据智能老司机15 小时前
构建具备自主性的人工智能系统——使代理能够使用工具和进行规划
深度学习·llm·aigc
红尘散仙15 小时前
六、WebGPU 基础入门——Vertex 缓冲区和 Index 缓冲区
前端·rust·gpu
红尘散仙16 小时前
四、WebGPU 基础入门——Uniform 缓冲区与内存对齐
前端·rust·gpu
量子位19 小时前
智能车速度刷新:仅 10 个月,首个纯端侧大模型上车量产!
人工智能·llm
技术你大飞哥21 小时前
【突破数据孤岛】MCP协议进化史:从 STDIO 到全双工流式 —— AI 应用开发效率提升 90% 的秘密武器
llm·ai编程·mcp
Goboy1 天前
零基础搞定 Trae 智能体配置 + MySQL MCP 集成!手把手教学
llm·ai编程·trae