vLLM源码之模型并行

原文：https://zhuanlan.zhihu.com/p/...

关于vLLM，之前介绍过vLLM框架（vLLM源码之框架执行）和PagedAttention的算子（vLLM源码之PagedAttention），本文主要结合代码，希望可以图文并茂的方式分析vLLM的模型并行（vLLM版本为v0.5.1）。

背景

笔者上周前学习了一下xDiT这个优秀的项目 https://github.com/xdit-proje... （https://github.com/xdit-proje...），并且对vLLM和xDiT做了一些个人的思考（手抓饼熊：从xDiT和vLLM引起对分布式系统的思考https://zhuanlan.zhihu.com/p/...）。回头看xDiT主页上的图笔者发现了一些神奇的事情。

xDiT整体架构

如上图所示，Tensor Parallel这个技术在xDiT里是灰色的，笔者觉得这个颜色，必有深意，于是便有了今天这篇文章，专门分析vLLM的模型并行实现，有了这个基础，Tensor Parallel这个技术在xDiT里是灰色的原因也就出来了。

本文首先分析一下模型并行的基础，在此基础上学习vLLM模型并行实现便更容易一下。

模型并行基础

Transformer架构

本节首先介绍Transformer模型结构以及推理的prefill和decode阶段。

Transformer结构

LLM基于自回归Transformer架构，如上图中简化展示。一个单独的Transformer层包括一个attention块，后面跟着一个两层的多层感知器（MLP），这些层被多次复制以形成完整的Transformer模型。Transformer模型是自回归的，意味着它们根据输入提示标记和先前生成的输出标记逐个生成新的标记。具体来说，Transformer模型的推理过程可以分为两个阶段。

prefill阶段

prefill阶段将整个用户提示信息计算成第一个输出令牌。在这个过程中，Transformer还为所有提示令牌生成键和值向量，这些向量将被存储以供将来decode使用。因此，对于这个阶段，输入的大小是整个提示的长度。该阶段与训练的模型forward阶段非常的类似。

decode阶段

自回归生成阶段逐个生成剩余的标记。具体来说，在每次迭代中，Transformer模型接收上一次生成的输出标记，并使用先前生成的所有键和值向量计算序列中的下一个标记。这个生成过程一直持续，直到序列达到最大长度（由用户或LLM指定）或返回一个结束序列标记。在此阶段，每次迭代输入的大小始终为一，因为只使用生成的最后一个输出标记。

这里的模型并行会同时作用于prefill阶段和decode阶段。

分布式矩阵乘法

Transformer的模型并行主要作用于attention、MLP、embedding和最后的softmax这些层，又以attention和MLP最为重要。无论是attention还是MLP的模型并行，都涉及到分布式矩阵乘法，故本节介绍一下分布式矩阵乘法。分布式矩阵乘法 Y = XA，分为行切和列切，这里的行切和列切指的是右侧A矩阵使用的是行切还是列切。

分布式矩阵乘法-行切示意图

如上图展示了行切的例子，如果A使用了行切，那么X矩阵必须是列切的矩阵，由上图可以看出，绿色的矩阵在卡0，蓝色的矩阵在卡1，执行矩阵乘法之后，2张卡的数据均是全量（最终输出矩阵shape对的）不完整的数据，需要进行AllReduce才能得到完整正确的数据。

分布式矩阵乘法-列切示意图

如上图展示了列切的例子，如果A使用了列切，那么X矩阵必须是全量未经过切分的矩阵，由上图可以看出，绿色的矩阵在卡0，蓝色的矩阵在卡1，执行矩阵乘法之后，2张卡的数据均是部分（最终输出矩阵shape是按照列切的）完整的数据，需要进行AllGather才能得到完整正确的数据。

通过上述分布式矩阵乘法，对于Y=XW的矩阵乘法我们得出了2个结论：

如果左边X为按照列切分的话，右边W必须按照行切分，结果需要经过AllReduce；
如果左边X如果没有切分的话，右边的W可以按照列切分，结果需要经过AllGather；

