Accelerating Production LLMs with Combined Token/Embedding Speculators翻译-CSDN博客

摘要

本技术报告描述了新型推测解码草稿模型的设计和训练，旨在加速生产环境中大语言模型的推理速度。通过同时以上下文向量和采样 token 作为草稿预测的条件，我们可以训练推测器高效地预测高质量的 n-gram，然后由基础模型接受或拒绝这些 n-gram。这使我们能够在每次推理前向传递中有效地预测多个 token，从而将高度优化的基础模型实现的实时推理速度提高 2-3 倍。我们将探索这些初步成果，并描述进一步改进的后续步骤。

1.介绍

随着生成式大语言模型 (LLM) 规模不断扩大，其能力也随之增强。然而，这种为了提升表达能力而增加的模型规模，却给生产环境的部署带来了挑战：大语言模型虽然功能强大，但也会产生巨大的内存和计算开销。像 Llama2-13B 这样的知名模型包含 130 亿个参数，使用典型的 16 位权重表示，大约占用 24 GB 内存。在推理阶段，这些巨大的计算量完全用于每次生成一个 0 到 32K 范围内的 token 索引。因此，在极端情况下，我们产生 24 GB 的开销，却只生成不到 2 GB 的信息！
　　纠正这种不平衡的一个显而易见的方法是一次预测多个 token。事实上，经典的自然语言处理理论已经证明，即使是简单的 2/3 元语法语言模型也具有很强的预测能力，这告诉我们，学习到的语言模型应该能够以合理的准确度一次预测多个 token。然而，我们也希望避免修改基础模型本身，因为众所周知的模型（例如Llama-13B）是已知量，具有理想的特性，在可预测性和可靠性至关重要的生产环境中不能冒险改变。
　　解决这个问题的一种方法是推测解码，即训练一个较小的草稿模型或推测器，使其能够根据给定的输入序列预测多个 token。这些推测 token 的生成成本较低，但准确率低于基础 LLM。然而，我们可以在 LLM 正向传递过程中利用 GPU 并行性，以最小的额外开销评估每个新 token 的输出。然后，通过将输出与推测的输入进行比较，我们可以接受所有与基础模型输出匹配的预测 token，同时拒绝所有不匹配的 token。通过这种方式，我们可以在每次 LLM 正向传递中以最小的额外成本预测多个 token。关于推测解码的更深入解释，请参阅 [3, 6]。
　　本文介绍了 IBM 研究院近期的研究成果，该研究利用新的架构和推测解码来加速生产级 LLM 推理。先前的方法使用单独的 LLM 作为草稿模型，消耗整个 token 输入序列，而 Medusa 表明，一个简单的多头 MLP 架构即可达到相同的性能。该架构仅使用来自最新 token 的 base 模型的嵌入向量，因为该嵌入向量隐式地包含了从前一个输入派生出的上下文信息。我们基于 Medusa 提出了一种改进的架构，其贡献如下：

我们表明，除了基础模型上下文向量之外，通过对采样 token 进行调节，可以大大提高推测器的输出质量。
我们引入了一种有效的两阶段训练方案，首先将我们的推测器与基础模型输入行为相一致，然后再与输出行为相一致。
使用这个推测器训练管道，我们将四个高度优化的生产 LLM 加速了 2-3 倍。
我们探索了推测解码在生产环境中的局限性，结果表明，随着基线计算和效率水平的提高，承诺的加速会减少。
我们概述了下一步措施和进一步调查的领域。

我们的工作代码是开源的，可以在 https://github.com/foundation-modelstack/fms-fsdp 上找到，我们的 13Bparameter 基础模型的推测器已在 HuggingFace 上发布：https://huggingface.co/ibm-fms。

2.Related Work

