图解大模型计算加速系列:分离式推理架构1,从DistServe谈起

大家好,最近Kimi开源了它的推理架构Mooncake的技术报告,让分离式推理架构的关注度一下升了起来。所以在这个系列中,我打算写一写关于分离式推理架构的一些有趣的优化知识。对于这个架构,我之前也只是了解一些,并没有深入探究过,所以在这个系列中我也和大家一起学习,一起发现好玩的东西。

本篇作为该系列的第一篇,选择DistServe这个分离式架构入手,选择它的原因是

  • 这篇论文中通过各种实验和数学建模,很好论述了“为什么要用分离式架构”这一点。很适合初次了解这个架构的朋友阅读。这也是本文强调的重点
  • 这篇论文代码是开源的,本文在写作时,也借鉴了开源代码的一些内容
  • 调度策略比较简单(FCFS),也没有做抢占之类的操作(所以本文的重点也不会放在这里)

【全文目录如下】
一、LLM推理的两阶段及评估指标
二、为什么需要prefill和decode分离
三、分离架构下prefill和decode的优化方向
四、实操:代码实践中的prefill和decode并行策略

(每个大标题下还有很多细节,大家可以看正文,这里偷个懒,就不把小标题打出来了)

【如有帮助,欢迎大家点赞收藏和在看~~~~~】

一、LLM推理的两阶段及评价指标

1.1 LLM推理的两阶段

image.png

  • prefill阶段:把整段prompt喂给模型做forward计算。prefill阶段结束后,模型产出第一个token(例如图中t3)
  • decode阶段一个token一个token地产出response。(例如图中逐一产出t4,t5,t6)

既然LLM的推理由这两个阶段组成,那么一个LLM系统的性能评估同样也和这两个阶段密切相关。解析来,我们来看两个分别用于评价这两阶段性能的指标。

1.2 prefill性能评估指标:TTFT

TTFT(Time To First Token),表示生成第1个token所用的时间,这是prefill阶段的重要评估指标之一。

在实践中,我们一般会人为设置一个TTFT SLO(Service Level Objective),这个指标是我们给系统提的要求,只有系统达到了这个指标,我们才认为它的性能达标了。举例来说,假设我们定P90 TTFT SLO = 0.4s,那么这意味着我们对该系统的要求是:“90%的request的TTFT值都必须<=0.4”

1.3 decode性能评估指标:TPOT

TPOT(Time Per Output Token),产出每一个response token所用的时间,这是decode阶段的重要评估指标之一。

同样,我们也有一个TPOT SLO指标,这是人为给系统定下的性能要求。

二、为什么需要做PD分离

在快速回顾LLM推理的两阶段及它们的评估指标后,现在我们来看一个重要的问题:为什么需要做PD分离?

2.1 分离架构速览

在一些推理框架中,prefill和decode是合并在一起的,你可以通俗理解成一块gpu上我既做prefill又做decode。我们以vllm框架为例:在一张卡上它的prefill和decode是交替进行的,即vllm的调度器会根据存储使用情况、请求池中的请求状态等要素,决定在模型推理的每1个阶段是做prefill还是做decode。

分离式的框架长什么样呢?我们直接来看DistServe提供的一张架构简图:

image.png

从这张架构简图中,我们不难发现,在分离式框架中prefill和decode不再共享一块/一群gpu,而是你做你的prefill,我做我的decode。prefill实例计算完毕后,把产出的KV Cache传送到decode实例,然后decode实例继续做推理。这里我们再明确一些术语:

  • Prefill instance:prefill实例。一个prefill实例中包含一个完整的模型,它可能占据1或多块gpu。
  • Decode instance:decode实例。定义同上,只是它与prefill实例不再共享gpu。

介绍到这里,你肯定开始疑惑了:从直觉上看,分离式架构相比于合并式架构,多加载了模型副本(耗显存),同时还涉及到gpu间的KV Cache传输(耗时间),怎么看都好像比不上合并式架构,那么我们做分离的原因何在呢?

接下来,我们就来回答这个问题。

2.2 TTFT和TPOT间的balance