Transformer模型并行

MLP的模型并行

我们知道MLP是由2个矩阵乘法组成的（矩阵乘法结束后还有一些激活函数如GeLU等，如下图所示，不影响整体分析），MLP的输入是一个完整的X，输出也是一个完整的Z，这样并很容易利用上述分布式矩阵乘的结论设计MLP的模型并行。

如上图所示，输入X为全量数据，此时第一个矩阵乘法的A矩阵按照列切，XA之后的结果数据是[XA1, XA2]是一个列切的（XA1在卡0、XA2在卡1），此时不必进行AllGather，因为还需要和B进行矩阵乘法。由于左边的XA是按照列切的，故右边的B应该是按照行切的，XA和B进行分布式矩阵乘法之后，最终的结果进行AllReduce即可使2张卡均得到完整的结果。

Attention的模型并行

attention逻辑如下几个步骤：

首先attention是多个head的，如上图左部分，即有2个head的attention，最终不同的head算出的结果如果需要进行合并，在最后一维列方向合并，
如果attention需要进行多卡拆分，那么天然的可以按照head切分，如上图所示，Q1、K1、V1、和最终结果Y1在卡0，Q2、K2、V2、和最终结果Y2在卡1，Y1和Y2是按照列切分的矩阵；
attention之后需要再乘矩阵o（图中是B），由于Y1和Y2是按照列切分的，B就是按照行切分的，YB之后的结果，进行AllReduce；

左图（transformer架构）右图（tensor并行transformer架构）

此时再来全局看一下引入模型并行之后transofmer的结构，如上图所示。

MLP首先会有一个列切的层、然后再有一个行切的层，之后会有一个Allreduce同步结果；
attention每张卡有不同的head，attention之后会有一个行切的层，之后会有一个Allreduce同步结果；
右图红色的部分是“驱动模型执行的调度器”，模型并行需要给不同的卡喂进去相同的数据，在训练的时候，其实就是trainner不断的读取相同的数据，将数据送到模型，执行模型的forward方法，这里模型并行维度为2，即有2个调度器驱动模型执行，调度器和worker在一个进程；

至此，模型并行相关的基础知识已经介绍完成，接下来介绍vLLM的模型并行实现。首先介绍一下vLLM整体组件的结构。

vLLM整体组件结构

如上图所示是vLLM早期的结构图，由LLMEngine驱动模型整体的实现，scheduler负责调度任务（选择哪些数据执行），Worker负责模型执行，这里如果模型并行度为2，那么就会启动2个Worker，每个Worker执行model的方法。

新版本的结构增加了Executor的抽象，我们看到注释，Executor是分布式runtime管理器，一个Executor里会有多个worker，虽然多了一层Executor，但是本质上和之前的版本类似。学习vLLM模型并行，核心需要搞清楚2件事情：

vLLM的模型并行模型层是怎么实现的（这个主要是model里的实现）；
vLLM的模型并行整体流程是什么，还记得Transformer模型并行最后一小节，文章中说训练的时候会有一个调度器，不断的驱动模型执行，推理阶段这部分内容和训练阶段差异较大，是由多个组件配合而成的；

下面2章即分别分析这两个技术点。

vLLM模型并行model实现

上图是关于模型的实现，我们可以看到和模型并行有关的是一个ParallelLinear的类（ParallelLinear不在赘述，直接看代码即可，torch层切weight），该类是行切和列切类的父类，attention和MLP中均使用了该类的子类。

vLLM并行MLP实现

前面关于MLP模型并行已经说的比较详细，这里直接看图和代码即可。

vLLM并行Attention实现

前面关于attention模型并行已经说的比较详细，这里直接看图和代码即可。

embedding层的模型并行使用的是按照行切的方式，通过mask input等实现，结果会经过一次Allreduce，和Megatron-LM基本类似，不详细分析。

这里和训练有一个区别是，训练的最后一层softmax也使用了高性能的分布式softmax，避免通信量较大，而这里没有使用，读者可以自行思考一下原因。

vLLM并行执行流程

并行计算任务的运行（这里的场景可以是模型并行、如2卡情况下，会有2个进程进行模型并行任务。）一般有2种方式，分别以tensorflow v1为代表的single-controller模型和PyTorch为代表的multi-controller。

如上图所示single-controller模式，最上面是一个调度器，下面2个则是具体的执行器，调度器驱动执行器执行，故调度器起一个协调作用，协调2个worker一起执行任务，早期的vLLM就是这种方式。

如上图所示multi-controller模式，我们可以看到没有调度器这个角色了，只有2个执行器，目前的Megatron-LM、Deepspeed等均是这种模式。没有调度器的话，也就没有角色协调这两个worker执行了，那么worker之间无法协调怎么办呢？唯一的办法就是这两个worker执行相同的流程（Megatron-LM的模型并行流程，不同的模型进程组执行相同流程）或者根据worker角色不同用if else判断（Megatron-LM的流水线并行流程根据流水线阶段判断）。

vLLM最新版本的模型看起来比较的奇怪，既不是上文说的single-controller、也不是像是multi-controller，整体的演进如下图所示。

左图是比较老的的vLLM版本，我们从图中可以看出，整体像是一个single-controller的模式，由Ray driver进行调度，调度完之后驱动2个进程进行同时进行模型执行。

右图是比较新的vLLM版本，从图中可以看出，驱动进程和Worker0进程是一个进程了，这样带来的好处是节省了一点的时间开销，但是看代码的时候，稍微会让人产生一些疑惑，因为需要通过一些 if else进行判断。

下面结合代码和进程栈对新版vLLM的执行模式再进行一些分析。

上述图描述了TP = 2的时候2个进程的状态，一眼看出2个进程的进程栈是有区别的，整体说明如下：

左图代表driver进程的状态，我们可以看出，左图线程栈的起始点是LLMEngine的step方法，该方法首先执行schedule、根据schedule的结果进行一次execute_model操作；
右图代表另一个worker1进程的状态，我们可以看出，右图的线程栈的起始点是start_worker_execution_loop这个函数，该函数直接执行execute_model方法，这个方法挺奇怪的，输入的execute_model_req = None，那worker1进程的数据是什么呢？

其实无论是driver还是worker1进程，最终模型都开始执行了execute_model方法，我们结合上面的图进行解释。

左图是driver的进程代码执行位置，我们知道driver的step方法在schedule结束之后，会进入这个方法，在driver进程中，程序执行路径是is_driver_worker为true的代码，也就是上面的逻辑，会经过一系列输入数据，最后经过broadcast_tensor_dict；
右图是worker1进程，start_worker_execution_loop方法是一个死循环，该循环一直执行execute_model方法，但是execute_model方法，会停在broadcast_tensor_dict这个逻辑（这是mpi程序的特点，broadcast_tensor_dict底层是一个broadcast集合通信，所有的进程都需要执行相同的方法），当driver执行该方法之后，此时worker1也在执行该方法，driver就把数据给个worker1，2者后续就可以愉快的一起进行模型并行的执行了；

从上述分析可以看出vLLM模型并行的执行模式还是有一些特点的。其实传统分布式系统研发对single-controller了解比较深，而大模型分布式系统尤其是训练方向的研发更熟悉multi-controller。新版vLLM的流程可以看成这两个的结合体，在模型层面大家执行的都是一样的代码（model forward），而具体控制着一块则是driver和worker协同的，这里通过判断各自角色来执行不同代码。

总结

关于vLLM的模型并行模型层实现和模型并行执行逻辑已经讲完，回到最初的问题，那为何 Tensor Parallel这个技术在xDiT里是灰色的呢？这里结论已经很显然了。

作者：手抓饼熊
来源：大模型新视界

推荐阅读

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

背景