ICWS 2024 _ 基于生成长度预测的大语言模型推理请求调度

随着技术的快速迭代，大语言模型（Larage Langugage Model, LLM ）在各种场景下都展示出强大的文本处理能力，越来越多的业务期待通过接入大模型服务，提升业务效果。区别于传统RPC请求服务时间相近，大模型请求服务时间受输出长度影响差异明显，同时每个请求所需的推理资源以及推理时间都无法事先感知，导致传统请求调度方案面临以下两个问题：（1） 当一个批次中请求的生成长度不同时，生成长度较短的请求需要等待生成长度较长的请求完成后才能一起返回，造成了计算浪费，影响了推理速度。（2）由于具有较长的生成长度的请求会产生更多的键值缓存，会占用更多的 GPU 显存。在不知道请求生成长度的情况下，静态批处理总是使用一个较小的批次规模（Batch Size）来避免显存溢出（Out of Memory, OOM）错误，无法充分利用GPU的计算能力。

本文尝试从请求调度的的角度提高 LLM 的推理性能，提出面向 LLM 推理的请求调度系统Magnus 。它通过对请求的生成长度进行预测，将生成长度相似的请求放在同一个批次（Batch）中进行处理，来降低计算浪费并增大批次规模，从而降低请求响应时间并提高大模型推理的吞吐量。实验表明，Magnus 可以将响应时间降低 89.7%，请求吞吐量提高 234%。在这项工作中，我们显著提高了静态批处理（Static Batching）的吞吐量，在未来，我们将进一步探索基于生成长度预测的请求调度方案在持续批处理（Continuous Batching）中的应用。

该工作目前已经被服务计算领域会议 IEEE International Conference on Web Services （ICWS） 录用，技术细节可以查看预印版： https：//arxiv.org/abs/2406.04785

研究背景

随着参数规模的增加，基于 transformer 的语言模型在各种自然语言处理（NLP）任务上都表现出强大的能力，这可以使许多应用程序受益。然而，由于成本过高，大多数应用程序开发人员无法负担训练和部署 LLM 的高昂费用。因此，诸如 OpenAI、谷歌和阿里巴巴等人工智能领域的科技公司将他们的 LLM 作为服务发布，并允许开发人员通过 API 访问，即语言模型即服务（Language Mode as a Service, LMaaS）。

如图 1 所示，在 LMaaS 场景中，应用程序将用户输入的文本附上指令作为请求，发送给 LLM 进行处理。例如，VSCode上的代码助手插件可以通过在用户的代码前加上指令“Fix bugs in the following code：”作为请求发送给 LLM 服务从而实现程序漏洞修复功能。在服务端，来自不同应用程序的提示被混合在一起然后分批，并由部署在图形处理器（Graphics Processing Unit, GPU）等加速硬件上的 LLM 实例进行批处理。

由于 LMaaS 场景中的应用，如机器翻译和程序修复，更多地关注生成质量而不是多样性。因此，通常采用贪婪采样和波束搜索（Beam Search）[1]。考虑到波束搜索的计算开销很大， LLM 往往采用贪婪采样方式生成文本，因此相同的请求的生成结果总是相同的。

现有的深度神经网络推理系统，如Tensorflow serving [2] 和Triton Inference Server [3] 利用先到先服务的方式以固定的批次规模处理请求，在本文中我们将这种调度策略称为朴素调度（Vanilla Scheduling），当使用静态批处理（如图2所示）时，朴素调度会导致两个严重影响批处理效率的问题。首先，当一个批次中请求的生成长度不同时，生成长度较短的请求需要等待生成长度较长的请求完成后才能一起返回，造成了计算浪费，影响了推理速度。其次，由于具有较长的生成长度的请求会产生更多的键值缓存，会占用更多的 GPU 显存。在不知道请求生成长度的情况下，静态批处理总是使用一个较小的批次规模来避免显存溢出 （Out of Memory, OOM）错误。因此，GPU 强大的并行计算能力无法得到充分利用，降低了系统总体吞吐量。

研究动机

很多应用的请求生成长度和用户输入呈现强正相关性

