苹果正在与英伟达合作,想让AI的响应速度更快
近日,苹果与英伟达宣布合作,旨在加速和优化大语言模型的推理性能。为了改善传统自回归LLM推理效率低和内存带宽小的问题,今年早些时候,苹果机器学习的研究人员发布并开源了一种名为ReDrafter的推测解码技术。目前,ReDrafter已经整合到英伟达的可扩展推理方案TensorRT-LLM当中,后者是基于TensorRT深度学习编译框架的专为优化LLM推理而设计的开源库,支持包括Medusa等推测解码方法。不过,由于ReDrafter所包含的算法使用了之前从未用过的运算符,因此英伟达方面添加了新的运算符,或者公开了现有的运算符,大大提高了TensorRT-LLM适应复杂模型和解码方式的能力。据悉,ReDrafter推测解码通过三个关键技术来加速LLM的推理过程:RNN草稿模型动态树注意力算法知识蒸馏训练RNN草稿模型是ReDrafter的核心组件。它使用循环神经网络,基于LLM的“隐藏状态”来预测接下来可能出现的tokens序列,其能够捕捉局部的时间依赖性,从而提高预测准确性。这个模型的工作原理是:LLM在文本生成过程中首先生成一个初始token,然后RNN草稿模型利用该token和LLM的最后一层隐藏状态作为输入进行束搜索,进而生成多个候选tokens序列。与传统自回归LLM每次只生成一个token不同,通过RNN草稿模型的预测输出,ReDrafter能够在每个解码步骤生成多个tokens,大大减少了需要调用LLM验证的次数,从而提高了整体的推理速度。图源:arXiv动态树注意力算法则是一种优化束搜索结果的算法。我们已经知道,在束搜索过程中会产生多个候选序列,而这些序列往往存在共享的前缀。动态树注意力算法会识别出这些共享前缀,并将它们从需要验证的tokens中去除,从而减少LLM需要处理的数据量。某些情况下,该算法能将需要验证的tokens数量减少30%到60%。这意味着使用动态树注意力算法后,ReDrafter能够更高效地利用计算资源,进一步提高推理速度。知识蒸馏是一种模型压缩技术,它能够将一个大型、复杂的模型的知识“蒸馏”到一个更小、更简单的模型中。在ReDrafter中,RNN草稿模型作为学生模型通过知识蒸馏从LLM中学习。具体来讲,蒸馏训练过程中,LLM会给出一系列下一个可能词的概率分布,开发人员会基于这个概率分布数据训练RNN草稿模型,然后计算两个模型概率分布之间的差异,并通过优化算法使这个差异最小化。在这个过程中,RNN草稿模型不断学习LLM的概率预测模式,从而在实际应用中能够生成与LLM相似的文本。通过知识蒸馏训练,RNN草稿模型更好地捕捉到语言的规律和模式,从而更准确地预测LLM的输出,并且因为其较小的规模和较低的推理计算成本,显著提高了ReDrafter在有限硬件条件下的整体性能。图源:阿里云开发者社区苹果的基准测试结果显示,在NVIDIA H100 GPU上对数十亿参数的生产模型使用集成了ReDrafter的TensorRT-LLM时,其贪心解码每秒生成的tokens数量提高了2.7倍。此外,在苹果自家的M2 Ultra Metal GPU上,ReDrafter也能实现2.3倍的推理速度提升。苹果的研究人员表示“LLM越来越多地用于驱动生产应用程序,提高推理效率既可以影响计算成本,也可以降低用户端延迟”。图源:Apple