Tabi：一种适用于大语言模型的高效多层推理系统

23年5月来自香港科技大学、中科大和福州大学的论文“Tabi: An Efficient Multi-Level Inference System for Large Language Models”。

当今的趋势是构建越来越大的大语言模型 (LLM)，同时推动自然语言处理的性能，但这显著增加了推理阶段的延迟。由于向 LLM 添加参数的收益递减，对于大多数查询，较小的模型可以做出与昂贵的 LLM 相同的预测。Tabi，一个具有多级推理引擎的推理系统，它使用小型模型和可选的 LLM 为要求苛刻的应用程序提供查询服务。Tabi 针对服务框架中的判别模型（即非生成 LLM）进行优化。Tabi 使用标定后的置信度得分来决定，是否以极快的速度返回小型模型的准确结果，或将它们重新路由到 LLM。对于重新路由的查询，它使用基于注意的词修剪和加权集成技术来抵消系统开销和准确性损失。

LLM的资源开销

2018 年发布的 BERT [18] 令人印象深刻地提高了由 9 项任务组成的 GLUE（通用语言理解评估）基准 [72]，从 69 分提高到 80 分。此后，BERT 提出的掩码语言建模 [18] 和其他自监督预训练目标，使得使用数十亿个单词组成的大规模数据集训练成为可能。

从海量数据中学习以充分理解语言需要大规模模型。例如，具有 3.4 亿个参数的 BERT-large 是提出时最大的 DNN 之一。Switch Transformer [24] 使用超过一万亿个参数在多个认知任务上取得 SOTA 结果。近年来，具有注意机制 [69] 和与 BERT 类似架构的 LLM 已经超越人类基线，并将 GLUE 排行榜推高至 90 以上。

另一方面，实现 SOTA 性能需要越来越多的计算资源，而回报却在减少 [61]。例如，SOTA Transformer 模型花费了 50% 的训练空间才实现了 0.3 的错误率降低 [53]；最近的 DeBERTa [35] 使用的参数比 ALBERT [43] 多 36 倍，但准确率仅提高了 6%。GPT-3 [13] 拥有 1750 亿个参数，是其前身的 100 倍，这表明当前的 LLM“最终可能会达到预训练目标的极限”。

如图 a 显示 5 个 DNN 在文本分类中的准确率和大小（从小到大）。可以看到，即使添加更多参数（条形高度差异），准确率相对于以前模型的改进（条形宽度）也变得不那么明显。对于要求不高的任务，尤其是在边缘设备上，当小得多的模型就足够时，不需要这么大的 DNN [27, 70]。在使用当前的 DNN 推理系统实现其 SOTA 性能时，开发更大模型的 NLP 现状将导致巨大的资源开销。为了满足苛刻的目标（垂直虚线），推理系统从可用的训练模型中选择 RoBERTa-large，其准确率为 96.3%，参数为 3.55 亿。工作负载的资源使用量通常相当于整个图形的面积，使用所选 DNN 的容量（高度）处理所有查询（宽度）。然而，在事后分析结果时，其他更高效的模型可以对超过 90% 的查询做出与 RoBERTa-large 一样准确的预测（例如，在图 a 中第二小模型 DistilBERT 的准确率为 91%），尽管它们都没有达到被选中的准确率目标。

理想情况下，如果小模型可以正确处理查询，就不应该调用更大的模型。这种场景类似于通过图 b 中从小到大的模型多级反馈队列 (MLFQ)，同时知道到达正确级别的捷径（虚线箭头）。因此，工作负载的最佳资源使用率是图 a 中所有条形图的右下角区域，这仅占实际使用率的 11%！将简单查询定义为无论使用小型或大型 DNN 都能获得相同结果的查询。而左侧灰色区域是不必要地使用过于强大的 LLM 提供简单查询的开销。

实际上，事先并不知道查询的最低要求，因此没有通往最高效模型的捷径；此外，通过多个模型会增加额外的延迟。这就是为什么当前的推理系统在应用程序级别匹配性能目标，而专注于资源/模型扩展的原因。然而，延迟减少、成本节省和环境影响的巨大潜力，在无模型系统中优化每个单个查询的 LLM 推理。

DNN 推理系统

最近的推理系统 [31, 57] 让用户仅指定高级性能要求（例如，情绪分析任务的准确性和延迟），而不是特定模型，并部署经过训练的模型而无需掌握 ML 或系统优化，这称为无模型推理。系统首先从注册的模型中选择适当的资源（例如，GPU 或 CPU）和 DNN（例如，针对准确性或效率进行优化）。然后，它将查询路由到选定的模型，运行 DNN 推理并返回结果。

