PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一...

西风 发自 凹非寺
量子位 | 公众号 QbitAI

PyTorch团队让大模型推理速度加快了10倍

且只用了不到1000行的纯原生PyTorch代码

项目名为GPT-fast,加速效果观感是这样婶儿的:

7c9540465eb8744efe22790819196c79.gif

通畅,属实通畅!

重点是,团队直接放出了代码以及详细“教程”。还是简笔画版的那种,特别好理解。

开发团队成员@Horace He表示:

我们不把它看作是库或者框架,更希望大家能把它当成个例子,根据自己的需求“复制粘贴”。

2f7f5d63821d48fe11fdc92b4dbb9222.png

网友直接炸开锅,英伟达AI科学家Jim Fan评价道:

这是自Andrej Karpathy发布的minGPT以来最棒的教程式repo之一!

开源世界需要更多minGPT、GPT-Fast这样的项目!

d172c5194ab184a5df056f45337e0190.png

那么GPT-fast究竟是如何给大模型提速的?

开盒大模型“加速包”

总的来说,用到这几种方法:

  • Torch.compile:一个专门为PyTorch模型设计的编译器,可以提升模型运行效率。

  • GPU量化:通过减少计算的精度来加速模型的运算速度。

  • 推测性解码:使用一个较小的模型来预测较大模型的输出,以此加快大语言模型的运算。

  • 张量并行性:通过在多个硬件设备上分布模型的运算来加速处理速度。

下面我们来一一展开。

开发团队一开始使用简单的PyTorch来实现,但效果不佳(25.5 tok/s):

6dcf6ee1482ef13c90ce4c14f8e60265.png

他们查看跟踪后发现,一个原因是推理性能由于CPU过多占用而受限。

36832bfe39d83aa0dd18d43f0d975277.png

那么如何解决呢?

可以想象这样一个场景,GPU是一个庞大的工厂(拥有大量可用的算力),而CPU则是一个小推车,来回为工厂“供货”。

在很多情况下,CPU无法足够快地“喂”GPU。

因此,开发团队建议给GPU更多的工作量,或者说一次性给它更大“块”的任务来处理。

e7727c6e519597a3cc7bc7e95b684a68.png

在推理过程中要做到这一点,可以引入torch.compile

torch.compile能够捕获模型中更大的区域,并将其编译成单一的编译区域。特别是当以“reduce-overhead”模式运行时,它非常有效地减少了CPU的开销。

效果立竿见影,性能直接提升了4倍,从25 tok/s提高到107 tok/s:

a70b5bc077e6920470b8848fd148ae4b.png

接下来,开发团队想进一步提升速度,但遇到了内存带宽瓶颈。

开发团队计算了模型的带宽利用率,结果已经达到了72%:

255fea4428a1a4130ee16e4785084982.png

也就是说进一步提高速度的空间可能有限。

重新审视上面的方程式,团队发现虽然实际上不能改变模型参数量,也不能改变GPU的内存带宽(至少在不花更多钱的情况下),但可以改变存储每个参数所用的字节数

8ef26e3969d66149b5e92be9646d0e00.png

这意味着,虽然无法改变模型的大小或者升级硬件来提高性能,但可以通过减少存储模型参数所需的数据量来提高效率。

通常可以通过量化技术来实现,即减少表示每个参数所需的位数。

c3a46f1dec922b61784fe07ae2f5b73b.png

由此,开发团队引入了下一个技术——int8量化

采用int8权重量化减少了内存负载,进一步提升了性能(157.4 tok/s):

4732207c363daea1b62ee88b4b4ff18b.png

使用量化后还有一个问题:要生成100个token,必须加载(或调用)模型权重100次。频繁加载模型权重也会导致效率低下。

1b1fb0c1153db642c4a8edef910841b5.png

乍一看,好像没有什么解决的法子,因为在自回归生成模式中存在着严格的序列依赖关系。

但开发团队指出,通过利用推测性解码可以打破这种严格的序列依赖关系。

82372c17f10380b031ed73386487fea7.png

再来打个比方,想象有一个资深工程师Verity,他在技术决策上总是正确,但编写代码的速度相对较慢。

同时,还有一个初级工程师Drake,和Verity相反,不擅长技术决策,但编写代码的速度更快、成本也更低。

1744fa010ada2fd93afc91b26474a0cb.png

那么如何利用不同人的优势来提高整体效率?

方法很简单,先让Drake编写代码,并在此过程中做出技术决策。接下来,将代码交给Verity进行审查,不对的地方就让Drake重做。

419310f25c54fa0985c52bd6813d027f.png

在Transformer模型推理中,大型的验证模型即为Verity角色,Drake则是一个更小的、能更快生成文本的草稿模型。