假设你现在有一个requests pool,里面装了状态不同的requests,有正在等待做prefill的,也有刚做完prefill等待做decode的,而你面前只有1张GPU。由于prefill和decode在做fwd阶段的差异性,如果不考虑特殊的优化方法,在模型每1个推理阶段,你要么全做prefill,要么全做decode。现在你需要做一个决策:在下1个推理时刻,我该做prefill,还是做decode?

你观察了一下requests pool,发现里面塞满了等待做prefill的请求,你心想“要不还是赶紧多做些prefill,好歹给用户返回第一个token,防止他们生气了直接不用我的产品了”。所以你在接下来的时刻中一直在做prefill,由于gpu资源是有限的,你做了prefill就无法同时做decode了,你其实是在牺牲TPOT保全TTFT,反之亦然。

所以这时,你产生了一个疑问:当prefill和decode合并在一起的时候,似乎无论调度策略怎么调整,TTFT和TPOT这两个指标都存在强耦合关系。因此如果我把这两个阶段分开,让TTFT和TPOT各自独立发展优化,会不会得到一个更好的结果?

DistServe做了一个实验来验证你的猜想,接下来我们就来看这个实验结果。

2.3 分离VS合并:一个实验结果

image.png

实验配置:

  • 1张A100 80G的gpu
  • 13B LLM,input_length = 512, output_length = 64

SLO(人为定义的系统性能达标要求):

  • P90 TTFT = 0.4s,人为定义在prefill阶段90%的请求的TTFT必须达到0.4s(上面那张图黑色虚线部分)
  • P90 TPOT = 0.04s,人为定义在decode阶段90%的请求的TPOT必须达到0.04s(下面那张图的黑色需要部分)

我们先只看图中的蓝线部分:

  • 在蓝线部分对应的实验中,我们采用PD合并的架构
  • 横轴表示每秒到达这块gpu上的请求数量,记为rps(requests per second)
  • 随着rate增加,我们统计每个请求的TTFT和TPOT指标,进而绘制出蓝线。
  • 将蓝线和我们设定的SLO(黑虚线)对比发现,如果我们同时达到设定好的TTFT SLO和TPOT SLO,我们的最大rps = 1.6,我们也记这个指标为goodput。

再来看图中的黄线和绿线部分:

  • 还是这块gpu,现在我们做两次实验,一次实验只让它处理需要做prefill的请求(黄线),一次实验只让它处理需要做decode的请求(绿线)
  • 按照前文介绍的流程,我们得到在这块gpu上只做prefill时,它的goodput=5.6;只做decode时,它的goodput=10。

从这个实验中,我们可以发现,一张卡只做prefill或者只做decode的goodput,都要比它既做prefill又做decode的goodput要高,这也验证了我们前文的猜想:合并式架构可能会同时损害TTFT和TPOT。

假设现在我们不考虑KV cache传输的网络成本,观察上述实验结果可发现:纯decode的goodput是纯prefill的goodput的2倍,所以我们可以估算:如果我拿2张卡纯做prefill,1张卡纯做decode,那么我这三卡构成的架构的goodput就为10,也即单卡goodput = 3.3。这是合并式goodput(值为1.6)的2.1倍!

再换一个角度粗糙地看:分离式架构中,我们三张卡承受起了10 reqs/s的流量;而合并式架构中由于单卡流量承受为1.6 reqs/s,所以我们需要6张卡才能承受起10 reqs/s的流量。因此回到2.1结尾中我们的疑问,乍一看分离式架构似乎更耗卡了,但实际中如果优化得当,它其实是比合并式架构更加省钱的。

三、分离架构下prefill/decode的优化方向

通过上面的分析,我们具象化体会到prefill和decode分离的好处。在本节中,我们再进一步,详细来看当我们确定可以使用分离架构后,prefill和decode阶段都有哪些优化方向。