当前的推理系统进行了各种优化。INFaaS [57] 选择针对不同设备和批次大小自动优化的最佳单一模型进行推理，并快速适应工作负载的变化。Cocktail [31] 采用集成学习和多个小模型，通过并行执行减少延迟；它还动态调整模型集成以最大限度地降低成本。

SOTA 推理系统为应用程序设定性能目标，并选择模型来均匀地服务于整个工作负载，而不检查所选模型是否对单个查询做的过火。尽管 Cocktail [31] 有一个使用真实标签的监控过程，这在现实世界的任务中很少见，但它无法实现细粒度的模型选择。当 DNN 性能随模型变大而显着提高（即每个查询的开销更少）时，这种一刀切（one-model-fits-all）的设计对于许多任务来说并不是一个令人担忧的问题。

NLP 应用程序对 DNN 容量的需求存在明显的差异：自然工作负载是难和容易样本的混合，许多容易样本可以通过小得多的语言模型正确服务。机器学习研究一直在探索训练和推理设置中的差异：难样本挖掘 [37, 63] 在复杂数据上花费了更多的训练资源，而自适应预测 [10, 68, 80] 开发了专门的 DNN，这些 DNN 针对简单数据只运行早期层。

Tabi

考虑到当今无模型推理系统为 LLM 提供服务的巨大开销以及注意机制带来的新优化机会，Tabi 提出一种多级推理引擎，有效地减少由实时查询反馈驱动的通用模型细粒度延迟。

如图说明了 Tabi 的工作流程和组件。模型存储库，存储用户提交和维护的训练 DNN。Tabi 进行离线分析获取单个模型和各种多级 DNN 配置（即候选模型）的统计数据。对于具有性能目标的新任务，控制器选择最佳候选模型并在工作器上运行 DNN 推理。推理控制逻辑由三个模块组成，并管理工作器上使用的 DNN。

多级推理引擎

Tabi 的核心是一个多层级推理引擎，它采用高效且高精度的 DNN 来为异构 NLP 查询提供相应的资源，如图所示。除了 DNN 之外，推理引擎还包括 (1) 基于概率的调度器（dispatcher），它使用标定的置信度分数快速返回准确的预测；(2) 使用现有的注意权重对输入数据进行修剪，加速重新路由查询的 LLM 推理，并减少系统开销；(3) 激活模型的一个加权多级集成，进一步提高最终准确性。

在处理复杂自然语言的应用中，DNN 不仅应该准确，还需要指出何时可能不正确。在 Tabi 中，过早返回小模型的错误结果可能会抵消 LLM 的任何优势，并有违反准确性目标的风险。一个概率调度程序，决定何时过早返回输出以及何时继续推理。获得预测置信度的一个直接方法是使用 softmax 概率，其总和为 1。流行的 Hugging Face 部署流水线 [22] 及其模型库都使用了这种方法。

然而，包括语言模型在内的现代 DNN 都“过于自信”[17, 32]。在测试中，小型模型不断做出 softmax 概率超过 0.9 的预测，这高于它们的实际准确率；此外，即使是错误的预测也经常是用高 softmax 概率做出的。

如图左侧部分，将二元分类工作负载的结果按置信度得分（范围从 0.5 到 1）分组到五个桶中，并将其平均准确率显示为橙色图案条。红色对角线说明置信度得分与实际推理准确率相等时的理想关系。显然，它们不匹配，高 softmax 概率（即高于 0.8）并不能保证高精度，因此很难理解推理质量。

为了可靠地调度查询并平衡准确性和效率，使用温度缩放（scaling） [32] 来标定 softmax 置信度得分，以便它能够与实际推理准确性相匹配。结果，表达后的置信度得分（图中的蓝色条）与理想的准确度分布相匹配，为 Tabi 提供可靠的依据来决定何时快速返回高效模型级的正确结果。

利用标定的置信度分数，计算将查询重新路由到下一级的概率，其为标准化置信度、ReLU 和缩放 S 型激活函数的组合。它结合软概率和硬截断来平衡准确性（挑战性查询的高召回率）和延迟（重新路由的高精度）。上图右侧部分说明了这一点，截断阈值0.9，只两个分类。对于置信度低于截断阈值的查询，调度程序总是将它们重新路由到下一级模型。对于更确定的查询，使用缩放 S 形函数优先考虑那些接近 100% 置信度的查询，并调整置信度略高于阈值查询的重新路由概率，获得更好的准确性（实线）。与线性降低重新路由概率（虚线）相比，该修正方案可以尽早返回更多查询以降低延迟，并在返回的查询中实现更高的准确性。

