不到1000行代码,PyTorch团队让Llama 7B提速10倍

PyTorch 团队亲自教你如何减速大模型推理。在过去的一年里,生成式 AI 发展迅猛,在这当中,文本生成一直是一个特别受欢迎的领域,很多开源项目如 llama.cpp、vLLM 、 MLC-LLM 等,为了取得更好的效果,都在进行不停的优化。作为机器学习社区中最受欢迎框架之一的 PyTorch,自然也是抓住了这一新的机遇,不断优化。为此让大家更好的了解这些创新,PyTorch 团队专门设置了系列博客,重点介绍如何应用纯原生 PyTorch 减速生成式 AI 模型。代码地址:,PyTorch 团队展示了仅应用纯原生

PyTorch 团队亲自教你如何减速大模型推理。

在过去的一年里,生成式 AI 发展迅猛,在这当中,文本生成一直是一个特别受欢迎的领域,很多开源项目如 llama.cpp、vLLM 、 MLC-LLM 等,为了取得更好的效果,都在进行不停的优化。

作为机器学习社区中最受欢迎框架之一的 PyTorch,自然也是抓住了这一新的机遇,不断优化。为此让大家更好的了解这些创新,PyTorch 团队专门设置了系列博客,重点介绍如何应用纯原生 PyTorch 减速生成式 AI 模型。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

代码地址:https://github.com/pytorch-labs/gpt-fast

在第一篇博客中,PyTorch 团队展示了仅应用纯原生 PyTorch 重写 Segment Anything(SAM)模型,比原始兑现快 8 倍。在本博客中,他们又为我们带来了新的内容,即如何加快 LLM 推理。

我们先来看看结果,该团队重写 LLM,推理速度比基线足足快了 10 倍,并且没有损失准确率,只用了不到 1000 行的纯原生 PyTorch 代码!

不到1000行代码,PyTorch团队让Llama 7B提速10倍

所有基准测试都在 A100-80GB 上运行的,功率限制在 330W。

这些优化包括:

Torch.compile:PyTorch 模型编译器, PyTorch 2.0 加入了一个新的函数,叫做 torch.compile (),能够通过一行代码对已有的模型进行减速;

GPU 量化:通过降低运算精度来减速模型;

Speculative Decoding:一种大模型推理减速方法,应用一个小的「draft」模型来预测大的「目标」模型的输出;

张量并行:通过在多个设备上运行模型来减速模型推理。

接下来,我们看看每一步都是如何兑现的。

6 步加快大模型推理

该研究表示,在没有优化之前,大模型的推理机能为 25.5 tok/s,效果不是很好:

不到1000行代码,PyTorch团队让Llama 7B提速10倍

经过一番探索后终于找到了原因:CPU 开支过大。然后就有了下面的 6 步优化过程。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

第一步:通过 Torch.compile 和动态 KV 缓存减少 CPU 开支,兑现 107.0 TOK/S

torch.compile 允许用户将更大的区域捕获到单个编译区域中,特别是在 mode=”reduce-overhead” 时(参考下面的代码),这一功能对于减少 CPU 开支非常有效,除此以外,本文还指定 fullgraph=True,用来验证模型中没有「图形中断」(即 torch.compile 无法编译的部分)。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

然而,即使有 torch.compile 的加持,还是会遇到一些障碍。

第一个障碍是 kv 缓存。即当用户生成更多的 token 时, kv 缓存的「逻辑长度(logical length)」会增长。出现这种问题有两个原因:一是每次缓存增长时重新分配(和复制)kv 缓存的成本非常高;其次,这种动态分配使得减少开支变得更加困难。

为了解决这个问题,本文应用动态 KV 缓存,动态分配 KV 缓存的大小,然后屏蔽掉注意力机制中未应用的值。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

第二个障碍是 prefill 阶段。用 Transformer 进行文本生成可视为一个两阶段过程:1. 用来处理整个提示的 prefill 阶段 2. 解码 token.

尽管 kv 缓存被设置为动态化,但由于提示长度可变 ,prefill 阶段仍然需要更多的动态性。因此,需要应用单独的编译策略来编译这两个阶段。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

虽然这些细节有点棘手,但兑现起来并不困难,而且机能的提升是巨大的。这一通操作下来,机能提高了 4 倍多,从 25 tok/s 提高到 107 tok/s。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

第二步:通过 int8 权重量化缓解内存带宽瓶颈,兑现 157.4 tok /s

