斯坦福博士提出超快省显存attention,gpt-2训练速度提升3.5倍,bert速度创纪录

QbitAl QbitAl     2023-03-13     301

关键词:

白交 发自 凹非寺
量子位 | 公众号 QbitAI

Flash is all you need!

最近,一个超快且省内存的注意力算法FlashAttention火了。

通过感知显存读取/写入,FlashAttention的运行速度比PyTorch标准Attention快了2-4倍,所需内存也仅是其5%-20%。

而它的表现还不止于此。

  • 训练BERT速度相较于MLPerf训练记录提升15%;

  • 训练GPT-2的速度提高3.5倍;

  • 训练Transformer的速度比现有基线快。

网友们纷纷表示惊叹:Great Job!这项工作对我来说很有用。

来看看这是一项什么样的研究~

FlashAttention

本文提出了一种IO感知精确注意力算法。

随着Transformer变得越来越大、越来越深,但它在长序列上仍然处理的很慢、且耗费内存。(自注意力时间和显存复杂度与序列长度成二次方)

现有近似注意力方法,在试图通过去牺牲模型质量,以降低计算复杂度来解决该问题。

但存在一定的局限性,即不能提升运行时的训练速度。

研究者认为,应该让注意力算法具有IO感知,即考虑显存级间的读写,比如大但慢的HBM(High Bandwidth Memory)技术与小但快的SRAM。

基于这样的背景,研究人员提出了FlashAttention,具体有两种加速技术:按块递增计算即平铺、并在后向传递中重新计算注意力,将所有注意力操作融合到CUDA内核中。

FlashAttention使用平铺来防止大的𝑁×𝑁注意力矩阵(虚线框)在GPU HBM上物化(materialization)。在外部循环中(红色箭头),FlashAttention循环通过K和V矩阵的块,并将其加载到SRAM。

在每个区块中,FlashAttention 循环Q矩阵的区块(蓝色箭头)将其加载到 SRAM,并将注意力计算的输出写回 HBM。

这样就产生了一种注意力算法,在实际耗时(wall-clock time)内,其内存效率和速度都很高,相比于标准的注意力算法可以更少地访问HBM。

结果比现有注意力算法都快

研究人员评估了FlashAttention来训练Transformer的影响,包括训练时间、模型准确性,以及注意力运行时间和内存效率。

首先在训练速度上。FlashAttention比MLPerf 1.1的BERT速度记录高出15%。

在实现GPT-2上,比HuggingFace速度高出3倍,比Megatron的标准Transformer速度高出1.8倍,FlashAttention将LRA(long-range arena)的基准速度提高了2.4倍。

模型质量,FlashAttention将Transformer扩展到更长的序列,并且质量更好。

长上下文的语言建模

如图所示,使用FlashAttention可以让GPT-2上下文长度增加4倍的情况下,训练时间还比Megatron-LM优化实现快30%,同时也获得了0.7的困惑度(困惑度越低,说明语言模型越好)。

长文档分类

对较长序列的Transformer训练可以提高MIMIC-III和ECtHR数据集的性能,比如序列长度为16K在MIMIC上比长度512多出4.3分。

MIMIC-III:包含重症监护室病人的出院总结,每个都有多个标签注释;ECtHR:包含欧洲人权法案的法律案件;两个数据集都包含很长的文本文件。

此外,还完成了第一个能在Path-X和Path-256任务中实现非随机性能的Transformer模型。

之后,研究人员还完成了基准测试,测量FlashAttention和块状稀疏(Block-Sparse)FlashAttention的运行时间和内存性能,并与带有40GB HBM的A100 GPU上的各种注意力基线进行了比较。

结果显示,FlashAttention的运行时间,比PyTorch注意力实现快3倍;在短序列情况下,FlashAttention在短序列中仍比近似和稀疏注意力运行得快;至于块状稀疏的FlashAttention,在所有的序列长度上都比现有注意力实现都快。