推测解码的原始实现依赖于一个规模较小、独立的 LLM 来为更大的基础模型提供草稿。虽然这种方法可行，但也存在一些实际限制。（1）首先，由于消耗了整个过去的 token 历史记录，即使是小语言模型也会迅速膨胀，导致开销巨大。（2）其次，较小的模型是基于真实文本进行训练的，因此与更大的基础模型的行为不一致。因此，草稿 LLM 可能会生成后缀候选，这些后缀候选虽然反映了真实训练数据，但相对于基础模型的行为（而这才是我们最终想要复现的行为）而言并不准确。（3）第三，如果没有使用相同领域、相同 tokenizer 进行训练的小型 LLM，那么从头开始训练它可能会耗费过多的时间和精力。
　　最新的 Medusa 模型旨在通过单个多头 MLP 生成草稿来解决这三个问题。该 MLP 将基础模型最新 token 预测的最终嵌入向量作为输入。这使得它无需完整的 token 历史记录即可访问语义上下文，因为基础模型已经为该 token 向量进行了上下文化，这使得 Medusa 能够快速推理任意长度的提示。Medusa 的训练也相对快速简便，因为它是一个带有丰富语义输入的简单模型。
　　我们发现 Medusa 架构的一个主要缺陷：由于仅使用基础模型中的上下文向量，Medusa 未能对采样的 token 选择进行条件化，导致 n-gram 预测质量不佳。EAGLE、Hydra 和 Recurrent Drafter 是并行运行的模型，它们各自独立地得出了相同的结论。这验证了我们的推测器设计。与这些模型相比，我们专门考察了此类模型在高度优化的生产环境中的性能。具体而言，未加速的基础 7B 模型每秒产生的 token 数量比 EAGLE 案例研究中等效 16 位精度基线模型（每秒 94.9 个 token 对比每秒 55.1 个 token）高出 70%，比 Hydra 基线模型（每秒约 50 个 token）高出约 90%，比 Recurrent Drafter 模型（每秒 35.6 个 token）高出 170%。这带来了截然不同且更有趣的部署和评估前景。

3.Speculator Design and Training

在本节中，我们描述我们的推测器架构，然后讨论有效的训练方案。

3.1 Speculator Design

