大语言模型的快速分布式推理服务

23年5月来自北京大学的论文“Fast Distributed Inference Serving for Large Language Models”。

ChatGPT 为代表的大语言模型 (LLM) 为新一代交互式 AI 应用程序提供支持。这些应用程序的交互性要求模型推理的作业完成时间 (JCT) 较低。现有的 LLM 服务系统为推理作业进行“运行-到-完成”处理，这会受到队头阻塞和长 JCT 的影响。

FastServe是一种用于 LLM 的分布式推理服务系统。FastServe 利用 LLM 推理的自回归模式来实现在每个输出token粒度的优先。FastServe 使用抢占调度，通过跳-加入（skip-join）多级反馈队列（Multi-Level Feedback Queue，MLFQ）调度器来最小化 JCT。基于 LLM 推理的半信息不可知设置，调度器利用输入长度信息，为每个要加入的到达作业分配适当的初始队列。比加入队列优先级更高的队列被跳过，减少降级。设计一种高效的 GPU 内存管理机制，可以主动卸载和上传 GPU 内存和主机内存之间的中间状态，进行 LLM 推理。基于 NVIDIA FasterTransformer 构建 FastServe 的系统原型。实验结果表明，与 Orca 相比，FastServe 将平均 JCT 和尾部 JCT 分别提高 5.1 倍和 6.4 倍。

现有的推理服务解决方案（例如 Clockwork [29] 和 Shepherd [59]）主要针对确定性模型推理作业（例如 ResNet [31]）。它们依靠准确的执行时间分析来做出调度决策，而对于执行时间可变的 LLM 推理，这不起作用。Orca [58] 是最先进的 LLM 推理解决方案。它提出了迭代级调度，在每次迭代结束时，它可以向当前处理批次中添加新作业或从中移除已完成的作业。但是，它使用先到先服务 (FCFS) 来处理推理作业。作业一旦被调度，就会一直运行直到完成。由于 GPU 内存容量有限且推理作业需要的 JCT 较低，因此当前处理批次无法在任意数量的传入作业中进行扩展。众所周知，“运行-到-完成”处理具有队头阻塞 [35]。对于 LLM 推理作业来说，这个问题尤其严重，因为 LLM 的大小会导致绝对执行时间过长。大型 LLM 推理作业（即输出长度较长）将运行很长时间，从而阻塞后续的短作业。

大多数现有的推理服务系统，例如 Tensorflow Serving [43] 和 Triton Inference Server [19]，与 DNN 模型无关。它们作为底层执行引擎之上的抽象，对到达的作业进行排队，将作业分派到可用的计算资源，并将结果返回给客户端。由于 GPU 等加速器具有大量并行计算单元，因此它们通常会批处理作业以提高硬件利用率和系统吞吐量。启用批处理后，来自多个作业的输入张量将连接在一起并作为一个整体输入到模型中。与单个作业执行相比，批处理的缺点是内存开销更高。由于激活内存与模型大小成比例增长，LLM 的大尺寸限制 LLM 推理的最大批次大小。

随着 GPT 模型的普及，推理服务系统已经发展到包括针对 GPT 独特架构和迭代生成模式的特定优化。 GPT 架构的主要部分是 Transformer 层的堆叠，如图所示。在 Transformer 层中，Masked Self-Attention 模块是将其与 CNN 等其他架构区分开来的核心组件。对于输入中的每个 token，它都会派生出三个值，即Q、K和V。它将Q与之前 token 的所有K点积相乘，从当前 token 的角度衡量之前 token 的兴趣。由于 GPT 是一种经过训练以预测下一个 token 的语言模型，因此每个 token 不应该看到其位置之后的信息。这是通过 Transformer 中的因果掩码实现的。然后，它将 Softmax 应用于点积获得权重，并根据权重生成V的加权和，作为输出。在高层次上，注意算子使输入中的每个 token 都知道其他的 token，而不管位置距离如何。

