利用零值的AI芯片突破

对人工智能模型而言，算力规模往往决定上限。

尽管一些AI从业者提醒，像大规模语言模型（LLM）这类技术的效果提升可能正在出现边际递减，企业依旧持续发布规模更大的AI产品。Meta 最近推出的Llama模型就配备了高达2万亿个参数。

模型越大，能力通常也越强。但与此同时，它的用电量和运行时间也会随之上升，进一步推高碳排放。为降低这些代价，人们转而采用更小但能力更有限的模型，并尽量用更低精度的参数配置。

不过还有另一条思路：既保留大模型的高性能，又能显著减少运行时间与能耗。关键在于“善待”大模型内部的零值（This approach involves befriending the zeros inside large AI models）。

在许多模型里，多数参数（权重与激活值）其实都为零，或接近于零，因此把它们视作零通常不会明显牺牲精度。这个现象称为稀疏性（sparsity）。稀疏性给节省计算资源带来很大空间：原本用于零的加法或乘法不必再做，相关计算可以直接跳过；同时无需在内存中反复存储大量零值，只需记录非零参数即可。

遗憾的是，现行主流硬件（例如多核CPU和GPU）并不能真正把稀疏性用到极致。要把稀疏优势发挥出来，研究人员和工程师必须重新审视并重构计算栈的各层设计，包括硬件、底层固件以及应用软件。

在斯坦福大学的研究团队中，我们打造了（据我们所知）首套能够高效处理各种稀疏与传统工作负载的硬件平台。不同任务带来的节能幅度并不相同，但平均来看，这款芯片能耗只有CPU的1/70，计算速度则提升约8倍。为了达成这一目标，我们需要从硬件、底层固件到软件整体“重来”，让稀疏性得到充分利用。我们也希望这只是开始，未来还能实现更节能的人工智能。

什么是稀疏性？

神经网络与输入数据通常以数字数组的形式出现。这些数组可能是一维的（向量）、二维的（矩阵），或者更高维的（张量）。稀疏的向量、矩阵或张量主要由零元素构成，稀疏程度因对象而异。当任意维度的数组里零元素占比超过50%时，就可以采用面向稀疏性的计算方式。反过来，如果一个对象几乎不稀疏——也就是与元素总量相比，零值很少——就称为稠密对象。

稀疏性既可能自然形成，也可能被有意引入。比如社交网络图本身就常常是稀疏的：设想一个图，其中每个节点代表一个人，每条边表示一段友谊，由于大多数人并非彼此都认识，把所有潜在连接都画成矩阵时，大部分位置都会是零。类似地，其他AI应用如图学习与推荐模型，也常内建天然稀疏结构。

除了天然稀疏之外，还可以通过多种手段在人为设计中让人工智能模型呈现稀疏性。两年前，Cerebras的一支团队证明：在LLM模型中，可以将高达70%到80%的参数设为零而不损失精度。Cerebras还在Meta开源的Llama 7B模型上验证了这一结论；但这些思路也同样适用于其他LLM，例如ChatGPT和Claude。

稀疏性论证

稀疏计算之所以高效，主要源自两个基础特性：一是对零值进行压缩的能力，二是零本身的数学性质。用于稀疏计算的算法以及配套专用硬件，都围绕这两点展开。

首先，稀疏数据能够被压缩，从而更高效地使用内存，具体做法是用稀疏数据类型按需存储——只保存必要内容。压缩也能降低处理大规模数据时的数据传输能耗。举例来说，假设有一个4×4矩阵，其中只有3个非零元素。若不压缩，矩阵要占用16个存储单元；但如果使用稀疏数据类型，把零元素去掉，仅保留非零部分，那么在同样的示例中，压缩后的矩阵只需要13个存储单元，而原本的稠密版本仍要16个。随着稀疏度与矩阵规模增加，内存节省会进一步扩大。

压缩不仅涉及值本身，还需要元数据（metadata）。要知道每个非零元素在行和列的具体位置，行列坐标必须被记录下来。通常，这会被实现为“fibertree”：先列出包含非零元素的行标签，再把行标签链接到对应的列标签，随后列标签再指向存储这些位置上数值的部分。

