点击蓝字 关注我们
一行代码让推荐大模型推理加速 1.5~2.5 倍
推荐大模型常用Beam search来生成top-K个item。然而，由于自回归解码的特性，标准的LLM推理是非常低效的，并且其时耗成本几乎随 beam size 线性增长。为了加速推荐大模型推理，我们提出了一种无损加速beam search 的方法——AtSpeed，并开源了相关 Python 包，使用起来非常方便，只需 import 再加一行代码即可 (Python包的使用说明附在文末)。
该方法针对 beam search 开发了相应 speculative decoding 算法并实现了 tree attention 算法。其算法不仅能应用于生成式推荐，还有潜力应用在LLM推理（如CoT 推理）等需要 beam search 的多种场景。
论文题目：Efficient Inference for Large Language Model-based Generative Recommendation
Arxiv链接：https://arxiv.org/abs/2410.05165
代码仓库：https://github.com/transcend-0/BeamSD
生成式推荐大模型中的 Beam Search
在生成式推荐中，模型需要用 beam search 自回归逐步解码得到 top-K 个推荐物品。然而，自回归解码需要串行执行，本身就是非常低效的，再加上 beam search 使解码的 sequence 数量增长为 beam size 倍，使资源消耗和时间开销几乎随 beam size 线性增长。因此，如何提高 beam search 的推理效率成为生成式推荐的一个重要问题。
Speculative Decoding 从 N-to-1 到 N-to-K
Speculative Decoding 是一种无损加速 LLM 推理的方法，其使用 draft-then-verify 的思想充分利用 GPU 的并行能力。简单来说，先用一个小模型快速生成多个 draft token，然后用大模型一次并行验证这些 draft token，来达到大模型一次生成多个 token 的效果，从而加速推理。
传统的 speculative decoding 通常用于 NLP 任务，遵循 N-to-1 验证（图1），即每一步生成 N 个 draft token，从中验证通过 1 个 target token。而在推荐任务中，需要通过 beam search 生成 top-K 个 item（即 K 个不同的 token sequence），这就引出了 N-to-K 验证（图2）。此外，beam search 下的 N-to-K 的验证目标不是简单地要求通过 K 个 token，而是对于每个验证步骤，要求通过 K 个 token sequence，这样才能保证从 draft 中得到的 top-K 个 token sequence 和 target LLM 的标准 beam search 得到的 top-K 个 token sequence 等价。
论文对 beam search 的 greedy search 和 sampling search 都开发了相应的 speculative decoding的算法，在greedy beam search上实现了生成结果上的无损加速（～2倍），并在sampling-based beam search上以非常微弱的性能损耗换取更强的加速性能 (最高接近2.5倍)。
Tree Attention
Speculative decoding 从 N-to-1 到 N-to-K 带来的另一挑战是验证的效率问题。由于 N 必须大于 K，而 K（target beam size）通常也比较大，例如 K=10，为提高通过率，N 可能需要设置为 40（在传统的 N-to-1 验证中，K=1，N 只需要设置一个较小的数例如4即可），这导致要验证的 token sequence 数量大大增加。我们观察到，这些 token sequence 的前缀有很多是共享的，因此可以用基于树的注意力（Tree-based Attention）机制加速计算。这一机制不只可以用于 speculative decoding，在普通的 beam search 上也可以使用，实现加速。
N-to-K 验证下 draft model 与 target LLM 的对齐训练
论文不仅开发了 beam search 的 speculative decoding 验证算法，还给出了在 N-to-K 验证下，分别针对 greedy search 和 sample search 的 draft model 与 target model 的对齐方法（AtSpeed-S 和 AtSpeed-R）。在该对齐方法下，greedy 和 sample 的 beam search 都能达到 2 倍以上的加速效果（图4）。
蓝色的 WS 是 walltime speedup，绿色的 AS 是 accept step）。
Python 包的使用方法
论文开源的工具包使用起来非常简单。
比如使用 speculative decoding 版本的 beam search，直接传入 target model, draft model 和 inputs 即可。
from atspeed.beamsd import beam_search_by_SD
outputs = beam_search_by_SD(target_model, draft_model, inputs)
如果不使用 draft model，仅启用 tree attention，只需 import 再加一行代码即可。如下，替换掉原本 transformers model 实例中的 beam search。
from beamsd import replace_beam_search_with_TreeAttn
model = replace_beam_search_with_TreeAttn(model)
然后就可以和平时一样使用 model.generate 了。
outputs = model.generate(**inputs, max_new_tokens=32, num_beams=5)
仅启用 tree attention，在 beam size=20 时，推理速度也可以提升超过 1.5 倍（图5）。
更多细节请参考论文原文和 GitHub 仓库。
Arxiv：https://arxiv.org/abs/2410.05165
GitHub：https://github.com/transcend-0/BeamSD
END
