【推理】PD分离的绝佳入门教程 —— DistServe

good

作为推理服务中PD分离的技术的经典方法之一，DistServe（OSDI 2024）无疑为广大学者提供了一份绝佳的入门教程。该工作详细分析了PD分离的技术动机、提出了通俗易懂的PD分离算法思想、并充分开源了代码，让每一个想了解PD分离技术的人都能够“亲自把玩”。
        本文为在精度论文并部署代码进行多组实验后整理的“玩后感”，会从PD分离技术动机、DistServe项目算法思想、code实现以及个人理解等方面展开介绍。本人初学，望看到本文的大佬们多多指点交流～
1. PD分离技术动机

        我们知道一条请求在推理的过程中需要经过两个步骤：第一步，根据输入序列的内容生成第一个token，该过程称为Prefill；第二步，根据输入序列以及之前已经生成的token生成下一个token，该过程称为Decoding。随着“以存代算”的KV-Cache技术的引入，大多推理过程避免了在Attention结构的重复计算，这导致Prefill与Decoding在执行时存在着较大的差异。
        为什么Prefill和Decoding存在较大的差异？
        首先我们从计算过程上去直观理解两个步骤的区别。在Prefill阶段，模型输入prompt的所有token，经过一步计算后得到第一个输出，与此同时每个token在计算时产生的KV-Cache被储存下来，而在Decoding阶段，每个新token的计算都会涉及到KV-Cache的读取与存储，而计算量相比Prefill而言被大大降低了，取而代之的是高频的存取操作带来的压力，这便是KV-Cache技术带来的Prefill和Decoding的差异化。
        这种差异是否会导致放在一个batch中处理会相互拖后腿？
        论文通过实验给出了肯定的答案。下图表示在全部是decoding状态的batch中加入一个做prefill的请求导致的延迟变化。蓝色线条的起点表示只处理一条prefill的延迟，蓝色后面的点与起点的差值都可以看做添加decoding状态的请求对该prefill状态的请求带来的延迟危害，而蓝色与橙色对应节点的差值可以看作添加的一个prefill状态的请求对decoding状态的请求batch带来的延迟危害。



截取自论文

          综合分析，两阶段混合推理带来的主要问题如下：

• 调度低效：在相同GPU上进行不同状态的调度会造成decoding等待prefill的现象，并且GPU在进行decoding会导致资源低利用率；
  
• 资源和并行策略的耦合：两阶段的差异对于资源和并行策略的需求不同，相同GPU上部署不可避免的会导致两阶段资源和并行策略的共享。
  

        至此，PD分离就成为了一种最为直接的解决思想。   
2. P/D阶段特性分析

        在正式提出算法前，论文对于P和D阶段的batch策略和并行方案进行了分析。对于batch策略，论文通过以下实验对比了两阶段的不同特点：



截取自论文

Prefill阶段并不是batch size越大吞吐越大，吞吐随batch size的变化取决于输入文本的长度，因此在设置batch size时要考虑输入长度。而增加batch size对于Decoding阶段是一个避免低吞吐的好方法。
         对于并行方式，为了分析Prefill阶段对于并行策略的偏好，论文首先利用M/D/1排队模型对请求队列进行了建模，进而TTFT就可以用请求在系统内的平均等待时间表示，加入Pipeline并行和Tensor并行的影响后，TTFT建模结果分别如下：
{\rm PP=2}:  \ \ \ \ \        D+\frac{RD^2}{4(2-RD)} {\rm TP=2}:  \ \ \ \ \ \           \frac{D}{K}+\frac{RD^2}{2K(K-RD)}
其中D为请求的执行时间，R为到达率，K为TP带来的影响系数，由于通信代价，通常 1<K<2 。当到达率R很小时，第一项执行时间为主要作用，这时TP的效果会更好，当R增加时第二项其主要作用，PP的效果则更好，结合实验结果看更清楚：



截取自论文

这一实验结论对根据不同场景下设置并行模式起到了参考作用。 而Decoding阶段的并行模式分析同理，TP可以降低延迟，但是过大的TP会引入过量通信，延迟降低的优势也就不再明显，PP虽不影响单个请求的执行时间，但可以加大吞吐。