在内存层面，事情会变得更复杂：除了存储每个非零值对应的行与列标签，还要额外记录“段”（segment）的信息，用来标明预计有多少个相关标签，从而把元数据与真实数据清晰地区分。

对于稠密且不压缩的矩阵数据类型，可以按顺序或并行访问数值，并且能用直接的简单公式推算位置。但对稀疏且压缩的数据而言，读取某个值往往要先定位行索引，再基于间接信息去找到列索引，最后才得到目标数值。由于稀疏数据的真实存储位置可能呈现较大随机性，这类间接查找结果与数据相关，并且需要动态发起内存查找操作。

其次，零值的两个数学特征让软件与硬件都能省下大量计算：任何数乘以零都等于零，因此不需要真的执行乘法；任何数加上零仍是该数本身，因此也无需进行加法运算。

在矩阵向量乘法中，这是一项AI工作负载里极常见的操作。除非同时涉及两个非零元素的情形，其他计算都可以跳过。举例来说，还是前面提到的4×4矩阵与一个包含四个数字的向量：在稠密计算里，向量的每个元素都要与矩阵中对应行的元素相乘，然后把所有乘积累加，才能形成最终向量。此时需要进行16次乘法与16次加法（或可理解为4次累加）。

在稀疏计算里，只处理向量中的非零元素。对每个非零向量元素，先通过间接查找找到对应的非零矩阵元素，再只对这些位置执行乘法与加法。在这个例子中，总共只需完成两次乘法，而不是通常的16次。

GPU和CPU的问题

遗憾的是，现代硬件并不总能有效加速稀疏计算。比如进行矩阵向量乘法时，最直接的做法是在单个CPU核心上依次完成：读取向量元素，逐一相乘并把结果写入内存。由于一次只能完成一次乘法，这种方式速度会很慢。于是人们通常会使用支持向量运算的CPU或GPU，让大量元素并行相乘，从而提升运行速度。但如果矩阵与向量都极其稀疏，向量化硬件仍会把大部分资源用于处理与零相关的乘法，产生大量无效计算。

新一代GPU确实能利用硬件稀疏性，但通常只适用于某一类方法：结构化稀疏性。它要求在每四个相邻参数中至少有两个为零。然而不少模型更受益于非结构化稀疏性，也就是任意权重或激活值都可为零并被压缩，且不强制其位置与相邻关系。GPU也可以通过cuSparse等库在软件层面运行非结构化稀疏计算，不过对稀疏计算的支持往往有限，导致GPU利用率不高，许多本该节省的高能耗计算仍被开销消耗掉。

当在软件中实现稀疏计算时，现代CPU可能反而更合适，因为其架构更灵活。但CPU上的稀疏运算经常会被间接查找非零数据的流程拖慢。CPU设计通常面向内存预取（prefetch），但对随机分布的稀疏数据而言，预取往往拿不到正确内容，于是CPU不得不消耗更多周期去等待正确数据。

苹果公司率先让A14与M1芯片中的预取器支持一种名为指针数组访问模式的技术，从而加速间接查找。虽然预取优化能让苹果CPU在稀疏计算上更有竞争力，但由于CPU架构仍要承担通用计算任务，它依旧存在专门稀疏计算架构所没有的基础开销。

其他公司也在尝试用硬件加速稀疏机器学习，例如Cerebras的晶圆级引擎（Wafer Scale Engine）以及Meta的训练与推理加速器（MTIA）。晶圆级引擎及其稀疏编程框架在LLM上展现出极高稀疏度，最高可达70%。但该方案的硬件与软件主要覆盖权重稀疏，尚不支持激活稀疏，而激活稀疏对许多场景同样关键。MTIA v2则宣称相比MTIA v1，稀疏计算性能提升了七倍，不过目前公开的关于MTIA v2稀疏支持内容仍主要停留在矩阵乘法层面，而非向量或张量。

尽管矩阵乘法在多数机器学习模型里占据主要计算时间，但对其他环节同样需要稀疏支持。为了避免在稀疏数据类型与稠密数据类型之间频繁切换，理想做法是让各类操作都保持稀疏友好。

Onyx

斯坦福团队没有选择这些权衡式路径，而是直接开发一款名为Onyx的硬件加速器。Onyx从设计之初就能同时发挥结构化与非结构化稀疏性的作用。它也是首款支持稀疏计算与稠密计算并存的可编程加速器，能把两类场景中的关键运算都提速。

