放大招了｜十亿参数大模型LLMs运行功耗仅需13W，内存使用量减少90%！

news2024/10/6 12:31:49

矩阵乘法（MatMul）历来是大型语言模型（LLMs）总体计算成本的主导因素，尤其在模型向更大维度嵌入和上下文长度发展时，这一成本呈指数级增长。

近期有一篇刚刚发表的论文中提出的方法完全去除了矩阵乘法操作，同时在十亿参数级别的模型中实现了与前沿Transformer模型相当的性能。研究发现，随着模型规模的增加，无矩阵乘法模型与全精度Transformer模型之间的性能差距逐渐减小。在GPU实现上，通过优化内核，训练时的内存消耗相比未优化版本减少了高达61%，而在推理阶段，模型的内存使用量更是减少了90%以上。此外，作者们还在FPGA上构建了定制硬件，利用了比GPU更轻量级的操作，实现了在13W功耗下处理十亿参数规模模型的能力。

大型语言模型，如ChatGPT，虽然能产生令人印象深刻的智能结果，但其高昂的运行成本和巨大的碳足迹也让人望而却步。据估算，运行ChatGPT 3.5每天的能源成本高达70万美元。然而，UC Santa Cruz的研究团队展示了如何在不牺牲性能的前提下，剔除运行大型语言模型过程中最昂贵的计算环节。

矩阵乘法是大多数神经网络中占据主导地位的操作，密集层涉及向量-矩阵乘法（VMM），卷积可以作为具有共享权重的块稀疏VMM实现，而自注意力则依赖于矩阵-矩阵乘法（MMM）。由于图形处理器（GPU）对MatMul操作进行了优化，MatMul的普遍性得以凸显。通过CUDA和高度优化的线性代数库如cuBLAS，MatMul操作可以被高效并行化和加速。这种优化是AlexNet在ILSVRC2012竞赛中获胜的关键，也是深度学习崛起的历史标志。AlexNet利用GPU将训练速度提升至CPU无法企及的程度，因此，无论是训练还是推理，都离不开MatMul操作。

为了进行矩阵相乘，需要将矩阵存储在物理分离的多个图形处理单元（GPU）上，这些GPU是专为快速处理大规模数据集设计的电路，主要由NVIDIA制造。当需要从不同GPU上的矩阵中提取数字进行相乘时，数据的移动成为神经网络时间及能耗成本的主要来源。

尽管在深度学习中普遍存在，但MatMul操作占据了计算开销的主要部分，通常消耗了执行时间和内存访问的大部分，无论是在训练阶段还是推理阶段。一些工作通过两种主要策略替代了MatMul，以更简单的操作取而代之。第一种策略涉及用基本操作替换MatMul，例如AdderNet在卷积神经网络（CNNs）中用有符号加法代替了乘法。考虑到对卷积的关注，AdderNet更适合用于计算机视觉而不是语言建模。第二种策略是通过近似方法减少MatMul的计算复杂度，如低秩分解和哈希技巧。这些策略通常会牺牲一定程度的精度以换取计算效率。

从上图对比发现，在370M参数规模下，MatMul-free Transformer++模型未能收敛，而本论文提出的方法在没有矩阵乘法（MatMul-free）的设置下成功收敛。这意味着，尽管MatMul-free Transformer++尝试在不使用矩阵乘法的情况下训练，但它未能有效地最小化损失函数，从而导致训练失败。这可能是因为它在去除矩阵乘法后，未能找到有效的替代机制来处理复杂的计算需求，特别是当涉及到大量的权重更新和梯度传播时。

相比之下，我们提出的方法在相同的MatMul-free条件下实现了成功的收敛。这表明我们的方法有效地解决了在没有矩阵乘法的情况下进行训练的问题，通过创新的算法设计或优化策略，确保了即使在去除矩阵乘法后，模型依然能够学习并改进其预测能力。

研究人员采用两种主要技术避免使用矩阵乘法。首先，他们将所有矩阵内的数字限制为三进制，即只能取-1、0或1三个值，这使得计算过程简化为加法而非乘法。其次，研究者调整了矩阵之间的通信策略，不是将一个矩阵中的每个数字与另一个矩阵中的每个数字相乘，而是设计了一种只执行最重要运算的策略，通过叠加矩阵并仅计算最关键的部分，达到与矩阵乘法相同的效果，但大大减少了计算成本。

尽管计算量减少，但通过引入基于时间的计算训练模型，研究者能够保持神经网络的性能，使网络能够“记住”处理的关键信息，从而提升表现。实验显示，他们的模型在数十亿参数规模下与Meta的顶级算法Llama取得了相同的性能。

研究团队最初设计的神经网络可在GPU上运行，以利用GPU在人工智能行业的普及性，让软件易于获取和使用。在标准GPU上，该神经网络的内存消耗减少了约10倍，运行速度提高了约25%，这可能为在内存较小的设备（如智能手机）上全速运行强大的大型语言模型铺平道路。

NVIDIA设计的GPU高度优化了矩阵乘法的性能，这使他们在行业中独占鳌头，成为全球最盈利的公司之一。然而，这些硬件并未针对三进制运算进行全面优化。为了进一步节省能源，研究团队与计算机科学与工程系助理教授Dustin Richmond以及讲师Ethan Sifferman合作，创建了定制硬件。在三周内，他们在一个高度可定制的电路——现场可编程门阵列（FPGA）上构建了硬件原型，充分利用了神经网络中编程的所有节能特性。

借助定制硬件，模型在仅13瓦特功率下实现了超越人类阅读速率的输出，相比之下，使用GPU大约需要700瓦特的功率，这意味着定制硬件的效率是GPU的50多倍。研究人员认为，随着进一步的发展，他们可以进一步优化技术，实现更高的能源效率。他们指出，当前的结果已经相当出色，但在更有效地利用资源方面仍有巨大的潜力可挖。如果能在13瓦特功率下做到这一点，那么想象一下在整座数据中心的计算能力加持下，我们能取得怎样的成就。

论文中提到了GPU和FPGA上的轻量化模型实现。通过融合内核，GPU实现加速了训练过程，内存消耗减少了61%。此外，使用低比特优化的CUDA内核，推理速度提高了4.57倍，当模型扩展到130亿参数时，内存使用量减少90%。FPGA上的定制硬件解决方案展示了在真实世界中，轻量化模型不仅能降低计算需求，还能减少能源使用。

总的来说，这篇论文展示了在大型语言模型中完全移除矩阵乘法的可能性，同时保持了模型性能，大幅降低了计算和内存消耗。此外，通过GPU和FPGA上的优化实现，展示了轻量化模型在实际应用中的潜力，特别是其能源效率接近脑部水平。这些成果为未来的模型设计和硬件优化提供了新的方向。

如果您看完有所受益，欢迎点击文章底部左下角“关注”并点击“分享”、“在看”，非常感谢！

精彩推荐: