大模型推理：从模型分析到计算优化(三)

首页

大模型推理：从模型分析到计算优化(三)

发布时间：2023-09-25 22:00

在上一篇文章中，我们看到了大语言模型在单卡的较小规模实现中各个推理关键指标如延迟(latency)、吞吐量(throughput)和推理成本与用户指定的模型设定和硬件参数之间的关联。而大规模并行化推理与之前的单卡推理有显著的不同：并行化推理会将模型在不同计算卡上进行拆分，推理时计算卡之间需要互相传递数据以完成模型的完整推理。具体的拆分和数据传递模式又和不同的并行化方式相关。

一般模型的并行方式可以分为数据并行（data parallelism），张量并行（tensor parallelism）和流水并行（pipeline parallelism）。

a.数据并行

数据并行是指是指每个计算卡分别保存了完整的模型，然后输入不同的数据(在大语言模型中即为提示)分别进行计算，在模型训练时这种并行方式涉及对不同数据产生的梯度进行整体求平均等的操作需要进行卡间数据交互，但在模型推理时这些计算卡间不需要任何关联，可以看成是单卡推理的简单重复，在这里不再讨论。

b.张量并行

张量并行是指对模型进行横向即层内的划分。从之前的文章中我们已经知道，大语言模型的主体是由很多层的transformer block构成，我们可以把每一层的网络结构都拆分至不同的计算卡中一起执行。我们仍然主要考虑模型的decode推理部分。

图1. attention block的一种张量并行推理方式

更具体地，对于transformer block中的输入与QKV权重矩阵乘以及attention block部分的并行方式如图1所示。这里我们仍以LLaMA-65B为例，该模型具有64个注意力头（attention head），假设这里我们有四张计算卡，那么在一开始的一个batch的输入与QKV权重矩阵相乘时，由于采用了模型并行，我们此时就需要把QKV权重矩阵都按照输出channel方向即竖向拆成均等的四份分别存入四张计算卡中。之所以按照这个方向切分是为了将输出结果可以按照不同的head均匀分配给四张计算卡，这样每张计算卡对于attention计算只需要独立处理16个head就可以，不需要进行卡间的数据交互。每张卡拿到16个head的数据后，再分配给不同的子计算单元进行并行处理。值得注意的是这里的每个head都包含了batch size份数据，每份数据都需要单独进行attention机制的计算，如图所示。进行完attention计算后，每张计算卡将各自的16个head的输出集合在一起，四张卡上就集齐了完整的attention输出，此时并不需要将这些输出做卡间的交换。

图2. Projection层的一种张量并行推理方式

随后会进入projection和FC1/FC2的一共三次全连接线性层及其之间的一些非线性激活层等。我们只考虑需要读取权重的线性层部分。在第一次的projection线性层中，由于传入的来自attention block的激活是按照列方向切分后分布在不同计算卡上的，因此这里的权重矩阵被按照行方向切分成了均等的四份，这样各个计算卡在进行矩阵乘运算时只需要使用自己卡上的激活部分与权重部分进行矩阵乘运算即可，如图2所示。其中权重矩阵使用粗线圈出，其余则为激活。不同颜色表明不同计算卡上的数据，图3，4与之相同。各个计算卡进行完矩阵乘之后的结果由于在中间reduce维度是不完整的，因此需要进行一次all reduce操作，这样每张计算卡都具有了相同的完整计算结果。

图3. FC1/FC2层的一种张量并行推理方式

最后再进入FC1/FC2的全连接层。由图3可以看到，FC1和FC2权重矩阵分别在列方向和行方向被均等地切分在了四张计算卡上。这样一来在两次FC层中间对于整个网络中channel数最大的激活就不需要做任何的卡间数据的交互（否则代价会很大），而只需要在FC2矩阵乘结束之后做一次all reduce即可。

因此可以看到，对于张量并行推理，所需的卡间数据交互为每个transformer block里的两次all reduce，如果使用ring all reduce [1]方法，则进行一次decode推理在卡间数据交互上需要增加的延时可以近似为

其中第一个2来自于两次all reduce，第二个2来自于每次all reduce需要的reduce-scatter和all gather的两种操作，BW_p2p为卡间传输的带宽。由此可以看到，虽然这一延时与张量并行的卡数本身关系不大（(N-1)/N随着卡数增加趋向于1），但batch_size在张量并行下由于我们往往希望得到更大的计算核利用率而增加（原因参见上篇文章），因此这一延时仍会随之增长。从系统吞吐量来看，在并行度不高时吞吐量会随着并行度增加而迅速增加(batch size增加的同时latency还在降低)，但并行度增加到一定程度时这一增长速度会被卡间通信的延时逐渐抵消。

c.流水并行

图4. 流水并行的batch分布示意图