至于在显存效率方面,FlashAttention比PyTorch注意力基线高20倍

在64k序列长度、其他所有算法都已经耗尽显存的情况下,FlashAttention的效率仍比Linformer高2倍

斯坦福博士一作

这篇研究来自斯坦福大学计算机系以及纽约州立大学布法罗分校。共同一作是两位斯坦福计算机博士生Tri DaoDan Fu

感兴趣的朋友,可戳下方论文链接了解更多~

论文链接:
https://arxiv.org/abs/2205.14135
GitHub链接:
https://github.com/HazyResearch/flash-attention
参考链接:
https://twitter.com/tri_dao/status/1531437619791290369

—  —

直播报名 | 自动驾驶的量产之路:

为什么“渐进式”路径先看到了无人驾驶量产的曙光?

自动驾驶领域一直以来就有“渐进式”和“跨越式”两种路径之争,前者以特斯拉为代表,后者以Waymo为领头羊。

特斯拉宣布2024年实现新型“Robotaxi”的量产,而另一边是Waymo CEO离职,商业化落地裹足不前。在此背后,为什么“渐进式”路径被越来越多的机构看好?“渐进式”技术发展路径是什么?自动驾驶量产离我们的生活还有多远?

点这里👇关注我,记得标星哦~

《图神经网络导论》,斯坦福尤佳轩博士讲授

...分享,如有侵权,联系删除转载于:专知来自斯坦福大学的尤佳轩博士在CS520上讲述《图神经网络导论》,包含图神经网络的动机、概念、应用等。值得关注。尤佳轩,斯坦福大学CS在读四年级博士生,师... 查看详情

使用/指定 attention_mask 使用 Trainer 和 TrainingArguments 训练 GPT2

】使用/指定attention_mask使用Trainer和TrainingArguments训练GPT2【英文标题】:TrainGPT2withTrainer&TrainingArgumentsusing/specifyingattention_mask【发布时间】:2021-07-2102:45:23【问题描述】:我正在使用Trainer&TrainingArguments来训练GPT2模型,但这... 查看详情

多图详解attention和mask。从循环神经网络transformer到gpt2,我悟了

transformaer原理文章目录transformaer原理1.Transformer的兴起2.图解Attention2.1seq2seq2.2循环神经网络的不足:2.3attention的引出(重点内容)2.4从机器翻译推广到attention的一般模式2.5Attention模型的改进形式2.6Self-Attention2.6.1Self-Att 查看详情

斯坦福博士论文|机器学习的模型解释和数据评估

点上方人工智能算法与Python大数据获取更多干货在右上方 ··· 设为星标 ★,第一时间获取资源仅做学术分享,如有侵权,联系删除转载于:专知机器学习正在医疗健康等各种关键应用得到实施。为了能够信任... 查看详情

斯坦福博士论文|机器学习的模型解释和数据评估

点上方人工智能算法与Python大数据获取更多干货在右上方 ··· 设为星标 ★,第一时间获取资源仅做学术分享,如有侵权,联系删除转载于:专知机器学习正在医疗健康等各种关键应用得到实施。为了能够信任... 查看详情

一键生成ppt,斯坦福博士生自制ppt生成神器火了

 Datawhale干货 AI工具:ChatBCG,编辑:机器之心谁没有过为PPT熬夜的经历?现在,免费的一键生成PPT神器终于诞生了。在即将过完的12月,相信很多人都在忙一件事:年终工作汇报。就像那句话说的:... 查看详情

让ai替企业做复杂决策真的靠谱吗?参与直播,斯坦福博士来分享他的选择|量子位·视点...

视点发自凹非寺量子位|公众号QbitAI疫情加速全球产业链深度重构,市场供需两侧不确定因素增加;全球经济滞胀导致出口需求衰退,工业生产过热增长不可持续;行业整体产能过剩,产品同质化严重,绿... 查看详情

硕士生和博士生应该如何选导师?

