LLM推理到底需要什么样的芯片?

2024-03-11 13:37:07 常识科普

顺着Groq公司推出的全球最快的大模型推理服务达到每秒输出500个token,如何看待这一技术?这个问题下的回答进一步延申一下,不讨论groq本身了,讨论一下LLM推理需求下对芯片和系统架构设计的基本逻辑。和过去的文章一样,我的观点一般比较激进,主要也是希望能有明确的观点和碰撞,各位看官酌情食用。

LLM本身的需求在回答里已经详细阐述了,核心需要考虑的其实就是两部分:权重以及和上下文对应的KV-Cache。

权重部分和过去的NN完全一致,不多赘述,固定大小,承载了大模型对世界的认知,在推理的时候是静态的。而上下文对应的KV-Cache是个全新的部分,承载的是对这个用户动态的权重,对上下文的认知,每个用户请求都是独特的,也是LLM独有的元学习能力的基础,这部分是传统NN所不具备的。我在三年前曾经在你为什么认为完全人工智能可以实现?中写过这种动态权重对实现AGI的途径和可能性。

战术层面,其实神经图灵机是一个非常具有潜力的idea。人的学习其实对应的不是深度学习的训练过程,而是inference的一个小步骤,人的一辈子就是一次inference过程,获取输入,提取经验,再接受新的输入,并根据之前提取的经验快速得到结论。“经验”其实就类似于attention机制中的权重,inference过程中动态生成,并用于给其他数据进行加权,而神经图灵机将经验存储起来。

只不过当时把问题想复杂了,觉得这种attention动态产生的权重得需要像神经图灵机那样复杂精细管理才能发挥起来,没成想OpenAI靠着“Scaling is all you need”的信念把复杂精细的算法管理变成了Self Attention这种简单算法+硬件暴力堆料的方式搞成了,直接O(n^2)的方式算上下文关联度就完事儿了,不需要什么精细结构从茫茫多的上下文中精挑细选相关部分,于是实现AGI的问题很大一部分简化成了芯片和系统设计的Scaling问题了,一下子就简单了很多,而芯片行业恰恰是人类所有工业体系里面Scaling技能点最牛逼的,还是我的老本行。等上下文长度Scaling起来又能把AGI涌现到什么程度我都不敢想。

十年内,芯片行业一定天翻地覆。

我去年写了一系列文章阐述芯片行业潜在的巨大变化(mackler:芯片军备竞赛新十年)、对NV的冲击(mackler:英伟达的破绽)、以及盲猜各家的下一步动作(mackler:盲猜一下NV的下一步、mackler:盲猜一下AMD的下一步),主要强调了内存带宽、互联带宽成为超过算力重要性的核心竞争力,这一点一定程度上已经有一定的行业共识。但通过groq问题的讨论,实际上更细致的取舍逻辑业内还比较混乱。

Token的成本实际上就是内存带宽成本。严谨一些讲,这里的内存指的就是“全量”存放这些权重和KV的存储介质,并不一定是最常规的DRAM内存,像Groq这样把SRAM当内存也可以,完全取决于芯片和系统设计期望的存储位置。

LLM的访存特征并不复杂,权重是所有请求所有Token都共享的,也是固定大小的内存占用量,一般是1GB~100GB量级,取决于模型规模,KV的内存占用量则是和模型相关,也和上下文长度成正比,并且每个请求独立的,并发的请求越多,KV需要占用的存储越大。今天在100GB~1TB区间上极限也就是做到100K量级的上下文长度,此时并发度往往是个位数甚至1。这个是很恐怖的,因为并发度提升一个数量级,或者上下文长度提升一个数量级,需要的KV存储也直接提升一个数量级,奔着1TB~10TB去了。而今天对于上下文长度的提升需求一定程度就是这样指数级的。

每个请求实际上对应了一连串的Token生成,并且这一串Token是串行生成的。每个Token生成的过程都需要权重和这个请求对应的KV一起参与计算,这个访存量和硬件架构无关,只和模型以及上下文长度有关。而硬件能以多快的速度完成这个访存量,就决定了Token生成的时间下限。

复用是一切花里胡哨的前提。一个系统同时并发处理大量这样的请求。权重部分的访存量所有Token都可以复用,而KV部分只有同一个请求的Token才共享,但这些Token的生成又是串行的,也没法共享。所以实际上硬件和系统设计层面来看可以花里胡哨的只有2种:

扩大并发度,提高batch大小,让不同请求的Token复用权重部分的访存。不过更多的并发请求意味更多的对应的KV存储占用,进一步意味着KV存储占用比例会大幅提升,从而导致可复用的权重部分占比实际上是下降了,最终导致权重复用的边际收益急剧降低。毕竟组一个更大的batch之后,虽然权重复用增加了,但也需要把这个batch对应的更多的KV一起读进来处理,最后大头是读的不可复用的KV。

投机执行,把同一个请求里的串行的Token通过预测并行执行,复用KV部分的访存。这个其实还是有一定作用的,不过预测准确率远不如CPU的分支预测,毕竟不是二分支预测而是三万分支预测(取决于词表大小),能准确预测个3~5步很不错了,复用度上限也就这个量级了。

其他的花里胡哨都是在贴近这个访存量下的上限,而常规的Memory Hierarchy抓访存局部性这种复用性在这个层面几乎没有任何帮助。最终硬件拼的就是能够全量放下所有权重和并发请求对于KV的那个内存介质的带宽成本。Token的质量取决于访存量。其实Token/$这种指标欺骗性很大,毕竟Token和Token的差异太大了。本质上我们需要的是相同质量的Token的成本。Token生成的质量其实抛开纯算法层面的因素,主要是取决于参与计算的权重和参与计算的KV,分别对应模型对世界的静态认知以及对当前请求上下文的感受野,从芯片和系统设计的角度看,在一定的时间尺度内,算法的差异还是噪音,整体上Token质量还是取决于上述两部分对应的访存量。以Token/$作为芯片设计的优化指标很容易带进通过降低Token的访存量来降低Token/$的歧途。如果我们可以固定Token的访存量来作为Token的质量,从而滤掉算法对质量的影响,实际上完全可以以TBps/$作为更好的量化指标,也就是带宽成本。当然反过来也简化了算法层面的花里胡哨对推理成本的影响,硬件的TBps/$固定的情况下,提高Token质量的方式就是更精细地利用访存量,让参与到Token计算的模型权重和KV更重要,除了量化压缩这些在质量和访存量上取平坦斜率的玄学外,这里面最显著的花里胡哨也就是稀疏化。

MOE算是稀疏化了权重,每个Token只访问了可能和这个请求更相关的权重,这样用更少的访存量就可以获得接近的权重质量,因为Token感受野里的世界理解以及上下文相关度更高。

KV的稀疏化未来大概率也会爆发,当上下文长度飙升到M甚至G量级的时候,只访问其中相关度更高的子集来减少访问量,其实一定程度类似于我最前面写的几年前对于神经图灵机这种精细化管理的理念。