在 GPT 推理的每次迭代中，对于每个 token，注意算子都需要其前一个 token 的K和V。一种简单的无状态实现总是在每次迭代中重新计算所有的K和V。为了避免这种重新计算的开销，fairseq [44] 建议将K和V保存在KV缓存中以供跨迭代重用。这样，推理过程可以分为两个阶段。如图说明了KV缓存在不同阶段的使用情况。在初始化阶段，即第一次迭代，处理提示以生成 GPT 每个Transformer层的KV缓存。在解码阶段，GPT 只需要计算新生成的 token 的Q、K和V。利用和更新KV缓存来逐步生成 token。因此，解码器阶段的迭代执行时间通常小于第一次迭代。其他针对 Transformers 优化的系统库，如 HuggingFace [56] 和 FasterTransformer [18] 也执行了同样的优化。

另一个优化是 Orca [58] 提出的迭代级调度。简单的作业级调度会执行一批作业，直到所有作业完成。提前完成的作业无法返回客户端，而新到达的作业必须等到当前批次完成。相反，迭代级调度会调用执行引擎每次只对批次运行一次迭代，即为每个作业生成一个输出 token。每次迭代之后，完成的作业可以离开这个批次，到达的作业可以加入批次。然而，最大批次大小受 GPU 内存容量限制，交互式应用程序的低延迟要求也会影响批次大小的选择。

如图说明 FastServe 的架构。用户将作业提交到作业池。跳跃-加入（skip-join） MLFQ 调度器利用一个分析器（profiler）根据新到达作业的启动阶段执行时间确定其初始优先级。它采用迭代级优先，并优先执行最少获得的作业，以解决队头阻塞问题。一旦选择执行某个作业，调度器就会将其发送到分布式执行引擎，该引擎为 GPU 集群中的 LLM 提供服务，并与分布式KV缓存交互，以在运行时检索和更新相应作业的KV张量。为了解决 GPU 内存容量有限的问题，KV缓存管理器主动将低优先级作业的KV张量卸载到主机内存，并根据工作负载的突发性动态调整其卸载策略。为了扩展系统以服务于 GPT-3 175B 等大型模型，FastServe 将模型推理分布在多个 GPU 上。调度器和KV缓存中添加了扩展以支持分布式执行。

如图突出显示核心调度操作，而算法显示了伪代码。调度器使用基本 MLFQ 框架和 skip-join 功能来处理新作业。𝑄1 的量子设置为最小迭代时间，𝑄𝑖 和 𝑄𝑖−1 之间的比率由参数量子比控制。默认设置为 2，实验表明 FastServe 的性能对此量子设置不敏感。在完成当前处理批次中的作业迭代后，调度器会优先这些作业 𝐽𝑝𝑟𝑒 并调用过程 𝑆𝑘𝑖𝑝𝐽𝑜𝑖𝑛𝑀𝐿𝐹𝑄𝑆𝑐h𝑒𝑑𝑢𝑙𝑒𝑟。此过程处理新到达的作业 𝐽in 并构建一批新的作业 𝐽out 以执行下一次迭代。

