闲聊AI芯片和存算

闲聊AI芯片和存算

引言

存算这个领域早在1990就开始被讨论^{[1][2][3][4][5]}（最早的建议可以追溯到1970^[6]），当时的主要思路还是停留在怎么把计算的logic和DRAM做到一起，受限于存储器工艺和逻辑工艺的不兼容问题，这方面的研究进展并不是很顺利。同时memory wall的问题对于主流计算机技术路线来说，完全可以通过堆memory hierarchy解决，从两级Cache到三级Cache到四级Cache......在相当长一段时期里存算并不是一个像现在这么火热的领域。

在这之后发生了哪些事情改变了这一格局呢？首先是AI技术的强劲需求端驱动，从2012年伟大的AlexNet为始，AI轰轰烈烈的十年之期开始了。一条AI领域几乎颠扑不破的真理是the bigger the better，更深更大的网络就更加牛逼，这条路一直通到了我们今天见到的百、千、万亿级参数巨无霸：LLMs。巨型网络的training和inference所需要的算力也急速膨胀，我们不妨看看Alisa Su在ISSCC2023上做的《Innovation For the Next Decade of Compute Efficiency》演讲里是怎么说的：

As an industry, we are now setting our ambitions on the technology required to achieve zettascale computing. Assuming the historical rate and pace of advances in performance and efficiency are maintained over the next decade, we could achieve a zettaflop of computing in ~10 years — but it would require a half a Gigawatt of electricity, the equivalent output of half a nuclear power plant (Figure 1.1.6). This is clearly not a viable solution!

如果只是简单的scaling方式来增加算力的话，那按照目前的算力需求增长水平，10年之后我们就得搞个核电站来给超算供能了。所以不仅仅是峰值算力问题，Energy Efficiency也是一个很重要的问题（用更少电提供更大算力），而这恰好是存算技术的初衷，即通过减少memory到计算之间的数据吞吐来数量级的提升Energy Efficiency。

其次是其他各种技术的涌现：首先是SRAM存算^[7][8]，一种观点是SRAM存算发源自BCAM/TCAM，后来为了AI的应用做了很多MAC方面的优化^[9]。相比DRAM存在的工艺兼容问题，SRAM显然是一个完全CMOS工艺兼容的方案，唯一的缺点可能就是面积太大，但架不住工艺节点scaling的威力，这股风先在体系结构那边的会吹了起来，接着就是电路级，这个转折点基本上就是2017-2018。我考古的结果是ISSCC2017^[10]尚且没有独立的存算session，到了ISSCC2018^[11]蹦了个Session31 Computation in Memory for Machine Learning 出来，有两篇SRAM的存算工作和两篇RRAM的存算工作，ISSCC2019^[12]则是给了一个SRAM and computation-in-memory的大Session，但里面做纯SRAM改进工作的和做存算工作的对半开。至于2020至今存算的发展盛况可以看我这篇博客(2023存内计算/近存计算宏电路进展调研 - sasasatori - 博客园 (cnblogs.com))。

其次是新型存储器，除了赛道里的老玩家ReRAM（从18年就在存算赛道里了），还有后续加入的FeRAM，STT-MRAM，PCM等一众玩家。

最后还有最老牌的DRAM，其实2010年商用3D堆叠芯片出现以后大家就对这个方案又开始感兴趣了，这样就消除了之前存储器工艺和逻辑工艺冲突的问题，只需要把存储器和逻辑各自做好之后TSV封装到一起即可。就目前来看三星^[13]和AMD都在这条技术路线上整大活，把HBM和PIM搞在一起。

总的来说就是存算现在被炒的火热。但我们都知道任何一个新兴事物都必然有其质疑者，让我们来听一听批评的声音。

一些批评