通过上文,我们已经看到应用 torch.compile 、动态 kv 缓存等带来的巨大减速,但 PyTorch 团队并不满足于此,他们又找了其他角度进行优化。

他们认为减速生成式 AI 训练的最大瓶颈是将权重从 GPU 全局内存加载到寄存器的代价。换句话说,每次前向传播都需要「接触(touch)」GPU 上的每个参数。那么,理论上我们能够以多快的速度「接触」模型中的每个参数?

不到1000行代码,PyTorch团队让Llama 7B提速10倍

为了衡量这一点,本文应用模型带宽利用率(MBU),计算它非常简单,如下所示:

不到1000行代码,PyTorch团队让Llama 7B提速10倍

举例来说,对于一个 7B 参数模型,每个参数都存储在 fp16 中(每个参数 2 字节),可以兑现 107 tokens/s。A100-80GB 理论上有 2 TB/s 的内存带宽。

如下图所示,将上述公式带入具体的数值,可以得到 MBU 为 72%!这个结果是相当不错的,因为很多研究很难突破 85%。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

但 PyTorch 团队还想将这个数值在提高一些。他们发现无法改变模型中参数的数量,也无法改变 GPU 的内存带宽。但他们发现可以更改每个参数存储的字节数!

不到1000行代码,PyTorch团队让Llama 7B提速10倍

因此,他们打算用 int8 量化。 

不到1000行代码,PyTorch团队让Llama 7B提速10倍

请注意,这仅是量化权重,计算本身仍然在 bf16 中完成。此外,有了 torch.compile,可以轻松生成 int8 量化的高效代码。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

不到1000行代码,PyTorch团队让Llama 7B提速10倍

就像上图所展示的,从深蓝色线(torch.compile + int8)可以看出,应用 torch.compile + int8 仅权重量化时,机能有显着提升。

将 int8 量化应用于 Llama-7B 模型,机能提高了约 50%,达到 157.4 tokens/s。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

第三步:应用 Speculative Decoding

即使在应用了 int8 量化等技术之后,该团队仍然面临着另一个问题,即为了生成 100 个 token,必须加载权重 100 次。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

即使权重被量化,一遍又一遍地加载权重也避免不了,这种问题该如何解决呢?事实证明,利用 speculative decoding 能够打破这种严格的串行依赖性并获得减速。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

该研究应用草稿(draft)模型生成 8 个 token,然后应用验证器模型并行处理,丢弃不匹配的 token。这一过程打破了串行依赖。整个兑现过程大约 50 行原生 PyTorch 代码。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

第四步:应用 int4 量化和 GPTQ 方法进一步减小权重,兑现 202.1 tok/s

本文发现,当权重为 4-bits 时,模型的准确率开始下降。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

为了解决这个问题,本文应用两个技巧来解决:第一个是拥有更细粒度的缩放因子;另一种是应用更先进的量化策略。将这些操作组合在一起,得到如下:

不到1000行代码,PyTorch团队让Llama 7B提速10倍

第五步:将所有内容组合在一起,得到 244.7 tok/s

最后,将所有技术组合在一起以获得更好的机能,得到 244.7 tok/s。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

第六步:张量并行性

到目前为止,本文一直是在单个 GPU 上最大限度地减少延迟。其实,应用多个 GPU 也是可以的,这样一来,延迟现象会得到进一步改善。

非常庆幸的是,PyTorch 团队提供了张量并行的低级工具,只需 150 行代码,并且不需要任何模型更改。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

前面提到的所有优化都可以继续与张量并行性组合,将这些组合在一起,能以 55 tokens/s 的速度为 Llama-70B 模型提供 int8 量化。

不到1000行代码,PyTorch团队让Llama 7B提速10倍

最后,简单总结一下文章主要内容。在 Llama-7B 上,本文应用「compile + int4 quant + speculative decoding」这一套组合拳,兑现 240+ tok/s。在 Llama-70B,本文还通过引入张量并行性以达到约 80 tok/s,这些都接近或超过 SOTA 机能。

原文链接:https://pytorch.org/blog/accelerating-generative-ai-2/

给TA打赏
共{{data.count}}人
人已打赏
AI

Meta教你5步学会用Llama2:我见过最简单的大模型教学

2023-12-4 11:30:00

AI

线上开售!大模型时期的向量数据库 AI 手艺论坛圆满收官

2023-12-6 15:20:00

0 条回复 A文章作者 M管理员
    暂无讨论,说说你的看法吧
个人中心
今日签到
有新私信 私信列表
搜索