上图算法伪代码中，调度器根据新到达作业的首次迭代时间（由输入序列长度决定）准确地为其分配优先级。具体而言，当作业到达时，使用 𝑔𝑒𝑡𝐻𝑖𝑔h𝑒𝑠𝑡𝑃𝑟𝑖𝑜𝑟𝑖𝑡𝑦 方法将其优先级设置为时间量大于作业首次迭代时间的最高优先级（第 7-8 行）。然后，（1）调度器将作业跳跃-参与（skip-join）加入到其相应的队列中，而不是简单 MLFQ 中的最高优先级队列（第 9 行）。对于被抢占的作业，调度器立即将新生成的tokens返回给客户端，而不是在作业完成之前返回整个响应，这优化了用户体验（第 12 行）。如果作业未完成且在当前队列中用尽了其时间片，则调度器将根据作业的当前优先级和下一次迭代时间，使用 𝑔𝑒𝑡𝐷𝑒𝑚𝑜𝑡𝑖𝑜𝑛𝑃𝑟𝑖𝑜𝑟𝑖𝑡𝑦 决定作业的降级优先级，然后（2）将其降级到相应的队列（第 17-20 行）。 跳跃-参与（skip-join） 和降级操作可能会导致输入长度或输出长度较长的作业遭受饥饿。为了避免这种情况，（3）调度器会定期重置作业的优先级，如果作业处于等待状态的时间超过一个提升阈值 𝑆𝑇𝐴𝑅𝑉𝐸_𝐿𝐼𝑀𝐼𝑇（第 22-26 行），则将其提升到最高优先级队列 𝑄1。如果提升作业的下一次迭代时间小于 𝑄1 的量子，则它将获得额外的量子，确保其下一次迭代不会被抢占。这会产生队头阻塞的可能性，因此 FastServe 的系统管理员可以调整 𝑆𝑇𝐴𝑅𝑉𝐸_𝐿𝐼𝑀𝐼𝑇 以在性能和饥饿之间进行权衡。最后，调度器选择一组优先级最高的作业，且不超过最大批处理大小，这是 GPU 内存容量的限制（第 28-31 行）。通过利用 LLM 推理的特点，跳跃-参与（skip-join）MLFQ 调度器可以更准确地调整作业优先级并减少降级。因此，它比朴素 MLFQ 更好地近似最短剩余处理时间（Shortest Remaining Processing Time，SRPT）方法，这是一种著名的最小化平均作业完成时间（JCT）抢占调度策略。

使用“运行-至-完成”策略的调度器可以容忍增加的GPU内存开销，因为正在进行的作业的最大数量不会超过当前处理批次的大小。如图展示了在一个合成工作负载下 GPT-3 2.7B 模型的 FCFS 和 跳跃-参与（skip-join）MLFQ 的KV缓存内存消耗。虽然选择一个相对较小的模型并将最大输出长度限制为 20，但 跳跃-参与（skip-join）MLFQ 的峰值 KV 缓存内存开销，可能是 FCFS 的 7 倍。在更现实的场景中，模型大小扩展到 175B 并且输出长度可能超过一千，跳跃-参与（skip-join）MLFQ 的内存开销很容易超过 NVIDIA 最新 Hopper 80 GB GPU 的内存容量。

FastServe 的主动KV缓存交换机制也是分布式的。由于 LLM 的不同阶段同时处理不同的作业，因此每个阶段可能会独立卸载或上传不同的KV张量。为了减少冗余控制，在处理上一个阶段发送的中间结果之前，当前阶段会执行与上一个阶段相同的卸载或上传动作。中间结果传输和KV缓存交换是并行发生的，因此KV缓存交换的开销进一步降低。如图所示，当中间结果被发送到下一个阶段时，下一个阶段会收到交换指令，并且可以根据需要同时交换KV缓存。KV交换机制只需决定第一个阶段的卸载或上传。当使用张量并行将第一个阶段拆分为多个块时，集中式KV缓存交换管理器会指示第一个阶段中的所有块卸载或上传同一个作业拥有的KV张量。

本文工作用 10,000 行 Python 和 C++ 代码实现 FastServe。基于 NVIDIA FasterTransformer[18]，一个具有自定义 CUDA 核实现的高性能Transformer库，构建分布式执行引擎。对其做一些修改，支持迭代级调度并与KV缓存管理器交互。还为其流水线并行添加扩展，因为其原始的作业级调度实现，不允许在当前运行的批次完成之前注入另一个批次。它只能将一批作业拆分为多个微批次 [33]，并把不同流水线阶段的执行跨微批次流水线化。这失去了在作业批次之间流水线执行的机会，而较小的微批次会降低设备利用率。在实现中，执行引擎可以在第一个流水线阶段完成执行后立即接收一批新作业，这意味着模型的每个分区都会处理其中一个批次而不会处于空闲状态。

以分布式方式使用 MPI [26] 实现KV缓存管理器，因为KV张量是在不同的 GPU 上生成和使用的。分布式设计使得可以在相应的 GPU 上保存和检索KV张量，从而最大限度地减少数据传输开销。还使用 MPI 传递消息以同步 GPU 之间的卸载过程，并利用多个 CUDA 流通过主动交换重叠计算。