...商业转载请联系作者获得授权,非商业转载请注明出处。斯坦福的一个教授在专业科研杂志Cell上讲如何挑博士生导师,完完全全的从学生的角度出发。想起我自己的一些经历,觉得如果很多年前看过这个文章就好了,把英文消... 查看详情

斑马识别成狗,ai犯错的原因被斯坦福找到了丨开源

...砸一些预测,肯定是在其中的某个环节出现了纰漏。斯坦福大学的两位博士生和教授JamesZou在一篇论文中,带我们探究了分类模型犯错的原因。随后,论文提出一种方法——反事实的 查看详情

glm国产大模型训练加速:性能最高提升3倍,显存节省1/3,低成本上手(代码片段)

作者|BBuf、谢子鹏、冯文2017年,Google提出了Transformer架构,随后BERT、GPT、T5等预训练模型不断涌现,并在各项任务中都不断刷新SOTA纪录。去年,清华提出了GLM模型(https://github.com/THUDM/GLM),不同... 查看详情

竞赛大佬在华为:网络专家出身斯坦福物理系,还有人“工作跟读博差不多”...

...。听上去有点离谱的是,拿下过国际物理奥赛金牌的斯坦福应用物理学博士,竟在菊厂搞起了网络。也不乏光学博士,敲起代码成了数据库团队的 查看详情

证明爱因斯坦错了!诺贝尔物理学奖今年颁给量子纠缠,潘建伟博士导师加冕...

杨净Pine发自凹非寺量子位|公众号QbitAI刚刚,2022年诺贝尔物理学奖揭晓!今年的物理学奖颁向量子纠缠的三位科学家,他们分别是阿兰·阿斯佩AlainAspect,约翰·克劳泽JohnF.Clauser以及安东·塞林格AntonZeilinger,以... 查看详情

证明爱因斯坦错了!诺贝尔物理学奖今年颁给量子纠缠,潘建伟博士导师加冕...(代码片段)

👇👇关注后回复 “进群” ,拉你进程序员交流群👇👇来源丨量子位(ID:QbitAI)杨净Pine发自凹非寺https://mp.weixin.qq.com/s/XpVS5puE_qvvABi8VPcpCQ刚刚,2022年诺贝尔物理学奖揭晓!今年的物理学 查看详情

深度学习attention的原理分类及实现(代码片段)

...层次2.4模型方面2.5相似度计算方式3Keras实现1原理1.1简介Attention是注意力机制,本质上对关注的多个对象有个不同的权重,从而关注重点不同。Attention是Transformer、Bert、GPT的基础。引入Attention机制的原因参数少:相对... 查看详情

lecun带两位uc伯克利华人博士提出「循环参数生成器」

...:新智元近日,LeCun带领两位来自UC伯克利的华人博士共同发表了一份关于如何减少参数冗余问题的论文,团队提出的RPG循环参数生成器,在减少骨干参数的同时࿰ 查看详情

快速串联rnn/lstm/attention/transformer/bert/gpt(代码片段)

参考:李宏毅2021/2022春机器学习课程王树森RNN&Transformer教程Transformer详解0.背景:序列数据及相关任务序列数据是由一组相互关联的样本组成的数据,其中任意样本对应的标记是由其自身和其他样本共同决定的;序列数据任务... 查看详情

learningfrominterpretableanalysis:attention-basedknowledgetracing

一篇短文LearningfromInterpretableAnalysis:Attention-BasedKnowledgeTracing|SpringerLinkInternationalConferenceonArtificialIntelligenceinEducation摘要知识追踪是个性化教育中一个公认的问题,也是一项重要的任务。近年来,人们提出了许多现有的工 查看详情

learningfrominterpretableanalysis:attention-basedknowledgetracing

一篇短文LearningfromInterpretableAnalysis:Attention-BasedKnowledgeTracing|SpringerLinkInternationalConferenceonArtificialIntelligenceinEducation摘要知识追踪是个性化教育中一个公认的问题,也是一项重要的任务。近年来,人们提出了许多现有的工 查看详情