3.1 算力与存储的独立优化

  • prefill阶段:拥有计算受限的性质(compute-bound),特别是在请求流量较大,用户的prompt也比较长的情况下。prefill阶段算完KV cache并发给deocde阶段后,理论上prefill就不再需要这个KV cache了(当然你也可以采用LRU等策略对KV cache的保存做管理,而不是一股脑地清除)。
  • decode阶段:拥有存储受限的性质(memory-bound),因为token by token的生成方式,decode阶段要频繁从存储中读取KV Cache,同时也意味着它需要尽可能保存KV cache。

对comput-bound和memory-bound不理解的朋友,可以参见我之前的这篇文章

正是由于这两个阶段这两种特性,我们在为其做硬件分配时,也有了更加灵活的空间,而不是只使用一套型号完全一致的硬件。也就是在分离式框架下,计算和存储可以朝着两个独立的方向做优化。

3.2 batching策略的独立优化

我们依然先展示DistServe所作的一个实验结果。

image.png

这个实验比较了在不同batch size下,prefill和decode阶段的吞吐量情况(tokens/s,即每秒能处理的tokens数量)不同颜色的线条代表不同大小的input_length。

  • prefill阶段:随着batch size的增加,吞吐量的增长趋势却趋于平缓。这正是因为prefill阶段是compute-bound的,当batch中的总tokens数大于一个阈值的时候,prefill阶段产生compute-bound,使得吞吐量的增长趋势变缓(这里我们不考虑显存问题,因为在该实验中硬件显存足够这些不同batchsize和input_length的数据使用)。
  • decode阶段:随着batch size的增加,吞吐量的增长趋势越来越显著。这是因为decode阶段是memory-bound的,即相比于计算,读写数据的时间要更多。所以在decode阶段中,如果我们能提升batch size,就能把计算强度提起来,吞吐量就上去了。

通过这个实验,我们不难得出结论:prefill和decode不适合用一样的batching策略,所以当两者分离后,我们可以在batching上继续做优化。

3.3 并行策略的优化

在合并式框架中,prefill和decode共享一个模型,这意味着它们共享一样的dp/tp/pp/....等并行策略。考虑2.3中的实验结果,我们根据实验指标,认为可以设计一个2个prefill实例(dp=2),1个decode实例(dp=1)的分离系统。所以我们自然而然想到,如果prefill和decode分离了,那是不是意味着它们的并行策略可以各不相同呢?

由于在前面的定义中,我们认为一个prefill/decode实例包含一个完整的模型副本。所以在本节的并行策略中,我们只考虑tp + pp(即1个实例只通过tp +pp划分)。这样一来,只要对这两个阶段,分别确认好最佳的tp + pp策略,我们就可以用这个策略复制多个prefill/decode实例了(dp维度)

明确了这一点,接下来,我们通过一个有趣的实验 + 数学建模的方式,来看一下prefill和decode在并行策略(tp, pp)上的倾向性。

(1)prefill阶段的并行策略

怎么样可以知道不同并行策略对prefill阶段的影响呢?最简单的方法当然是通过做实验。实验我们肯定是会做的。但在展示实验结果之前,我们尝试一个有趣的思路:能否拟合一个数学模型,来评估不同并行策略对prefill阶段的TTFT指标的影响?

如果我们能做这样的数学建模,同时用实际实验结果去论证数学模型的有效性,那么之后我们的系统在选择最佳并行策略时,就不需要通过做实验(比如模拟假数据,送入模型实际做fwd)的方式去搜索最佳的并行策略了。我们完全可以用自己的数学模型去估算一个结果,这样岂不是能降低系统在初始化时的耗时?

在DistServe中,在研究不同并行策略对prefill阶段TTFT的影响时,就使用了“排队论”建立了这个数学模型。不过由于这个模型中一些强假设性(我们马上会讲这点),它并没有用在最终的代码实践中(代码实践换了一种建模方式,我们马上也会谈)。不过因为我觉得这个基于排队论的建模挺有意思的,所以我在这节中也会详细来讲。

在下面的讲解中,对于prefill阶段,我们分别对于单卡,2卡pp,2卡tp的情况做基于排队论的建模(没错,这里我们先不谈tp + pp 的混合情况,这里我们主要想看对于prefill阶段,它对纯tp和纯pp的倾向性)。

单卡下prefill TTFT