稀疏化带来的影响是达到相同的模型质量,总的权重和KV存储量相比稠密模型会大幅度提升,同时每个Token访问量相比稠密模型会有所下降。那么对于芯片的需求就是更大的内存以及适当放缓了对于内存带宽成本的需求。当然如果稀疏度极速提升,一定程度也会对架构选择产生很多质变,例如稀疏度如果达到了2个数量级以上,并且不同权重和KV被请求到的频次呈现出显著的长尾和热点时,Memory Hierarchy可能又会重新派上用场。当然,目前距离这种状况还很遥远。完全按内存带宽计价得出的结论往往很反直觉,但并没有什么问题。因为半导体世界的存储介质基本都是带宽越高的容量越小(本质是因为容量又小带宽又小的废柴材料都被淘汰了,容量和带宽总得占一个),按照容量计价往往越高速越贵,但如果按照带宽计价有可能是反而是更便宜的,所以我们有一定动力选择传统意义上按容量计价更贵的介质。很多人觉得hbm很贵,实际上hbm的带宽性价比很高,同样groq采用的sram很贵,但sram的带宽性价比更高,甚至我们极端一点用寄存器算带宽性价比,性价比简直起飞了,所以这一点上groq是没问题的,但问题出在容量上。紧耦合的容量很重要。我们首先要保证“全量”的权重和KV都能放进去。由于推理硬件的总吞吐是可以无限复制跑一个完整模型的最小单元,所以我们对“全量”的定义可以简单理解为完整部署一份模型所需的权重存储和对应并发度所占据的存储量。其实每个请求对应的kv访存量是没有那么夸张的,一般常见的百亿模型4k上下文也就是每个请求100MB~1GB这个量级,当然长上下文以及千亿模型肯定成比例起飞,提高一两个数量级轻轻松松,但无论如何百亿模型权重总归是10GB~100GB量级。因此一个请求跨越的访存至少会覆盖到100GB这个量级的存储介质。如果大量请求并发,kv数量急剧膨胀,这个最小单元的尺寸还需要再增加一个数量级。容量可以堆,但这个最小单元必须是紧耦合的。我们本质上是要的100GB~1TB量级内存容量下的紧耦合系统的内存带宽性价比,当然这个紧耦合体系的容量需要随着算法发展以及长上下文的需求水涨船高。而内存性价比高的小容量介质应对100GB~1TB量级的紧耦合内存容量需求,需要更高数量级的介质紧耦合地连到一起。互联是耦合性的瓶颈。本质上我们前面的讨论都是围绕一个巨大的统一内存讨论的,如果我们可以无限堆料,把内存带宽性价比高的介质用独立的内存通道挂到一个统一的巨大内存下,让一个中央的计算单元可以拿到全部这些带宽,上面所有高速介质都是最美妙的选择。但巨大的统一内存不难,CPU内存可以做到这个当量,但没有独立内存通道,带宽没有同步堆料,性价比就会急剧下滑,现实的系统往往只能退而求其次,寻求分布式计算下的内存堆料,类似NVidia的单机8卡,或者groq的500+芯片,cebras的单个wafer上的海量sram。例如groq一个芯片sram是230MB+80TB/s,如果想把500+个这样的sram挂到一个统一的计算部件上,容量扩大到100+GB是没啥问题,但带宽也扩大到40PB/s几乎是不可能的,有这天顶星的技术还有啥memory wall呀,所以只能诉诸于分布式的计算部件。但实际上把500多个80TB/s带宽的节点连起来达到互联不瓶颈,对互联系统的挑战仍然是巨大的。分布式场景下的互联要求取决于分布式并行的模式。infra层面经常讨论的pipeline并行、tensor并行、expert并行都有不同的适用范围和对互联的要求,进而由于现实的互联所能实现的带宽延迟,并行度永远不可能无限扩展,一定会限制在一定尺度内,甚至非常有限的尺度内。Pipeline并行的优点在于对互联带宽要求不高,缺点是延迟巨大,而且需要相当多的用户请求来填满流水。因为实际上每个时间段服务一个请求的只有pipeline的一个stage所覆盖的内存带宽。本质上不利于groq去刷那个单用户500token/s的指标,但这是互联带宽太低,而sram太小导致并行尺度又得做得巨大无奈选择。实际上按照sram的容量带宽比,吞吐刷到几万都有可能。Tensor并行的优点在于吞吐和延迟都很好,缺点是带宽要求很高。互联带宽的要求至少需要匹配单节点的内存带宽,否则节点内计算时间按照节点内的内存带宽计算,互联通信时间按照节点间的互联带宽计算。如果通信带宽远小于内存带宽,通信会直接成为新的bound。这些是Infra层面在NVidia GPU的尺度找到的比较好的并行pattern,实际上还有更多的选择,只不过在这个尺度下效果太差就被淘汰了。其他的还有类似数据流并行,吹数据流的人很多,和Pipeline并行也有些许关系,狭义的数据流本质上还是一种局部贪心策略,对硬件拓扑和带宽也缺乏感知。广义的数据流希望通过图编译找到全局最优策略,本质上是一种把编译器当万金油的惰性做法,深度学习框架在系统调度这种比较粗放的尺度围绕数据流做了这么多年的自动并行化,最后业界主流实际上的并行策略还是预设的这些Pipeline、Tensor并行的组合,而不是编译器搜出来的自动化的并行策略,一定程度也是一种印证。而如果面对的是颗粒度更细,延迟带宽更敏感的芯片内等场景,最后肯定瓶颈就是硬件互联了,当然可以甩锅给编译器。互联的重点还是带宽。和传统CPU的NoC对延迟有苛刻需求不同,LLM对互联系统的核心诉求还是带宽,本质上是为了把分散的内存带宽能够发挥出来,不要因为互联系统的带宽阻塞了内存带宽。知乎上有很多做芯片的人吹dojo和tenstorrent那一类mesh互联,觉得对称性无与伦比,而且可以无限扩展铺满整个平面。但mesh根本不是个什么新东西,教科书上几十年的老黄历了。mesh在高带宽需求也是个相当糟糕的互联,除了物理走线短之外几乎一无是处,尤其在LLM这种带宽优先的场景下,连极其规整的集合通信都没法把宝贵的互联带宽资源打满。对称性也无从谈起,毕竟不是一个无穷大平面,边界和中央的流量是完全不均衡的,简直是并行策略的噩梦。拓扑层面ring、torus、fat-tree、dragonfly等等都是对称性好得多的拓扑,对宝贵的物理带宽资源也都能更好地利用起来。对应的并行模式,Tensor并行也优于Pipeline并行进一步由于数据流并行的策略。互联带宽层面,无论如何都逃不掉一个接近内存带宽级别的物理带宽。所以事实上又和前面形成了一个反过来的取舍逻辑,高速低容量的存储介质具有更好的内存带宽性价比,但因为容量小了,要形成一个LLM可以运行的紧耦合体系,需要的分布式节点就更多了,而节点多,单节点带宽高,都对于互联系统的拓扑和带宽提出了极端的要求,实际上又反过来促使我们使用低速高容量的存储介质,因为可以用更少的节点达成LLM可以运行的紧耦合体系,同时单节点带宽又不是那么高,互联系统还是比较容易达成的。而且互联世界的光谱也是节点越多,空间尺度越大,带宽越小,带宽成本越高。这种相反的取舍逻辑最终会达成一个平衡点。本质上取决于内存体系的带宽容量和成本的光谱与互联体系的带宽与成本的综合考量,这两个都是非常庞大的光谱,架构选择的空间非常宽泛。但总体来讲,内存和互联这两个光谱有一些基本的交点:同一个物理尺度下,局部的互联带宽上限通常还是小于内存带宽上限的。这里“内存”同样是广义的,芯片内部的sram与芯片内部的互联、芯片间的互联和片外内存带宽等等。毕竟同一个物理尺度下如果计算节点之间的带宽都能飙起来了,计算节点和存储介质之间的带宽为啥不能也用这种技术飙起来呢?所以groq这种选择一个很尴尬的地方在于,内存带宽选择的sram是芯片内的尺度,而互联又不可避免走向更低速的芯片间、甚至板卡间、服务器间的尺度,互联自然远远跟不上内存带宽,从而把内存的优势打回原形。进一步讲,100GB~1TB这个容量尺度的系统,互联大体上肯定是在芯片&板卡间这个级别的物理层的,那么内存大概率还是选择这个层级,也就是DRAM更合适,当然了,DRAM的形态光谱也非常宽,芯片&板卡间的互联光谱也非常宽。互联带宽没有显著高于内存带宽时,任何网络拓扑都没法扩展到非常多节点的紧耦合。简单来讲,如果没法达成互联带宽没有显著高于内存带宽时,分布式节点太多一定会造成带宽问题。NVidia创造的紧耦合系统目前主流的尺度是单机8卡,当然也做过16卡的形态。大力出奇迹的规格是256卡的GH200,但是对应的全局带宽堆得相当夸张,否则只是个松耦合系统,松耦合系统自然无所谓,现在的千卡互联万卡互联都是全局带宽收敛的,不在我们讨论范围,之前写过一篇mackler:AI云还是AI大型机?讨论过这个紧耦合系统的问题。这个并行尺度对于其他物理尺度下仍然适用。这个尺度对LLM芯片的影响实际上还是如何实现紧耦合的100GB~1TB这个量级。这个量级下紧耦合分布式的碎片化程度不能太高,如果搞到几百甚至上千的分布式,互联带宽根本做不上去。这两个交点一定程度上能帮我们收敛出大致的量级范围。单节点统一内存10~100GB量级,节点数量1~10量级,互联带宽接近单节点带宽的形态,此时系统的总内存带宽和总成本的比例越高越好。在这个范围内,还是有大量空间的。当然了,毫无疑问,NVidia的诸多选择也在这个范围内,其实这是必然的,因为上面说的很多尺寸规模都是算法在英伟达的互联、容量、带宽规格下筛选出来的,毕竟今天NV的紧耦合节点就是DGX,所以LLM的紧耦合尺寸基本都在围绕在DGX的总显存带宽量级。而NV的DGX总带宽性价比是16~25TB/s / $200k~350k。内存带宽与互联的配平比内存性价比更重要。当然比较这个指标的前提其实是上面讲的范围尽量落在这个范围内,否则一系列严重的互联问题会直接直接打崩所有内存性价比的优势。而落在这个范围之后,实际上NV的利润率是比HBM性价比富余得多的空间。所以今天其实对于LLM架构设计而言,这方面的配平远比单纯介质的性价比重要得多。一定程度上讲,大多数芯片架构师都是微架构背景出生,对微架构或者架构形态的重视程度更高,很多对微架构技术形态或者工艺技术的重视程度更高。其实大家不要太过技术思维的从技术路径和技术形态的角度选择自己的屁股,体系结构最重要的是tradeoff,是尺度范围。一个架构形态在合理的tradeoff下可能惊为天人,在另外的场景下某些数字超出了合理量级就会产生严重的问题,不可能有永远完美的架构形态,选择架构形态作为屁股是很危险的。无论是存算一体、数据流还是sram、3d dram等等都不是LLM芯片最关键的东西,而是架构的配平,这个配平在传统NN上是一种,在今天的LLM上又是另一种,尺度的数量级非常重要和敏感。Groq就是单节点做太小,并行度太高,互联带宽太低,当然第一性原理的内存带宽性价比肯定是OK的。数据流存算一体架构在传统NN时代虽然软件层面还是有太多问题暂时不讨论,但硬件架构选择层面还是有不错的竞争力的,无数公司例如graphcore、cerebras、tenstorrent等等都一起涌入了靠SRAM干掉内存访问的路径。但graphcore的CTO去年就公开表示了这个架构不适合LLM了,也表达过类似DSA已死的观点。除了groq这种激进的方案外,目前的常见的柔和改良方案基本都是考虑在这样一个mesh众核的架构基础上搞3D DRAM来解决sram太小的问题。但实际上除了软件层的问题外,互联配平的问题仍然非常严重。因为靠很多人觉得优雅的mesh网络是完全没法匹配3D DRAM的内存带宽的,同时因为众核的数量往往是100~1000量级,虽然是一颗DRAM芯片,但实际上是大量容量很小的小DRAM bank通过mesh网络连起来,mesh网络糟糕的带宽利用率以及100~1000这种夸张的碎片化程度,实际上对mesh网络上全局带宽提出了极其变态的要求,即使是芯片内部也很难达到。最后大概率只能回退到糟糕的pipeline并行面对和groq类似的内存带宽利用率问题或者回退到更糟糕的数据流并行并且把问题丢给“万金油”编译器,我之前在mackler:专用架构与AI软件栈(6)中详细讨论过这类架构的软件问题,其实软件的困难在今天LLM上不仅有耦合性的老问题,还有物理带宽事实上完全不够的问题。存算一体有一个很大的缺陷就是对互联的重视程度太弱了。其实存算一体的架构我在学校的时候也做过很多研究,存算一体听起来很美好,数据在哪里就在存储数据的地方计算。但问题是计算不一定只有一个操作数,很多时候有多个操作数,这些操作数不在一起的时候还是需要通过通信来传递数据。最后发现通过存算一体搞出了实际上最极致的内存带宽,但是互联完全跟不上,mesh对于大带宽需求的互联真的太糟糕了。我之前在ASPLOS上发表的一篇关于存算的架构设计最后甚至是用FPGA芯片架构里面的可编程连线这种暴力的方式来实现尽可能高的互联带宽来配平。其实你不能说这是存算一体的问题,因为这个概念是很模糊的,NVidia把HBM和GPU合封到一个芯片里其实也算某种意义上的存算一体,当然不用纠结是存算一体还是近存,都是文字游戏(2.5D合封不算的话3D合封算不算?SRAM和ALU也是分开的算存算一体还是近存?)但NVidia是实打实把NVLink的带宽实打实打到显存带宽量级。而现在大量mesh众核的互联瓶颈也是实实在在的。未来长远来看,MOE和KV稀疏化是加速AGI Scaling从算法层面最有效的途径,其实也是从粗放式逐渐过渡到精细化管理静态和动态权重,从而可以创造在芯片Scaling基础上进一步更快加速超长上下文和超大模型的低成本Scaling,而对于硬件的容量需求会进一步扩大,同时随着稀疏化程度提高,使得Memory Hierarchy可能重新变得在系统层级更有意义。

版权说明: 本文版权归原作者所有,转载文章仅为传播更多信息之目的,如作者信息标记有误,请第一时间联系我们修改或删除,多谢。