要理解Onyx，需要先弄清什么是粗粒度可重构阵列（CGRA：coarse-grained reconfigurable array），以及它和更常见的CPU或现场可编程门阵列（FPGA）有何不同。

CPU、CGRA与FPGA之间，本质上是在效率与灵活性之间做取舍。CPU的每个逻辑单元都面向特定任务设计，因此执行对应功能时效率更高。相对地，FPGA因为每个比特都可配置，灵活性很强，但效率往往偏低。CGRA试图在尽量保持CPU效率的前提下，获得接近FPGA的灵活度。

CGRA由一组高效且可配置的单元构成，通常包括内存单元与计算单元，并且这些单元会针对特定应用领域进行优化。其关键优势在于：程序员可以在较高层级重新配置CGRA的内部连接与结构，从而比FPGA更高效，同时又比CPU更灵活。

Onyx包含灵活可编程的处理单元（PE：processing element）以及存储单元（MEM）。存储单元负责保存压缩矩阵与其他所需数据格式，而处理单元则直接在压缩矩阵上完成运算，从源头消除那些无效且低效率的计算。

Onyx编译器会把软件指令转换为CGRA配置。首先，把输入表达式（例如稀疏向量乘法）映射成抽象的“内存节点”和“计算节点”组成的图。以该例为例，内存节点用来承载输入向量与输出向量，计算节点则用来完成非零元素交集的定位与乘法运算。随后，编译器决定这些抽象节点如何映射到CGRA上的MEM与PE，并建立节点之间的数据传输连接。最后，编译器生成满足目标功能的CGRA配置指令集。

由于Onyx支持可编程，工程师可以把大量不同运算（如向量-向量元素乘）或AI中的关键任务（如矩阵-向量、矩阵-矩阵乘法）映射到加速器上。

我们用“能耗×计算时长”的乘积指标来衡量硬件效率，也就是能量延迟积（EDP）。该指标直接体现了速度与能耗之间的权衡：如果只追求最低能耗，设备往往会变得很慢；如果只追求最高速度，又可能导致芯片面积变大、功耗偏高。

相较于使用专门稀疏库的CPU（我们使用的是12核Intel Xeon CPU），Onyx的能量延迟积最高可提升565倍。Onyx也可以像GPU或TPU那样配置，用于加速常规密集型程序：当计算属于稀疏场景时，它会采用稀疏原语；当计算偏向密集时，它会重新配置以更好利用并行性。这样的架构迈出了在同一芯片上同时加速稀疏与密集计算的一步。

同样重要的是，Onyx还有望催生新的算法思维。稀疏加速硬件不仅能提高AI的性能与能效，也能帮助研究人员与工程师从更高视角探索潜力巨大的新算法。

稀疏性的未来

斯坦福团队已经在研发基于Onyx的下一代芯片。除了矩阵乘法，机器学习模型还会做许多其他运算，例如非线性层、归一化、softmax等。团队正在把对这些计算的支持加入下一代加速器与编译器中。考虑到稀疏模型可能同时存在稀疏层和密集层，团队也在努力把密集与稀疏加速器架构更高效地融合到芯片里，从而减少不同数据类型之间的切换成本。此外，他们也在探索如何用更高效的方式分解稀疏数据，以更好管理内存限制并支持在多块稀疏加速器芯片上运行计算。

团队还在研发能够预测这类加速器性能的系统，这将帮助设计出更适合稀疏AI的硬件。展望未来，我们关注的是：高度稀疏的特性是否会在更多模型类型中被采用，以及稀疏加速器能否得到更广泛的规模化应用。

把硬件做成能处理非结构化稀疏数据，并真正把零值优势用起来，只是起点。有了这套能力，AI研究人员和工程师就能用更新颖、也更有创造力的方式构建利用稀疏性的模型与算法。我们相信，这对于应对人工智能不断增长的运行时间、成本以及环境影响至关重要。

END

欢迎留言区写评交流！

行业社群

注：本文题图来自freepik、作者自制、AI 生成（微信公众平台AI配图、豆包）、媒体公开资料、皆已授权。