我们假设队列中的请求服从排队论中的M/D/1模式,这个模式的含义如下:

  • M:请求的到达过程遵从柏松分布。即请求的到达系统是相互独立的(不存在互相干扰)且单位时间内各请求到达系统的概率是等可能的。
  • D:假设所有请求做prefill的处理时间相同。这个处理时间是指请求被喂给模型一直到模型吐出结果的时间。这就是我们前文说的“强假设”情况。
  • 1:可以理解成只有1张gpu

在这个模式下,单卡下prefill阶段的avg_TTFT可以被建模成:

image.png

  • D:即前文所说的单个请求处理时间
  • R:即遵从柏松分布的请求到达率(每s有多少个请求到达gpu)上。

我知道你现在一定对公式(1)一头雾水,不要急,我们这就来详细解释它。想象一下,对于队列中的某个请求,它的TTFT要怎么计算呢?稍一动脑,我们不难从直觉上给出:

一个请求的TTFT = 这个请求在队列中等待的时间 + 单个请求被处理的时间 = 排在这个请求之前的请求数量 * 单个请求被处理的时间 + 单个请求被处理的时间

image.png

最后注意,在M/D/1理论的假设下,RD一定是小于1的,如果RD大于1,则意味着这个系统的队列长度会趋向无限(即处理速度比进入速度慢,且进入是不停歇的),那么这时候就不再适合用这个排队理论了。

2-way PP下prefill TTFT

现在,我们假设有两张卡,我们对模型做2-way pp,则此时prefill阶段的TTFT建模如下:

image.png

2-way tp下prefill TTFT

同样假设有2张卡,做2-way tp,则我们有:

image.png

这里的K可理解成是一个折损因子。理论上,如果不考虑因为tp产生的通讯消耗,那么此时在多个gpu上做不同head的并行计算,可能是比单个gpu同时做这些head的并行计算要快的。但是如果考虑了通讯消耗,那么这个加速时间就有所折损,我们设1<K<2。不同的设备拥有不同的折损因子,它对TTFT也是有影响的。

建模 VS 实验结果

建模建完了,这下我们可以用真实的实验结果看看建模拟合效果怎么样了。在DistServe中,确实做了这样的一个实验,同时也提到了建模效果基本拟合的试验曲线的趋势(但是没有单独给出建模效果的趋势图)。我们来看下实验效果:

image.png

如图,左侧为真实实验的结果。关于右侧的图,我这里有些迷惑,它似乎是在论证tp模型中取不同折损因子K值时对TTFT的影响(黄线),但是如果是这样,这个曲线应该是用建模结果表示的,但它现在却长得和真实实验结果一致。

但是anyway,回到左侧的真实实验结果上,我们不难发现,对于prefill阶段,在rate较小时,更适合用tp;在rate较大时,更适合用pp。也即prefill阶段在不同条件下,对并行方式有倾向性。

(2) decode阶段的并行策略

这里DistServe不再额外建模了,但是它从另外一个角度做了模拟实验的结果:

image.png

  • latency:模型处理完一个请求的耗时
  • throughput:单位时间内模型能处理的token数量

可以发现,对于decode阶段,随着gpu数量的增加,如果采用pp的方式,能产生更高的吞吐量(理由上面讲过,PP对batch的处理时流水线式的);但是如果用tp,则能产生更低的延迟(即虽然单位时间内处理的tokens数量有限,但是处理单个请求的时间变快了)。

四、实操:代码实践中的prefill/decode并行策略

在上述的讲解中,我们通过实验的方式,比较直观感受到prefill和decode阶段对于不同并行策略的倾向性。我们还首次尝试通过数学建模的方式,来拟合在不同并行策略下系统的性能曲线(不过正如前文所说,我们初次尝试的这个数学模型具有一些强假设,使得它不适合用在实践中。因此在实践中我们可能要考虑别的建模方式,我们马上来谈着点)。在本节中,我们来具体看DistServe在代码实践中是如何为prefill和decode阶段分配最佳并行策略的。