虽然提前返回的概率（实线上方的区域）似乎有限，但查询的 CDF（绿色虚线）通常集中在高置信度一侧，因此潜力巨大。在评估中，调度程序可以在第一个有效模型之后返回大约 50% 到 70% 的查询。这是因为模型选择过程可以使用一个有能力且高效的模型作为第一级来分担 LLM 的工作量。

在调度之后，这些具有挑战性的查询仍然需要 LLM。使用超小型模型处理此类查询的一个担忧是，即使提高了平均和中位数性能，重新路由的查询也会增加尾部延迟。相反，小模型的结果可以减少延迟。利用小模型对查询的实时语义理解（即注意权重），可加速其 LLM 推理，方法是从输入数据中修剪不必要的单词并抵消多级延迟开销。

对于那些重新路由的查询，不会直接输出最终级 LLM 的预测，而是采用加权集成学习将其与前几级的预测相结合，以提高准确性。这种技术不会增加额外的 ML 计算，因为在多级结构中所有中间 softmax 输出都随时可用。

一个集成是一组模型，其各个决策与权重相结合以共同做出预测。事实证明，在现有的大型模型中添加独立模型（即使是弱模型）比单独使用单个大型模型更准确，因为可以减少方差和偏差 [40, 49]。

模型候选的配置和选择

要探索 Tabi 的优化空间，需要正确配置系统，例如，应该使用多少级和哪些模型，如何自信地尽早返回预测，以及一个词对于修剪有多不重要。抽象出一组 Tabi 配置作为具有准确度和延迟性能的模型候选。候选在逻辑上等同于其他 DNN 服务单元，例如 INFaaS [57] 中的单个模型和 Cocktail [31] 中的集成，这样 Tabi 就可以与现有的推理扩展系统兼容。

Tabi 包括离线分析和在线选择过程。分析会生成一系列不同优化的候选。当一个任务来了，Tabi 会选择一个综合性能以最低延迟满足目标的候选。Tabi 的候选分析和选择通常遵循 SOTA 推理流程 [31, 57]，同时适应 Tabi 的内部机制。

候选模型的性能由基础模型和 Tabi 超参共同决定。首先，用代表性数据集对已注册的单个模型进行分析，就像 INFaaS 所做的那样。保存它们的统计数据（例如准确率和延迟）和 Tabi 特定的中间变量，包括模型的 softmax 输出（#class-length 向量）、标定的置信度分数（浮点数）和注意机制的重要性权重（#token-length 向量）。这些变量增加的存储开销可以忽略不计（即每个模型几兆字节）。

接下来，采用提前停止（early stopping）提高效率来获得候选模型的综合性能。Tabi 使用已注册的模型自动生成一个可能的多级 DNN 组合列表，从 2 级开始，准确率按升序排列。为了消除不切实际的选择（例如，在所有级都放置 LLM 或使用太多模型），仅保留第 75 个尾部延迟在最慢模型延迟 1.5 倍以内的候选模型。对于每种组合，测试一系列超参值。对于调度器截止阈值，测试不同值不会调用额外的 DNN 计算，因为它们本质上是重新排列已保存的 softmax 结果。将阈值从 0.95（1 表示始终使用 LLM）开始减少，步长为 0.05，并在组合准确度低于所用模型的平均值时停止。将集成权重设置为与每个模型准确度成比例。由于词剪枝会修改输入，因此分析剪枝尺度仅对 LLM 需要额外轮的 DNN 推理。用代表性高效模型 DistilBERT 的已保存注意权重，并将修剪尺度从 0.5（0 表示无修剪）以步长 0.1 增加，生成修剪掩码，直到 LLM 的准确度损失超过 2%。这些设计减少了标准一次性分析的开销。最后，将 Tabi 候选模型的元数据保存在模型存储库中，就像单个模型一样。

借助已分析的性能，Tabi 按照现有标准 [57]，使用模型存储库中已保存的元数据，选择在线任务的候选模型。选择实现最低延迟的候选模型，满足任务的性能目标（准确度和延迟）。然后，控制器在工作机上启动选定的模型实例，服务于此应用程序的查询。工作重点是优化单个数据推理，同时仍然适用现有的资源和模型扩展系统。例如，Tabi 候选可以水平地扩展到多个实例，或者垂直地切换到不同的内部模型或参数，处理像单个模型一样的工作负载变化。