我们再来看下流水并行。如图4中右侧图所示，与张量并行需要进行的将单层算子也拆分的策略不同的是，流水并行是将模型中的不同层进行分组并分别放入到不同的计算卡中。然后流水线上靠前的计算卡处理完对应层的任务后将中间的激活数据传递给下一个计算卡继续处理直到最后一张卡处理完毕。

很显然，与张量并行一样，如果有N张计算卡，则每张计算卡需要存储的权重变成了原本的1/N，这样留下了更多的空间来存储KV_cache，应该支持更大的batch size以提高计算核的利用率。但流水并行带来的另一个问题是，为了使得流水线上的每个计算卡同时都有任务处理，这些不同计算卡任务必然只能属于不同batch的输入（图中用不同颜色表示）。然而对于这些不同batch的输入在它们结束各自的全部长度的decode推理之前，它们各自的K/V_cache数据都需要完整保存，因此整个流水并行系统中，每张计算卡都同时保留了N个batch的K/V_cache数据，导致实际的batch size并不会增加到如张量并行般的倍数，而是张量并行时的1/N。

但另一方面，如果从一个batch数据的流动来考虑的话，当该batch完成从最上方的输入到最下方的输出的整个流程后，伴随着的其实还有另外N-1个batch的输出，因为该batch每流过一个计算卡，最下方的计算卡都会有一次别的batch的输出。因此对比张量并行的话，如果二者的latency相同，而流水并行的batch size由于上述原因只有张量并行的1/N，但二者的吞吐量仍然是相同的。不过流水并行的总延时一般难以与张量并行相媲美，这是因为张量并行可以使一个batch所需的权重同时被所有的计算卡读取，带宽因此理论上增加为了N倍，但流水并行对于各计算卡上的任意batch来说，其权重的读取速度仍然是单张计算卡的显存带宽。

那流水并行的好处是什么呢？其中一项当然是模型按层来切分相对容易，另一点是卡间的交互压力更小。与张量并行不同的是，由于计算卡间每推理完一个batch的数据时只需要传输一次中间的激活（不同计算卡间同时传输的则是属于不同batch的不同推理阶段的激活），如图3中所示。因此这部分的延时可以近似为

可以看到Latency_p1是它的layers * 2 * 2 /N倍，而layers数通常是要比N大很多的，因此这一延时比Latency_p1低很多。当分布式系统变得很大的时候，张量并行相比于流水并行的延时优势会慢慢被自身的卡间通信延时给消耗掉，具体程度当然取决于batch size的设置，通信带宽等条件。

总结来说，张量并行切分对于算法的分布式拆分较为复杂，但是相对延时低，同时对卡间传输带宽要求较高，因此并行度不宜设计的太大，适合追求更低延时的场景。而流水并行对于算法的分布式拆分较为容易，同时对卡间传输带宽要求较低，对于提高吞吐量比较方便，但是延时相对较高。

因此在更大规模的推理场景下，可以将张量并行与流水并行同时混合使用，在一定范围中寻求最佳的延时和吞吐量。

依照上述的估计方法，我们可以作出一些关于不同并行度下各项性能指标的变换。我们以LLaMA-65B为例，在多卡A100下，考虑一个较大规模的并行度，为了考察张量并行与流水并行之间的异同，我们先设定流水并行度为8层，然后每层中张量并行度取1，2，4，8，10，12，14，16张GPU，由于张量并行度较大，我们先按照节点间的网卡传输带宽30GB/S作为参考。

1）张量并行度变化时的各指标变化情况

图5. 张量并行度不同时的系统延时与吞吐量变化

图5是在此设定下的系统延时与吞吐量随着张量并行度与batch size的变化关系。这里隐含的一个关系是，由于张量并行和流水并行都不会带来权重本身的重复保存，随着并行度的增加一定会带来更多的显存空间可用于存放K/V_cache，因此batch size可以随之提高以获得更大的吞吐量，但对于延时来说，提高batch size则需要同时读取更多的K/V_cache，所以延时必然会提高，所以在显存允许的情况下提高batch size对于这两个性能指标来说是互相矛盾的，我们需要根据场景需求找寻一个合适的解。

图
6. 吞吐量估计图的俯视视角

所以对于图5，这里的batch size的可取范围是被图中对应的张量并行度限制的，张量并行度越大，batch size可取范围就越大。因此如果从正上方的投影视角来看，这些曲面都将是直角三角形，如图6所示。我们回到图5，左侧图中可以看出，延时基本与batch size成正比，这是因为在大规模并行推理中，由于batch size相对较大，decode推理的大部分时间都用在了读取与batch size成正比的K/V_cache上而不是权重上。另一点值得注意的是，该曲面的顶部投影直角三角形的斜边在三维度空间中是逐渐升高的，这反映出在不同的张量并行数下，若batch size总是取最大值，即显存每次都用满的情况下，延时也在逐渐变大。由于读取每张计算卡读完进行一次推理即读完显存中的所有内容一次的时候时间应当是近似相等的，因此这里斜边的升高只能是由于卡间传输带宽不足。图中则可以看到吞吐量随着张量并行度和batch size的提高会相应提高，但提高趋势逐渐平缓，这是因为虽然batch size随着并行度在接近正比地增加，由于吞吐量=batch_size/延时，延时的相应增加会抵消掉一部分吞吐量增长的趋势。