在开始讲解之前,我们先明确实操中DistServe的整体优化思路:

  • (1)对prefill和decode分别做goodput最大化,也就是对他们分别寻找最佳的并行策略(确定一个实例/一个模型副本要怎么分割)
  • (2)根据系统中的流量大小确定prefill/decode下一共要部署多少个实例,然后按照(1)中确定的并行方式做实例部署
  • 在DistServe中,管这样的策略叫placement。

接下来,我们再来看两种不同的硬件场景:

  • 场景1:node间带宽十分高,这样我们可以忽略KV cache的传输时间。我们使用algorithm1来寻找在这样的场景下对单个prefill/decode实例的最佳并行策略。值得说明的是,单node推理也属于该场景
  • 场景2:node间的带宽并不高,不能忽略KV cache的传输时间。这种情况下,我们尽量把prefill和decode实例对应的模型层分在同一个node中(也就是prefill和decode的pp维度必须相同)。这种是实践中更普遍的场景。

【本节根据论文伪代码及DistServe开源代码库撰写而成。https://github.com/LLMServe/DistServe/tree/main/simdistserve

4.1 伪代码:寻找最佳并行策略

image.png

左图对应场景1,右图对应场景2。

(1)场景1

  • N:表示一个prefill/decode实例最多能占用多少台node
  • M:表示一个node最多有多少张卡
  • C:集群中GPU的存储总量
  • W:系统的负载情况。这个指标和模型类型、input_length等指标密切相关。在3.3节做排队论建模中,我们曾假设模型处理单个请求的时间是一致的(这也意味着我们假设所有请求的input_length长度一样),这在现实中是不可能的。在现实生活中,你的这个系统可能部署了不同的模型,而用户对每个模型的使用有倾向性,同时在每个模型中,用户传入的input_length也是不同的,所以我们可能很难对系统的负载情况建模。但是我们可以通过模拟实验的方式,先收集一段时间的用户真实数据,然后用一个模型计算出用户真实请求服从的分布。接下来我们只需要从这个分布中做采样,就可以采用出一批数据帮助我们做并行策略的profiling了。我们可以把这批采样出来的数据理解成是系统的负载W。真实代码中的load_workload就在做这件事。
  • R: traffic rate,即单位时间内到达这个系统的请求数量。
  • best_plm: 最佳策略,等于(n, config_p, m, config_d)
  • config_p:  prefill阶段最佳的并行方式(每个config_p装一套完整的模型,我们称为一个prefill实例)
  • config_d: decode阶段最佳的并行方式(每个config_d装一套完整的模型,我们称为一个decode实例)
  • n: 理解为prefill阶段的dp值,即要部署多少个prefill实例
  • m:理解为decode阶段的dp值,即要部署多少个decode实例

了解了这些指标,看懂算法1就很容易了。整体采用网格式搜索的方法,分别对prefill和decode找到满足最大goodput(这个指标表示满足认为设定的SLO值的最大rate值,rate值是单位时间内到达这个实例上的请求数量)。

这里我们只额外介绍一个有趣的函数simu_prefill

我们以prefill阶段为例,这个函数的输入是当前我们搜索到的一个并行策略,以及模拟出来的负载W。在DistServe的代码实践中,它是用一个“二分搜索”的方法来找到这个goodput值的,我相信很多朋友一定有所疑惑:欸,为什么找goodput值要用二分搜索的办法?

我们先来想一下,当我们决定采用simu_prefill时,有哪些数据是明确的:

  • TTFT SLO:这是你认为设定的指标,即你认为系统要达到这个指标才算合格,因此它是确定的,它和算法无关。
  • 并行策略G
  • 负载W:我们前面说过,这是我们根据历史数据拟合出来的输入数据的特征分布(input_length之类),然后我们用这个分布进行采样,得到了一批数据。接下来我们就要考虑把这批数据按照什么样的节奏发送到你的prefill实例上。

那这个“节奏”是指什么呢?我们再回想一下2.3实验结果中的那张图(还记得吗,我们就是通过这张图定义了goodput的含义的),不难理解,这个节奏其实就是指发送rate,也就是goodput的备选值。

image.png

