70B模型,秒出1000token,换算成字符接近4000!
研究人员将Llama3进行了微调并引入加速算法,和原生版本相比,速度足足快出了快了13倍!
不仅是快,在代码重写任务上的表现甚至超越了GPT-4o。
这项成果,来自爆火的AI编程神器Cursor背后团队anysphere,OpenAI也参与过投资。
要知道在以快著称的推理加速框架Groq上,70B Llama3的推理速度也不过每秒300多token。
Cursor这样的速度,可以说是实现了近乎即时的完整代码文件编辑。
有人直呼好家伙,如果把Cursor魔改后的Llama3放到Groq上,是不是每秒能跑出上万token了。
更是有人激动地说,在大模型领域,我们正在消除“延时”这一概念。
引入全新推理加速算法
作者此次设计的加速方法,主要是用来解决一种名为“Fast Apply”的任务,即对代码内容进行快速修改并应用。
首先需要说明的是,虽然说任务最终实现的效果是代码的局部修改,但是实际操作过程中,输出并非是只有变化的内容,而是直接全局重写。
这样做的原因,是团队在预先测试后做出的选择——他们发现,除了Claude-3-Opus,大多数模型在真·局部修改任务上的表现都不理想。
之所以会这样,主要有以下三个原因:
- 首先是直接重写时会输出更多token,使得有更多的前向传递来确定正确的解决方案。
- 其次,模型的训练数据也大部分都是完整代码,对局部修改相对陌生。
- 此外,大模型糟糕的数学运算也无法保证能在输出差异时正确处理行号。
(不过作者认为这仍然是一个有潜力的未来研究方向。)
确定了采用全局重写的方案后,Cursor团队使用了任务相关的数据对Llama3进行了微调。
所采用的数据有真实编辑数据与合成数据两大来源,按照1:4的比例进行了混合。
其中合成数据是指用GPT-4生成代码编辑的建议,然后用其他模型将这些建议“应用”到原始代码上。
为了提高数据集的质量,作者还对小文件、重复文件和无变化样本进行了下采样。
为了评估这些模型的表现,作者让它们处理了450个代码编辑任务(每个都不超过400行),并用Claude3-Opus对输出进行了打分。
最终,作者微调出的70B Llama3模型,表现几乎与Claude3-Opus-diff匹配,并且优于GPT-4-Turbo和GPT-4o。
至此的微调解决了性能问题,但不难看出此时的Llama3速度依然很慢,每秒只能输出不到300个字符(注意是字符,不是词也不是token)。
而让改写工作快到飞起的,还有另一项秘密武器。
针对代码改写任务,Cursor团队专门引入了一种名为预测性编辑(speculative edits)的算法。
这种方式用一种先验算法来对多个后续token进行预测,然后再用本体大模型进行验证,降低了大模型的调用次数,从而减轻了运算量。
这种先验算法来自于代码任务的一个特点——相比于其他文本,其词表更小,且语法结构、缩进规则等拥有更高的确定性,利用先验知识可以更精准预测未来的token。
这样的做法也与GPT-4和Meta有着共通之处——
传统的语言模型推理推理速度较慢的原因,主要是预测下一个token的过程通常是自回归的,即模型在生成每个token时,都要考虑之前生成的所有token。
为了降低运算量,以GPT-4为代表的大模型,使用了名为预测解码(speculative decoding)的加速算法,通过小的近似模型提前进行预测,然后再让本体大模型对预测结果进行验证。
Cursor和GPT-4的区别就在于,前者的小“模型”是一种更确定的算法,而后者只是模型规模减小,本质上仍是概率预测。
Meta这边则是推出了一次性预测多个后续token的算法,用n个独立的输出头并行预测n个未来token,结果发现在编程任务上表现尤其优异,原因是由于编程语言的逻辑结构更严谨,知识的内在联系更紧密。
当然,Cursor对这种特点利用更为充分,没有用注意力头,而是直接拿更确定的算法来做多token预测。
最终的结果就是,预测算法为70B的Llama3带来了近13倍的速度提升,而测评表现没有任何损失。
此外,作者还与企业AI模型基础设施平台fireworks.ai合作,利用其优化的推理引擎和定制化的硬件环境,进一步提高了模型的运行效率。
未来,团队还计划进行知识蒸馏,并把预测编辑算法迁移到更小的8B Llama3,并扩展到更多的编程语言和任务。
同时,对于Cursor团队研究过但并未采用的真·局部修改(Diff)算法,作者也计划进行改进。
One More Thing
在实验当中,作者不仅用预测算法加速了Llama3,也实现了对GPT4-Turbo的加速。
不过作者并没有介绍具体在GPT当中如何实现,而是留做了思考题,还搞了一场“有奖竞猜”。
能够正确解答的人将获得1个月的Cursor会员;如果能在vllm和TensorRT-LLM中实现预测加速,将分别获得半年和一年的会员。
如果你感觉有思路的话,不妨挑战试试(手动狗头)。