图7. 张量并行度不同时的性价比指标变化

图7则反映了这一情形下的性价比指标。左侧展示的是平均每张计算卡产生的吞吐量，可以看到，其随着张量并行度的增加是呈现下降趋势的，若从batch size维度上看，单卡吞吐量在一开始迅速增长，然后在batch size并不大的时候就就开始趋于稳定，无法在batch size增长的大范围内获得持续的相应增长。图中右侧则是产生每个token时的卡-毫秒数即相对成本，它与单卡吞吐量应当互为反比，当然会随着张量并行度增加而上升。

2) 提高传输带宽下的变化

图8. 更高带宽下的延时和单卡吞吐量

我们假设张量并行的多节点卡间带宽可以从30 GB/s提高至100 GB/s，如果再对延时和单卡吞吐量做估计，如图8，可以看到在延时图中该曲面的顶部投影直角三角形的斜边在三维度空间中的升高趋势相比之前大大缓解了，说明此时就算batch size一直取最大，延时也不会受很大影响，这样就进一步提升了吞吐量。而由右侧的单卡吞吐量图则能看到，提高了多节点多卡传输的带宽也能使得单卡的吞吐量不再受张量并行度的增加而大幅下降，大大提高了张量并行时的性价比。

3) 流水并行度变化时的各指标变化情况

图9. 流水并行度不同时的系统延时与吞吐量变化

我们仍然回到1）中的硬件设定，不同的是此时让张量并行的卡数恒为8，而流水并行度取1，2，4，8，10，12，14，16层的并行情况。如图9中左侧图所示，可以看出即使是多节点间的卡间带宽为30 GB/S的情况下，此时的延时也基本只取决于batch size而非流水并行数，这说明延时受卡间带宽的影响非常小。因为对于流水并行来说，不管流水并行度是多少，如果我们可以不考虑卡间的传输带来的时间损失，那么理论上总的延时在batch size相同的情况下应当是一致的，而这与图中的规律十分接近。

一旦延时的分布与流水并行度关系不大，对应右侧图中总吞吐量也就同流水并行度成正比了，这是因为流水并行在推理时由于并行度越高，同时处理的不同batch的数量也就越多。对比图8可以看出，虽然流水并行的对应batch size只有图5中的大约一半，但流水并行在延时稍有落后的情况下（流水并行的卡间带宽只有张量并行的不到三分之一），其吞吐量也差不多落后相同的比例，这也反映出图4中说明的流水并行吞吐量的一个重要特征。并且对于流水并行，由于同时处理的不同batch都将占用各自的K/V_cache空间，其每个batch的batch size与并行度也就不再具有较强的依赖，由于并行度提高带来的更多显存空间只是用在了存放更多份的batch产生的K/V_cache之上，这也就是为什么图9左侧图中看上去其图形的上投影面不再是如张量并行图中的直角三角形而是接近于矩形的原因了。

图10. 流水并行度不同时的性价比指标变化

图10则展示了这种情况下流水并行的单卡吞吐量和单token相对推理成本的关系。可以看到就算在30 GB/S的较低节点间带宽下，单卡的吞吐量也没有随着流水并行度的增加而像张量并行一样有显著降低，相对推理成本也就可以保持稳定。这也直观反映出流水并行可以稳定地将系统吞吐量随着并行度增加实现线性增长。

图11. GPT3-175B的推理成本估计

如果我们想估计推理的绝对成本，可以通过图上的相对推理成本乘以一定系数获得。由于图中采用的是卡-毫秒数(chip-ms per token)来估计单位token的decode推理成本，我们只需要对该值乘以单卡在一毫秒运行内的综合成本即可。图11中是对于GPT3-175B模型在同样硬件下的流水并行成本估计，其卡-毫秒数在batch size较大时会大范围处在在5左右的区间。而单卡的实际运行成本我们参考[2]中下表1给出的估计，8卡的年度总使用成本约5万美元，单卡的每毫秒成本即为约2e-7美元，则按照我们的估计每1000 token的推理成本约0.001美元，当然加上未经估计的提示部分还会有少量升高，而且这里的总体估计与实际情形相比肯定有很多未能考虑的因素，但按照OpenAI关于ChatGPT的0.002美元每1000 token的定价来看是总体合理的。