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


参考

相关推荐
bastgia1 天前
Tokenformer: 下一代Transformer架构
人工智能·机器学习·llm
新智元1 天前
李飞飞谢赛宁:多模态 LLM「空间大脑」觉醒,惊现世界模型雏形!
人工智能·llm
探索云原生1 天前
在 K8S 中创建 Pod 是如何使用到 GPU 的: nvidia device plugin 源码分析
ai·云原生·kubernetes·go·gpu
RWKV元始智能2 天前
RWKV-7:极先进的大模型架构,长文本能力极强
人工智能·llm
zaim12 天前
计算机的错误计算(一百八十七)
人工智能·ai·大模型·llm·错误·正弦/sin·误差/error
张拭心2 天前
Google 提供的 Android 端上大模型组件:MediaPipe LLM 介绍
android·人工智能·llm
带电的小王2 天前
whisper.cpp: Android端测试 -- Android端手机部署音频大模型
android·智能手机·llm·whisper·音频大模型·whisper.cpp
凳子花❀3 天前
市场常见AI芯片总结
ai·gpu
带电的小王3 天前
whisper.cpp: PC端测试 -- 电脑端部署音频大模型
llm·whisper·音视频·音频大模型
Ambition_LAO3 天前
LLaMA-Factory QuickStart 流程详解
llm·llama