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。