大模型推理:从模型分析到计算优化(二)

发布时间:2023-09-04 22:00

 

 

在上一篇文章中,我们对大语言模型的推理模式和细节过程做了介绍,并且引申出几个推理时需要关注的几个关键的指标,即延迟(latency)、吞吐量(throughput)和推理成本(tokens/$)。延迟是指前后连续产生关联token之间的时间,吞吐量是每秒能够产生的新token数,推理成本则是考虑设备成本与运营成本之后使用单位成本(1美元)可以产生的token数。这次我们将分析这些指标背后的评估方法。

 

1. 性能指标计算

 

由于吞吐量完全依赖于延时和人为参数batch size,我们主要分析延时就可以获得吞吐量的表现。模型运行的延时一般可以从两个方面进行考虑,分别是读写延时和计算延时。而在较大规模的模型计算中,读写和计算过程通常可以分布执行并相互遮盖,最后我们看到的延时将是其中数值更大的那个。比如对于一个卷积运算,它所需的数据包括待卷积的张量和卷积核,它们都来自于外部存储器,我们可以将它们的数据总量除以外部存储器的带宽来估计这部分的读取延时,而计算延时则可以用总的计算量除以对应计算核的算力来获得。

 

因此我们可以看到,推理过程延时既和模型的计算需求有关,又和底层硬件(比如GPU或AI加速卡)的处理能力有关。下面的分析假设底层硬件由计算核和外部的存储器组成,其中计算核的向量处理能力为Pvector,矩阵乘处理能力为Ptensor。由于LLM推理对于存储器带宽要求很高,这类假设外部存储器为HBM(high-bandwidth memory),其访问带宽为BWhbm

 

按照上篇文章所介绍的,模型的推理分为用户输入提示后产生完提示对应的K/V_cache阶段也就是context推理阶段和随后的依次产生新的输出的阶段也就是token推理阶段。前期的context推理是一次性完成的,其处理效率强烈依赖于所采用的矩阵乘算法在各尺寸矩阵上的优化情况。我们主要来看后期的token推理阶段,它是模型推理的主体部分。在这部分中,每个layers中都包含一次的transformer block又是其计算主体,其计算流程如图1所示,具体过程上篇文章已经介绍过。

 

 

图1. 大语言模型主体transformer block推理流程

    

  a)矩阵乘部分

  图的上方的橙黄色块展示了各个环节所需要的权重,它们通常保存在计算卡的HBM(high-bandwidth memory)也就是通常所说的显存上。例如对于LLaMA-65B模型来说,如果权重部分采用fp16格式进行保存就需要65G*2=130G的HBM空间。当模型进行decode推理时,batch size个input token向量以及之后的中间向量需要组合成矩阵与所有的权重都乘一次,因而也就需要将这些所有的权重都从HBM上读取一次,这些计算中由权重读取来的延迟就容易估计出来,我们记为latency1,即

 

 

 

 

    其中bytew和BWhbm分别指每个权重元素所占的字节数和HBM带宽。此处为什么没有考虑矩阵乘时由输入tokens造成的读取代价呢?这是因为在GPT类模型的decode推理中,每次的输入tokens只是单步的,就算这样的tokens有batch size个,每步矩阵乘时其通常的体积batch_size*hidden_dim*bytea相比于通常的权重矩阵的体积hidden_dim2*bytew相对于权重仍然是可以忽略的(因为受限于要保存大量的K/V_cache,batch size通常比hidden_dim要小多了)。

 

    而将这些权重读取到计算核上进行计算时的计算时间也可以按照矩阵乘的计算量和计算核的矩阵处理能力估计出来,其中计算量应为batch_size*2*weight元素数量=24*batch_size*hidden_dim2,其中的2是由于矩阵乘运算内做reduce的数量和乘法的数量一致。因此这部分的延迟为

 

 

 

 