在这里插入图片描述
　　我们从 Medusa 架构入手，实现了三个并行的 MLP 预测头，它们使用用于预测 token $n$ 的基础模型的嵌入向量来预测 token $n + 1, n + 2, n + 3$ 。然而，在观察预测结果后，我们很快发现了一个令人担忧的行为：该 Medusa 架构预测的候选序列通常是非常差的三元组。这是因为 Medusa MLP 仅使用 token n 的基础模型嵌入向量作为其唯一输入，而完全忽略了从该向量生成的分布中采样的实际 token。这意味着，Medusa 对 token $n + 1, n + 2, n + 3$ 的预测实际上根本不以 token $n$ 为条件，而是同时以 token n 所有可能样本的期望为条件。同样，token $n + 2$ 的预测也不以 token $n + 1$ 为条件，而是以 $n + 1$ 所有值的期望为条件，token $n + 3$ 的预测也类似。这导致了违反直觉的行为。
　　举一个简单的例子，假设基础模型生成了一个以“follow these steps:\n”结尾的输入。假设该输入最可能的三种延续方式如下（无需过多的怀疑）：
　　• “Step 1:”
　　• “\n Step 1:”
　　• “\n \n Step 1:”
　　此外，在每一个延续中，“Step” 出现的可能性都很高（因为我们知道它必然会出现），而“\n” 的置信度可能较低。在这种情况下，Medusa 对接下来三个 token 的预测是“Step Step Step”，因为这显然是 $n + 1, n + 2, n + 3$ 位置上最有可能出现的三个 token。然而，这显然也无法很好地预测基础模型的行为。
　　为了缓解这个问题，我们将 Medusa 架构修改为多阶段多层感知器 (MLP)，其中每个阶段预测草稿序列中的单个 token。MLP 的每个阶段都以状态向量和前一阶段采样 token 的一个（或多个）嵌入作为输入（基础模型可视为阶段 0）。每个输入对（状态向量加上采样 token 嵌入）被投影并通过 LayerNorm/GeLU 激活函数，形成一个新的状态向量。该状态向量用于预测下一组草稿 token，该组草稿 token 与新的状态向量一起作为下一预测阶段的输入。因此，该模型会生成一棵完整的候选 token 树，我们允许用户从该树中选择 top-k 个最有可能的预测值用于实践。图 1 展示了该 3 阶段模型架构的可视化效果。
　　在推理阶段，我们实现了多候选解码和树状注意力机制，以便并行评估选定的 k 个候选序列。我们并非针对单个候选序列比较输入 token 和输出 token，然后仅提取正确的 token，而是同时对所有候选序列执行此评分过程，并仅保留得分最高的候选序列的 token。这是一种提高推测器准确度（从而增加每步生成的 token 数量）的简单方法，但代价是增加了并行度。因此，当我们报告每步生成的 token 数量结果时，我们也会报告并行评估的候选序列数量。
　　通过对先前的 token 以及基础模型的内部状态进行条件调整，我们的推测器能够生成更贴近现实的 n-gram 草稿。同时，它在预测基础模型行为方面也更加准确。对于 Llama2-13B，我们只需根据先前采样的 token 以及基础模型状态对每个 MLP 头进行条件调整，就能将上述三头沙盒 Medusa 架构的速度从每步 2.28 个 token 加速到每步 2.63 个 token（取 10 个候选结果中的最佳值）。
　　请注意，图 1 可以很容易地描绘一个 RNN，实际上，我们可以通过简单地绑定不同阶段的权重，将多阶段 MLP 实现为 RNN。虽然这种权重共享会减少推测器的参数数量，但它不会在推理时提高推测器的速度。因此，为了最大限度地提高适应性，我们在实现推测器时没有使用任何权重共享。因此，最终的参数数量看起来很大（对于一个 7B-13B 规模的基础 LLM，参数数量为 800M-1.1B）。尽管我们的推测器的参数数量约为基础 LLM 的十分之一，但由于其深度较浅且架构简单，其执行时间不到基础 LLM 的三十分之一。

3.2 Speculator Training

为了训练我们的推测器，我们利用了 [2] 中介绍的高效的 Foundation Model Stack 分布式训练平台。我们调整了 Llama 训练脚本，不再冻结基础 Llama 模型，而是在其上训练推测器。我们保留了相同的数据加载器、检查点实用程序和其他便捷的流水线组件。
　　由于推测器是从零开始训练的，我们有机会将其与基础模型完全对齐：推测器无需预测真实文本，而是直接根据生成的基础模型输出进行训练。遗憾的是，使用 Llama 进行文本生成无法高效并行化——这正是我们训练推测器的初衷。因此，我们采用了两阶段训练范式。在第一阶段，我们使用小批量、长上下文长度（4k 个 token）进行训练。基础模型并行为每个 token 生成嵌入向量，推测器使用这些向量以及真实 token 来预测后续的真实 token。这项任务与标准因果语言模型训练类似，只是我们在每个位置预测多个 token（推测器 MLP 的每个阶段一个 token）。
　　在第 2 阶段，我们切换到由基础模型本身生成的短序列长度（256 个 token）的大批量训练。通过仅生成短的 token 序列，我们可以继续以相当高效的方式并行训练，同时对推测器进行微调以使其精确匹配基础模型的输出。我们没有观察到任何对短序列长度的持续过拟合——事实上，在某些情况下，推测器在较长的输入上表现得更准确。在我们的设置中，第 1 阶段和第 2 阶段的步长比大约为 5:2，相当于实际时间大约为 1:1，因此为了方便起见，我们使用这个比例。推测器一半时间在真实文本数据上进行高效训练，另一半时间则用于效率较低的模型对齐。
　　在第二阶段训练结束时，我们能够将三头 Llama2-13B 推测器的速度从第一阶段结束时的每步 2.41 个 token 加速到 2.63 个（10 个候选中最佳）。图 2 给出了两个训练阶段的训练曲线。