这是在涂锋斌博士的知乎文章(ISSCC'22奇妙之旅 - 知乎 (zhihu.com))底下翻到的一个评论：

这个评论发布的时间是2022年3月，客观的说我在那个时间段上比他还要悲观。别说AI芯片了，我对AI本身都不看好。是的没错，一直到22年年底为止，我一直是一个马库斯式的人物，我对目前以神经网络模型主导的AI技术的态度一直是看空胜于看多，我一直觉得这东西首先是个黑盒，应用极度受限，其次没有完整的理论，全靠玄学调优，因此注定就是个泡沫。这是我一月份上完自然语言处理课之后写的综述里的summary：

语言模型的发展脉络令人印象深刻，当历史推动主流范式从统计语言模型转向神经网络语言模型时，其过程既令人惊愕却又是如此自然，使得我在学习课程时深深的为之震撼。如今的超大规模预训练语言模型已经能够用于生成以假乱真的文章、诗词，机器译文也非常流畅，我们不由的期待将更多的自然语言处理的任务通过超大规模预训练模型进行统合（实际上GPT系列正有此意），同时通过进一步扩大模型和数据量从而实现性能的进一步提升，一切都看上去非常promising。

但歌舞升平的同时，我们也得看到神经网络方法的一些根源性的问题：可解释性差，我们无法像传统方法那样追究每一个结构，每一个步骤的具体数学原理，这导致其难以应用于一些要求高可靠性的系统，并且难以debug；极度依赖于数据集，质量差的数据集对神经网络的性能会产生严重的影响，甚至导致其彻底无法工作，而数据量不足同样对神经网络有着致命的影响；有知识而无智慧，尽管学习了海量的信息，但仍然缺乏架构化的知识，甚至缺乏常识。

当前知识图谱的研究受到了很多人的关注，人们希望将结构化的知识图谱和数据驱动的深度神经网络相结合，从而使得新一代的人工智能能够具备强大的自我思考和推理能力。

也许神经网络也会像传统方法那样，有朝一日来到它的尽头，而届时另一种全新的方法将会引领学科步入下一个新的纪元。

然后2022年年底我先后见证了stable diffusion和ChatGPT。我意识到我错了。我的意思不是指神经网络模型解决了自己的这些根本技术问题，而是我在以一种落后于技术本身的观点看待一门先进技术，应该被改造的是我看待这个新事物的方式，这更近似于一个哲学层面的理解问题的方式，我的意思是指，如果好好地去看一下OpenAI的技术领导者的采访视频的话，会发现他们对自己已经做过了充分的自我批评，我前面批评他们的问题他们都清楚，但是并不妨碍他们站在一个更高的角度上（这个角度是我从未采取过的）去审视自己的work，他们在一个新的高度上发展了一套新的看待AI可靠性和有效性的方法论，并在那个高度上做出了成绩，而直到我改观之前，我都停留在上世纪末的那批diss神经网络的人的高度上，固步自封。

回到AI芯片的角度上，我想表达的问题是，22年年底的重磅炸弹改变了很多东西，就目前LLMs带来的算力需求飙升来说，英伟达无疑首先是最大赢家，但其他选手也在面对迅速做大的蛋糕，英伟达并不能独自全部吃下迅速扩大的市场，君不见Jim Keller都下海整活了，我去看了一下Jim Keller的访谈，同样是醍醐灌顶。另一方面来说，现在这类强大的模型因为其算力需求过大只能放在云端，但是必然在未来会为了安全性或者实时性向端侧逐渐转移，这是一个非常困难的挑战，但同时也提供了无比巨大的机会。未来我们如果想要用上能够加载强大智能的机器人（这是Tesla的愿景），这个问题是无论如何都要被解决的。

但最重要的问题是提升自己的认识，和领域里最顶尖的一拨人站到一个高度上去，你的critical thinkings别人大概率不仅有过，而且比你更加深刻，但是人家看待新事物所建立的方法论大概是你从未有过的。

回归正题，我在这一节写了一下我对AI芯片的批评的一些理解，看上去有点跑题，这是因为存算本身是AI芯片解决方案的一种，还有许多其他方案竞争者，但是首先我得把母题给回答一下，为什么AI芯片有必要存在。

AI芯片的格局

就目前情况来说，Nvidia稳稳站住先发优势，软件生态护城河极深，几乎算是一己之力定义了整个行业。Google，IBM，Intel和AMD是相对成熟的追逐者，再这之后跟着的一众startups。这一节我们重点来聊聊Nvidia究竟做对了什么，让他能够稳稳坐在王座之上。

首先是起步阶段，我们都知道最开始Nvidia主要做的是游戏显卡，2006年Nvidia发布了CUDA架构为GPU提供了编程手段，这是Nvidia的第一着妙棋，彼时其他的显卡厂商并没有对GPU的可编程性做出足够的重视，这直接导致了转折点到来之时他们发现自己完全处于在局外的状态。这个转折点就是2012年的AlexNet，AlexNet采用两路GTX 580 3GB GPU并行训练，解决当时使用CPU训练速度缓慢的问题，而能够这么做的前提条件就是CUDA提供的GPU可编程性。

接着是真正意义上构筑护城河的阶段，Nvidia敏锐的战略意识，催动其深度参与了整个AI软件栈的构建，目前主流的三大AI框架：Pytorch，Tensorflow，MXNet都在Nvidia的硬件上可以得到良好的支持。基于Nvidia的CUDA生态，主流的AI框架都得到了良好的支持，AI开发者对于硬件的选择自然也毋庸置疑的以Nvidia的GPU硬件为主。与之形成对比的是Google希望将Tensorflow框架与自家的TPU硬件形成绑定，然而由于Tensorflow在AI框架中的竞争中逐渐式微，同时又缺乏对其他AI框架的良好支持，导致其最终在整个竞争过程中严重落后。

最后是不断地自我创新，我们提到过Nvidia的历史，他们创业之初根本没有AI加速这一概念，发家之际靠的是游戏业务，中间很长一段时间又侧重在加密货币业务上（狠狠的卖矿卡），Nvidia在硬件上开始为AI做专用化设计要追溯到2017年的Volta架构，在GPU中引入Tensor Core结构实现对AI的专业加速，这对于传统的GPU架构是一次很大的变化，也就是相较于传统的SIMT的架构，从此Nvidia的GPU走上了SIMT+DSA的混合，之后的Nvidia又数度创新，包括创造新的数制，挖掘稀疏性，到了H100架构更是直接对Transformer架构提供了Transformer Engine以示重视。从这里可以再次看出Nvidia的战略意识，敢于为了开拓新的业务对自家的成熟架构动大刀子进行架构革新，使得Nvidia GPU的硬实力相比其他AI专用芯片毫不逊色。

与Nvidia相比，其他后入局者面对着的问题很多，一方面是Nvidia积累十几年的软件护城河，另一方面又是硬件性能上难以做出绝对优势，最后芯片本身是个规模效益很明显的领域，越大规模的生产越能把成本给降下来，不到Nvidia的体量想和他打价格战难上加难，何况Nvidia自身已经是整个半导体领域Fabless市值的绝对第一了。

但上一小节我们提到过了，一方面是蛋糕总量在做大，Nvidia无法一家吃下全部市场，何况还有许多市场以外的因素（比如美国对中国的AI芯片禁售），市场上剩下的空间是新兴AI企业的绝佳机会。此外Nvidia的AI芯片业务主要集中在云侧，而端侧的AI芯片市场还没有被充分定义，还是一片待开发的蓝海。从我的角度看来，一个聪明的startup有几个可以选择的破局方向：如果要去卷云端，软件端需要花大量的功夫，不仅是支持主流AI软件框架，最好能够有自己独立的AI软件工具，建立起对Nvidia的比较优势，进而为自己的硬件提供机会。端侧的生意看起来机会更多一些，但是重点在于能不能做对应该做的事情，这里包括如何构建软件生态，如何找对合适的AI Application，如何定义端侧AI芯片行业等一系列复杂的问题。

一些引人注意的新架构

现在我们来谈谈一些引人注目的新架构。2022年的hotchips上dojo^[14]和cerebras^[15]这两个工作。这两个工作的共性特点都是将scale摆到了一个极其重要的位置上，以小的node为单位，通过2d mesh网络将node链接起来构成大型算力芯片。使用分布式SRAM作为片上存储，因为每个node实际上是平等的，所以和GPU里还存在层次性的Cache不同，通过GB量级的总片上SRAM存储直接取消掉了多级Cache的概念，每个node的计算单元直接和node内的SRAM存储耦合。

在具体的微架构上，单个dojo的node可以看到包含了1.25MB的SRAM，标量计算单元部分和向量单元部分，以及负责node2node的NOC Router。

cerebras的单个node的SRAM就小很多了，只有48kB。

但他们做的一个有趣的事情是把细粒度非结构性稀疏给放到了架构设计的考量之中：

这两个工作虽然是由公司提出，但我作为一个正在混学术界的人仍然感觉到了很多很有意思的东西，在下一个小节我将重点谈一谈我的见解。

一些个人不成熟的想法

首先是SRAM存算的机会。尽管之前没有专门关注过hotchips2022上dojo和cerebras的工作，我自己在23年2月份开始做的多核存算芯片和他们share了许多理念，也采用了2d mesh连接多个node的结构来实现可拓展性。此外大规模片上SRAM的使用，很大程度上是SRAM存算极佳的应用舞台，因为计算单元和SRAM被摆到了一个相当近的物理位置，从架构设计的角度来说，所谓的SRAM存算宏其实就是存储单元+SIMD单元，在node中应用SRAM存算技术使得SIMD部分计算的访存Energy Efficiency得到进一步的提升应该是非常有吸引力的选项。

其次是关于巨型芯片的问题，我对cerebras的晶圆级芯片的概念并不是很感冒，做过ic的人都知道芯片面积和芯片良率之间的反比例关系，这样大面积的芯片最终坏点会有多少，如何有效的应对这个问题，目前来看cerebras还没有给出一个令人满意的回应。对我来说，反倒是通过chiplet将小芯片封装成大芯片看上去更加具有吸引力，这里涉及到的问题是，以学术界的经费来说，制造大面积的芯片是极其困难的事情，如果想要染指大算力领域，那么制造小芯片然后通过封装的方式构成大芯片是一个有吸引力的方案。当然这也涉及到chip2chip之间的指令/数据通信如何设计的问题。

最后是关于新型存储器。我自己在以前的博客(基于存算一体技术的新型计算机体系结构 - sasasatori - 博客园 (cnblogs.com))里也做过简单的设想，将新型非易失存储器给引入到AI芯片的设计中来，这样带来的优势是，我们不需要在断电开机之后为整个超算重新load一遍数据（这种大规模片上SRAM的load开销将非常夸张），同时如果突发断电，也不用担心数据的丢失（如果要定期将大规模SRAM的数据写到片外，那也太噩梦了）。如何在这种小node内部加入新型非易失存储器并解决相应的体系结构上的理论问题，也是一个我非常感兴趣的方向。

引用