我们发现在 LMaaS 场景下，有许多流行应用，其请求生成长度与原始用户输入文本长度呈正相关。例如机器翻译（Machine Translation, MT）、语法纠错（Grammar Correction, GC）、文本去毒（Text Detoxification, TD）、代码翻译（Code Translation, CT）、漏洞修复（Bug Fixing, BF）和代码注释（Code Comment, CC）。为了证实这一结论，我们从现有数据集中为每个应用构造了 2000 个请求，并将这些请求使用 ChatGLM-6B[4]、Qwen-7B-Chat[5] 和 Baichuan2-7B-Chat[6] 三个 LLM 进行处理。我们将请求生成长度和用户输入长度进行可视化展示（如图 3）。可以发现，在这些应用中用户输入文本的长度与生成文本的长度具有显著的正相关性。

表I中展示了用户输入长度和请求生成长度的皮尔逊系数。从中我们可以看出，对于三种大型语言模型 ChatGLM-6B, Qwen-7B-Chat 和 Baichuan2-7B-Chat，大多数应用请求的用户输入长度和请求输出长度的皮尔逊系数都大于0.8，表明存在强烈的正相关性。因此，对于这些应用，用户输入长度可以极大地帮助预测请求的生成长度。

生成长度预测可以显著提高 LLM 推理性能

图4：朴素调度 vs. 基于生成长度预测的调度方案Magnus

为了初步验证基于长度预测方案的可行性，我们在 NVIDIA V100 32GB GPU 上部署了一个 ChatGLM-6B LLM 实例进行实验。在实验中，请求长度和生成长度均为 1000 左右的 "大 "请求和请求长度和生成长度均为 10 左右的 "小 "请求按 图 4（a） 所示的顺序到达。在实验中，我们利用 huggingface-transformers 作为推理引擎来加载和运行 LLM 。如图 4（b）所示，朴素调度按照达到的顺序对请求进行批处理，批次规模固定为 7，总处理时间为 242 秒。然而，基于生成长度预测的调度方案 Magnus ，将小请求和大请求分别分为两个批次。由于小请求占用的显存少，因此可以使用更大的批次规模，从而充分发挥 GPU 的计算能力。如图 4（c） 所示，大请求和小请求的批次规模分别为 18 和 3，总处理仅仅时间为 60s，大大优于朴素调度。这说明生成长度预测可以显著提高 LLM 推理性能。

系统设计

基于上述的初步数据分析和实验探索，我们提出了一种基于生成长度预测的请求调度系统。Magnus 包含四个核心组件：（1）请求生成长度预测器；（2）适应性批次组装器根据请求长度预测的结果将请求分为不同批次；（3）批处理时间估计器（4）最高响应比优先批调度器根据估计的批处理时间为多个批次确定处理顺序。通过这四个核心组件相互配合，Magnus 可以提升 LLM 批处理的请求吞吐量，并降低请求的响应时延。

Magnus 的工作流程如图5所示。当请求到达时，生成长度预测器预测它的生成长度，并将①请求及其预测结果发送给适应性批次组装器，适应性批次组装器将②该请求插入到具有相似长度和生成长度的批次中，以减少计算浪费。之后，批处理时间估计器将根据③批次的批次规模，长度、预测的生成长度来预测其推理时间。当一个 LLM 实例完成处理，变得空闲时，批次调度器将使用最高响应比优先算法，根据④估计的批处理时间从队列中选择一个⑤批次调度到 LLM 实例进行处理。除此之外，Magnus 周期性地利用新收集的⑥请求信息（如请求及其实际生成长度），以及从日志数据库中收集的⑦批次信息（如批次规模、批次中请求的长度、批次中请求的生成长度，和批处理时间），通过持续学习技术来不断提高生成长度预测器和批处理时间估计器的精度。接下来，我们针对每个核心组件详细介绍对应的设计细节。

生成长度预测器

生成长度预测器由一个基于 LaBSE [7] 的语义特征提取模块、一个压缩模块和一个随机森林回归器组成。图6展示了生成长度预测器的架构图，首先，语义特征提取模块将用户输入和指令作为输入，分别提取用户级和应用级语义特征，并生成两个嵌入向量 $v_{user},v_{app}\in {\mathbb{R}}^d$，其中$d=768$。其中应用级别的语义特征用于帮助回归器识别不同应用任务，从而为每个任务学习请求输入和输出的相关性。除此之外，由于语义相似的请求往往有相似的输出，因此具有相似的生成长度，我们利用用户级的语义特征来帮助回归器利用用户请求的语义相似性来提高预测精度。为了控制回归器的复杂性，压缩模块将$v_{user}$和$v_{app}$通过分组压缩的方式进行压缩，首先将它们平均分为 $d_{user}$和 d a p p d_{app} d