蚂蚁开源新算法,给大模型提提速!-技术圈

ChatGPT Plus 一个月就要 20 美元，这是笔不小的费用。虽然已经是尊贵的 Plus 用户，但每次输入一个问题后，还是需要等待一段时间才能拿到结果，感觉就像逐字蹦出来的一样。这是什么原因呢？

其实，回答的响应速度与服务成本，关键在于 推理效率 ，这也是大模型目前主要技术攻坚方向之一。

近日，蚂蚁开源了一套新算法—— Lookahead 推理加速框架 ，它能做到效果无损、即插即用，实现无痛推理加速，支持包括 Llama、OPT、Bloom、GPTJ、GPT2、Baichuan、ChatGLM、GLM 和 Qwen 在内的一系列模型。该算法已在蚂蚁大量场景进行了落地，实际获得了 2.66-6.26 倍的加速比，大幅降低了推理耗时。

下面，让我们一起来看看这个算法的加速效果和背后的原理。

加速效果

耳听为虚眼见为实，下面就以蚂蚁百灵大模型 AntGLM 10B 为例，实测一下面对同样的一个问题，生成答案的速度区别。

左图（当前主流算法）：耗时 16.9 秒，token 生成速度为 33.8 个/秒。

右图（蚂蚁 Lookahead 推理算法）：耗时 3.9 秒，token 生成速度为 147.6 个/秒，速度提升了 4.37 倍。

AntGLM 10B 模型	耗时	生成速度
主流算法	16.9s	33.8 token/s
Lookahead 算法	3.9s	147.6 token/s

接下来，再以开源的 Llama2-7B-chat 模型为例：

左图（当前主流算法）：耗时 15.7 秒，token 生成速度为 48.2 个/秒；

右图（蚂蚁 Lookahead 推理算法）：耗时 6.4 秒，token生成速度为 112.9 个/秒，速度提升了 2.34 倍。

Llama2-7B-chat 模型	耗时	生成速度
主流算法	15.7s	48.2 token/s
Lookahead 算法	6.4s	112.9 token/s

技术原理

当下的大模型基本都是基于自回归解码，每一步解码仅生成一个 token，这种操作方式既浪费了 GPU 的并行处理能力，也导致用户体验延迟过高，影响使用流畅度。

业内有少量的优化算法，主要集中在如何生成更好的草稿（即猜测生成 token 的序列）上，但是实践证明草稿在超过 30 个 token 长度后，端到端的推理效率就无法进一步提高，但是这个长度远没有充分利用 GPU 的运算能力。

为了进一步压榨硬件性能，蚂蚁 Lookahead 推理加速算法采用了 多分支的策略 ，即草稿序列不再仅仅包含一条分支，而是包含多条并行的分支，多条分支在一次前向过程中进行并行验证。因此可以在耗时基本不变的前提下，提高一次前向过程生成的 token 个数。

再进一步，蚂蚁 Lookahead 推理加速算法利用 trie 树存储和检索 token 序列，并将多条草稿中相同的父节点进行合并，进一步提高了计算效率。为了提高易用性，trie 树的构建不依赖额外的草稿模型，只利用推理过程中的 prompt 及生成的回答进行动态构建，降低了用户的接入成本。

上面的专业术语有点多，非专业的读者理解起来有点困难。换个通俗易懂的说法，再解释一下：

原来的 token 生成过程，就像早期中文输入法，只能一个字一个字“敲”出来，采用了蚂蚁的加速算法后，token 生成就像联想输入法，有些整句可直接“蹦”出来。

最后

蚂蚁 Lookahead 推理加速算法在检索增强生成（RAG）场景及公开数据集进行了测试。

在蚂蚁内部的 RAG 数据集上，AntGLM10B 模型的加速比达到 5.36，token 生成速度 280 个/秒；
在 Dolly15k 及 GSM8K 数据集上, 多个开源模型都有 2 倍以上的加速比，与此同时，显存增加和内存消耗几乎可以忽略。

该算法现已在 GitHub 上开源，相关论文也已公布在 ARXIV。感兴趣的同学可以阅读下相关论文了解更多技术细节，运行下源码查看效果。

论文地址 ： https://arxiv.org/abs/2312.12728

代码仓库： https://github.com/alipay/PainlessInferenceAcceleration

- END - 👆 关注「HelloGitHub」收到第一时间的更新 👆