开发团队使用草稿模型生成8个token,然后使用验证模型并行处理,丢弃不匹配的部分。

由此一来,打破了串行依赖,再次提高速度。

值得一提的是,推测性解码不会改变输出的质量。只要使用草稿模型生成token+验证这些token所需的时间少于单独生成这些token所需的时间,这种方法就是有效的。

而且使用原生PyTorch实现这种技术实际上非常简单,整个实现过程只需要大约50行原生PyTorch代码。

9760236331c42e5320a0aa11df08b8ac.png

由于AMD也支持Triton和torch.compile后端,因此之前在Nvidia GPU上应用的所有优化也可以在AMD GPU上重新应用。

开发团队观察到int8量化的加速从22 tok/s达到102 tok/s:

03144fc15615e3c350ebd16d242359c0.png

之后开发团队又用了int4量化,进一步提升速度,但模型准确性有所下降。

因此使用了分组量化和GPTQ降低权重大小。

9e260649e7e6c9878ca41a54ca301e01.png

最后在保证准确性的前提下,速度提升至202.1 tok/s:

88a92d0a95f286930cc987214eb8cf2f.png

将以上技术结合使用,达到更高速度244.7 tok/s:

8ddc5aa629260b58389c83157935de4f.png

到目前为止,研发团队一直都是在单个GPU上提速。但其实很多情况下是可以使用多个GPU的。

而使用多个GPU可以增加内存带宽,从而提高模型的整体性能。

db279a2e4b2e40c5e9f2d36495135290.png

在选择并行处理策略时,需要在多个设备上分割一个token的处理过程,所以需要使用张量并行性。

而PyTorch也提供了用于张量并行性的底层工具,可以与torch.compile结合使用。

开发团队还透露也正在开发用于表达张量并行性的更高级别的API。

然而,即使没有更高级别的API,添加张量并行性也很容易,150行代码即可实现,且不需要对模型进行任何改变。

6ace8367068b5190ddc6827e6d21b377.png

之前提到的所有优化都可以与张量并行性相结合。将这些优化结合起来,能够以55 tokens/s的速度为Llama-70B提供int8量化。

49b4248a9b8033212140130155998d12.png

最后总结成果,忽略量化,仅用766行代码(model.py 244行代码,generate.py 371行代码,tp.py 151行代码),就实现了快速推理、推测性解码和张量并行性。

对于Llama-7B,使用compile+int4量化+推测性解码速度达到241 tok/s。对于Llama-70B,通过加入张量并行性,达到80 tok/s。

这些性能都接近或超越了当前SOTA。

参考链接:
[1]https://pytorch.org/blog/accelerating-generative-ai-2/?utm_content=273712248&utm_medium=social&utm_source=twitter&hss_channel=tw-776585502606721024
[2]https://twitter.com/DrJimFan/status/1730298947376443698
[3]https://twitter.com/cHHillee/status/1730293330213531844

MEET 2024大会定档!

最新嘉宾阵容公布

12月14日,量子位「MEET2024智能未来大会」不容错过!点击报名线下现场

李培根院士、李开复博士及十余位AI各领域领先企业核心负责人已确认出席!戳此了解嘉宾详情:第二批嘉宾来袭!报名MEET2024的理由,今天又多了一个

a60557dd8d4986fb760f1cde945696be.jpeg

16db184c9804ef1f1c689606d5ba8277.jpeg

ce721eeb8739610d061c4314e9cca01e.jpeg

8d09c91fb36641ef255f9075ce600699.jpeg

3d3ab6c04acbf8d3fc680edff007f73f.jpeg

287cad169a693cab6ba2fa1f731bb8e0.jpeg

5b8abb4d065594d4a8b2cdbe63fe48c8.jpeg

3da352b4d25076d63c595a5f2148e03f.jpeg

b4d925148d7fbbad114a0b9af56ca6eb.jpeg

0fcd4565ac8a19cd0eeb280b719e0cdb.jpeg

fb979e9aafe9ed1434321f4fff9a58ea.jpeg

464cf222d1c8fc47ee284d52606e8d9f.jpeg

083ae00d1a4c0af35ae46a8ccc2fc9e8.jpeg

50a7802b8412b78f27c476269ada121e.jpeg

9688d96a8cc24ea4a9778393a7f7800b.jpeg

43213080f46f74893a5e5c3f086db35a.jpeg

2e57dd7d018e3b95666bfbf092a7fd21.jpeg

a52a901d217e7e73c40efec9b4076656.jpeg

< 左右滑动查看嘉宾海报 >

点击“预约”按钮,一键直达大会直播现场!

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

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值