对于矩阵乘来说,这两部分的延时可能都相当可观,不过计算核的利用率往往会受限于传输带宽,而这一利用率可以用latency2/latency1来大致获得。 但这两部分呈现的依赖特征又有很大差别。由上方公式可以看到,由权重读取部分带来的延迟始终是恒定的,与batch size无关,这是因为来自一个batch里的多个输入会同时执行与同一块权重进行的矩阵乘,因此权重始终只需要读取一次。而此时的计算量显然是和batch size成正比的。另外就像之前提到的,矩阵乘的执行通常是按照一定大小进行切块后进行的,这两部分过程其实是相互交叉并行的,理想情况下我们可以将它们看作可以完全并行,在这样的假设下总的延迟就取决于两者中更大的那个。而实际在batch size较小时,latency1通常会大于latency2,也就是计算核会处于有闲置的状态,此时由于前述的两种延时对于batch size呈现的不同特性,我们通过提高batch size就可以提高此时计算单元的利用率

 

  b)attention部分

 

  但是对于attention部分的运算来说则完全不同。相比于其他的神经网络部分来说,大语言模型中的attention部分具有一个显著的特性是它需要保存K/V_cache,这两个张量与推理的总token数量成正比,即正比于batch_size*sequence_length,而当这一数量较大时,K/V_cache的大小会变得相当可观!还是以LLaMA-65B模型为例,如果推理的token长度达到模型允许的最大值2048时 ,它们的总大小将高达

 

 

 

 

其中batch_size取32,seq_length即为token数2048,bytea是每个激活的元素所占的字节数,这里也取2。170 GB这一体积甚至超过了模型权重本身的大小(130 GB)。因此对于这部分的读取造成的延迟将会是至关重要的,我们记作latency3

 

 

 

    如前所述,这一部分的矩阵向量乘计算的计算量比矩阵矩阵乘要小得多,它的算力通常要小很多,但是同样地,它需要执行的任务量也要更小。在这里它需要执行的计算量为2*2*batch_size*seq_length*hidden_dim,因此需要的延时为

 

      

基于同样的原因,这两部分的延时也可以近似地取最大值来分析。值得注意的是,这部分的计算量始终是K/V_cache读取量的两倍(因为都和batch size成正比),因此batch_size改变不了这部分的计算与读取的时间之比,计算核此时的利用率也将大致维持在一个恒定的水平,这就与之前的向量矩阵乘部分呈现的与batch size相关的特性有很大的不同。

 

   c)总体分析                                                                                                              

我们知道,对于上述的latency1中的权重矩阵来说,不管batch size是多大,它们在一次推理中都只需要被读取一次,而就像前面提到的,这里做矩阵乘时的总读取量近似等于权重读取量(写入量相对于权重读取量也可以忽略),因此提高batch size理论上对这部分latency的影响几乎可以忽略,则这部分的吞吐量会和batch size呈现近似的正比关系(在latency2不超过latency1的前提下),同时也意味着计算核利用率的提高。

 

 另一方面我们需要保证K/V_cache可以顺利容纳在HBM除了权重以外的剩余空间中,这会使得在给定的硬件环境下batch size以及模型允许推理长度即sequence length都有一定的限制。

 

 更大的batch size以及sequence length使得K/V_cache体积变大,也就使得attention部分的读取延时边长,推理的总延时也就变得更长。这就意味着在sequence length不变的情况下,在允许范围内(受K/V_cache体积限制)增加batch size在一开始很可能可以带来吞吐量的迅速提升和延时的增长,此时延时与吞吐量是互为权衡的两个性能指标。而随着batch size继续增加到一定程度,总延时的增长会逐渐抵消batch size增加带来的吞吐量收益(吞吐量正比于batch_size/latency),吞吐量会增长放缓并接近某个极限,尽管此时的HBM空间可能还允许batch size的继续增加,但并不会带来吞吐量的进一步收益,同时还在增加延时。但如果HBM本身较小,可能batch size天然无法增加到这一程度,就看不到这里所说的吞吐量极限,这种情况下从吞吐量的优化角度考虑就使batch size越大越好。

 

2. 大语言模型在单卡GPU上的性能指标分析

 

 下面我们使用单卡GPU模型来具体分析上述的性能指标关于模型参数和硬件参数的影响。硬件参数上这里的理论分析可以带入任何GPU的硬件参数,我们这里以A100(80GB HBM2e版)为例,其HBM带宽为1935 GB/s,矩阵乘算力为312 TFLOPS(fp16),向量处理能力624 19.5 TFLOPS(fp32)。

 

图2. LLaMA-13B推理性能估计

 