截取自论文

        不过对于这一部分的分析，本人也做了粗略的实验比较了一下两个阶段TP与PP的影响，发现对于给定P/D两阶段的卡数的情况下似乎并行策略的影响不大～但是论文这部分的分析确实具有一定的启发性。
3. DistServe项目算法思想

        接下来正式介绍DistServe项目的算法思想。整个项目分成两大部分：设计并行模式、执行分离式推理。
3.1 并行模式设计

        论文的附录给出了对Prefill和Decode阶段的耗时建模方法，而在实验中，Prefill延迟模型为：
A+B*bs*l\_in +C*\sum_{i=1}^{bs} l\_in_{i}^2
Decoding延迟模型为：
A+B*bs*l\_in+C*bs
其中 bs 为batch size， l\_in 为输入的平均长度， l\_in_{i}^2 表示第 i 个batch的请求平均输入长度的平方，通过采样实际执行时间，则可作为训练数据来分别拟合两个延迟模型的参数 A、B、C 。
        总体寻优的思想也是比较清晰的。首先对不同并行策略下的耗时建立线性模型，再采集多组样本，对不同并行策略的线性模型进行拟合。值得一题的是，论文对Prefill和Decoding的耗时分别建立不同的模型，特别对于Decoding阶段又进一步区分了大Batch size和小Batch size两种场景下的建模。这或许可归因于第二节中对于Decoding阶段的分析，由于Decoding阶段的Batch size影响较大，区分建模能保证更高的精度。
        模型建成后的寻优准则则是该论文与其他论文方法不同的地方。DistServe全篇更加强调对于TTFT和TPOT分别达到SLO的目标，因此寻优算法也遵从这一准则。算法如下：



截取自论文

前6行在做的就是遍历所有可能的并行模式，第7-13行开始利用拟合好的模型预测效果，筛选最佳方案，筛选标准是：（1）. 在Prefill和Decoding均满足SLO，（2）. 并行方案对应的最大支撑到达率大于输入到达率，对于每个并行方案的支撑到达率则采用二分查找法确定。最后在所有满足条件的方案中，找到需要GPU数量最少的方案输出。
        论文同样给出了跨机的方法。跨机面临着缺少NVLink的高速连接，这时并行策略的性能则受制于机间通信带宽。不同于其他PD分离的方法，为了避免KV-Cache传输以及TP走机间通信，该论文在设置并行模式时，需要限制相同模型层的P和D需要在同一节点。示意图如下：


3.2 分离式推理

        确定好并行模式后，就可以开始部署推理引擎了。



截取自论文

        DistServe的整体逻辑遵从“生产者-消费者”模式，总体Engine负责管理Prefill、Decoding引擎，两个引擎分别实现了各自的调度器，Prefill引擎为生产者，Decoding引擎为消费者，到达的请求会首先分配到Prefill调度器的等待队列，一旦当前Prefill实例的资源满足该请求所需，就会被调度到Prefill引擎执行该阶段的计算，执行完Prefill的请求会被存储到一个中转队列中，该队列中的请求会等待Decoding调度器的调度，一旦Decoding调度器检测到中转队列请求所需资源能够满足，就会被调度到Decoding引擎，进而执行Decoding阶段的计算。
4. Code结构

        整个DistServe项目的代码可以分成两个部分。一个部分是项目基座，由C++和CUDA语言编写，用于定义推理服务中涉及到的模型计算、通信、资源开辟释放等，该项目名为SwiftTransformer，主要包含了对Transformer架构模型的优化算子等。第二部分则是论文重点内容所涉及的调度层，主要由Python编写。而在真正执行服务时，还有一层与用户交互的api层，同样由Python编写。



代码结构说明

5. 总结与讨论

        该项目作为PD分离的经典工作，非常适合刚接触该技术的初学者进行了解和学习。DistServe的实践简单、没有涉及过多的其他优化技巧，适合作为验证PD分离这一单一技术性能的工作，代码可读性比较高。然而也是因为实践方式较为简单，特别是对于跨机的方案限制比较多。
        另外，PD分离技术的使用场景也有相应的限制，例如至少需要两张GPU，因此考虑PD分离技术的场景通常需有相对充足的GPU资源，才能让该技术发挥出更大的空间。

原文地址：https://zhuanlan.zhihu.com/p/19666783423