所以,goodput不是我们通过simu_prefill函数直接算出来的:

  • 它首先是一串备选值(比如goodput = 5就意味着rate在0~5间都可以)。
  • 然后我们以不同的goodput(节奏)发送模拟的数据(负载)到prefill实例上,按照当前设定的并行策略实际运行一次系统做推理,得到P90 TTFT指标。
  • 将P90 TTFT指标和我们的P90 TTFT SLO做对比,如果这个小于SLO,则意味着可以加大rate,反之同理。
  • 那既然现在goodput是一串备选值,那么我们就可以使用二分法,根据P90 TTFT和SLO的比较结果,改变我们搜索的上下界,找到最终的goodput(leetcode medium之魂燃起)

好,现在我们已经理清了simu_prefill的运作流程,到这一步为止,我们是通过模拟实验的方式找到goodput,进而确定最佳并行策略的。而实验的目的则是获取prefill阶段的TTFT指标。但是,正如前文所说,如果我们能找到一种更确切的建模方式,对TTFT进行建模,那么我们就可以拥有函数式的simulator,这样在后续就可以避免每次都跑模拟实验了。再详细一点,我们可以先跑几次实验结果,然后用这个实验结果作为真值,对我们的数学模型中的参数进行拟合。关于这一点,我们放在4.2节中详细介绍。

(2)场景2

理解场景2伪代码的关键就是明白,在场景2下node间的带宽较低,因此KV cache的传输时间不能被忽略。所以我们需要让decode和prefill实例对应的模型层在同一个node里,这也意味着它们的pp维度必须一致。在这个基础上理解2的伪代码就不难啦(脊椎病犯了,实在是敲不动了)。

(3)并行策略结果

最后,让我们来看下DistServe实操中对于prefill和decode找到的最佳并行策略吧(注意,在175B的模型中,一个实例肯定是要跨机的)

image.png

4.2 实操中对TTFT和TPOT指标的建模

image.png

现在我们来回答4.1中遗留的问题:如果排队论的假设太过理想,那么在实操中,要怎么对TTFT和TPOT指标建模呢?

我们直接放出DistServe论文的appendix部分,这一部分写得非常非常好,短小精炼。但是对于对FLOPs,IO复杂度计算和flash attention不太熟悉的朋友,可能要花一些时间里了解。所以我把这一页都摘出来做说明。

首先我们明确一点,在attention score的计算环节,我们使用了flashattention做特殊优化,因此它和别的环节的TTFT估计还是有一定差别的,所以我们需要把他们分开来说,我们以prefill阶段为例,具体而言:

  • 观察A.2中的table,它是去除attention score计算后剩余的GEMM矩阵。shape of M/N分别表示这个GEMM计算的两个矩阵的维度。例如对于QKV linear,它指代的就是我们的输入数据通过Wq/Wk/Wv矩阵产出QKV值的过程。只是这里(t, h)表示我们把所有batch的数据拼接成一条,(h, 3h)表示我们把Wq/Wk/Wv矩阵拼成一个。这个table的其余列同理可推
  • 根据这个table列举的GEMM计算,我们可得总FLOPs(对这个计算方式不明确的朋友,可以参看3.1节中给的链接文章),这样就得到T1计算公式中括号里的那一坨
  • 由于table中的GEMM计算都是compute-bound的,因此在考虑TTFT时我们可以忽略读取数据的时间。这样,我们设因子C1表示计算单个FLOPs的耗时,那么table中的TTFT消耗就是T1所示。C1则是我们要拟合的模型参数。
  • 观察T2,由于attention score的计算是memory-bound的(尽管我们已经用flashattention做过了优化),所以我们不关心它的FLOPs,而关心它的IO复杂度(对这个计算有困扰的朋友依然可以参考3.1链接文章)。然后在这个基础上我们再设因子C2,表示单次读取的耗时。
  • 综合T1和T2,建模出最终的prefill TTFT模型。对decode阶段也是同理,不再赘述。
作者:猛猿
来源:GiantPandaCV

推荐阅读

欢迎大家点赞留言,更多Arm技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。

推荐阅读
关注数
18701
内容数
1280
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息