a)主要性能指标与batch size的关系


    图2中展示了LLaMA-13B模型的推理性能估计结果,横坐标为batch size,纵坐标分别为延时(latency)和吞吐量(throughput)。我们可以看到,随着batch size的增加(上限35为推理长度为2048时的极限),K/V_cache加载量随之增加,因此延时基本呈线性增长,此时吞吐量也不断增长,但增速会被延时的增长抵消一部分,因此增速呈现不断下降。

 

 

 

图3. LLaMA-7B推理性能估计

 

而在图3 LLaMA-7B推理性能估计中,由于模型自身权重所占的HBM总占用量更低,受K/V_cache的更大影响,此时吞吐量随着batch size的增加而增速放缓的迹象更为明显,理论上在允许更大的HBM上,吞吐量的增长最终会趋于停滞,此时再增加HBM和batch size就不会带来吞吐量的收益了,同时还会增加延时。另外在多卡环境下由于会引入卡间的通讯代价,这一趋势会变得更为明显,具体将在下篇文章中讨论。

 

b)主要性能指标与sequence length的关系

 

 如图4所示,不光是batch size,模型推理长度即sequence length对推理性能的影响也非常大。由于K/V_cache随着推理长度的不断增长而相应增长,而模型推理出每个token都需要将K/V_cache全部读取一遍,因此延时也将不断增长,吞吐量则相应下降。这也是大语言模型推理时的一个重大挑战。

 

 

 

 


图4. 推理长度对性能的影响

 

c)矩阵乘计算核利用率


    另外在这些模型设定下,矩阵乘计算部分的计算核利用率如何?我们通过比较latency2和lantecy1的比例计算了不同batch size下LLaMA-7B的decode推理部分的矩阵乘中的计算核利用率,如图5所示。可以看到,在单卡A100 80 GB的HBM空间下,无论batch size如何扩大,计算核的利用率也难以超过一半。我们难免会想在多卡系统下该利用率能否进一步提高?下篇文章中我们还将分析多卡时的情形,并会给出更普遍的规律。

 

图5. LLaMA-7B计算核利用率

 

 

d)相对推理成本与延时的关系

 

尽管在单卡上我们无法考察大规模并行推理时的推理成本与卡数的关系,但我们可以先考察一下此时的相对推理成本与延时之间的关联。由前述规律可以知道,为了实现更小的延时,我们需要设定更小的batch size,但此时latency也会降低,因此单位token的推理成本肯定会上涨。我们用每推理出一个token所需的卡-毫秒数(chip-ms)来表示其相对推理成本,结果如图6所示。

 

 

 

图6. 不同大小模型相对推理成本与延时的关系

 

图6中我们展示了在固定推理长度(2048)下LLaMA-7B和LLaMA-13B两个模型的相对推理成本与延时之间的关系,可以看到随着延时的增加(通过增加batch size),两个模型的推理成本都在下降,但前期的成本下降更为显著,后期的成本下降变得越来越不明显。这体现了延时与推理成本之间的相悖关系,在具体场景下则需要综合考虑二者的关系来决定一个合理的batch size以平衡场景下所需的延时和成本需求。另外可以看到,在低延时下两个模型的推理成本相差较大(延时同为15ms左右时),随着延时的增大(通过提高batch size),二者的推理成本在下降的同时也变得相对接近了。

 

小结:通过给出几个常用性能参数的估计方法,我们估计了在单卡系统上典型的大语言模型的推理性能。可以看到典型的性能如延时、吞吐量和推理成本等可以和batch size还有sequence length等模型参数建立起较为清晰的关系,在权衡延时和吞吐量以及推理成本时我们则往往需要在中间做一些取舍。多卡系统上,由于会加入模型切分,流水并行以及卡间通讯,一个大的系统的推理性能会呈现更复杂的变化,我们将在下一篇文章中带来相应内容。

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

上一个: 大模型推理:从模型分析到计算优化(三)

下一个: 多模态大语言模型:让AI看图说话

近期文章

AI 智能体:应用前景与对算力需求的影响

人工智能技术的迅猛发展,尤其是 AI 智能体这一核心部分。本文讨论了AI智能体的定义、应用前景和其对算力需求的影响。

2023-11-13

查看更多