近日,加州大学发布了一项研究,提出了一种名为MatMul-free LM的新型大语言模型(LLM)架构。不同于ChatGPT等依赖于矩阵乘法的传统LLM,该架构完全不使用矩阵乘法(MatMul),有望摆脱传统LLM对GPU的强烈依赖,为构建高效、硬件友好的LLM开辟了新的可能性。
MatMul-free LM新型大语言模型
自ChatGPT问世以来,大语言模型(LLM)将人工智能推向一个前所未有的高度。然而,训练一个大语言模型可能消耗数百万千瓦时的电力,训练GPT-3就需要消耗相当于500户家庭一年的用电量,大语言模型(LLM)的巨大算力需求和相应的能源消耗问题不容忽视。
传统上,GPU凭借其强大的并行计算能力,成为了训练和部署LLM的首选硬件。然而,GPU的架构设计,尤其是对矩阵乘法(MatMul)的高度依赖,也成为了其能耗过高的原因。
矩阵乘法(MatMul)通常主导大语言模型(LLM)的总体计算成本,矩阵乘法作为LLM中最核心的计算操作,需要进行大量的乘法和累加运算,以及频繁的内存访问,这些都会消耗大量的能量。
随着LLM扩展到更大的嵌入维度和上下文长度,此成本只会增加。
MatMul-free LM通过三项核心架构创新,彻底移除矩阵乘法(MatMul)运算,从根本上降低了能耗:
MatMul-free LM核心原理
三元权重:传统的LLM为了实现复杂的模型表达,通常采用高精度浮点数来表示模型参数。而MatMul-free LM采用了一种更为简洁的方式——三元权重,即将模型参数限制在{-1,0,1}三个值。将复杂的乘法运算简化为高效的加减法,大幅降低计算复杂度和能耗。
MatMul-free Token Mixer
MatMul-free Token Mixer:为了替代需要矩阵乘法的自注意力机制,提出了一种基于门控循环单元(GRU)的MatMul-free Token Mixer(MLGRU)。MLGRU仅依赖逐元素乘积来混合时序信息,从而降低了矩阵乘法运算带来的高昂计算成本。
MatMul-free Channel Mixer:基于门控线性单元(GLU)和三元权重的MatMul-free Channel Mixer,通过门控机制和高效线性变换,进一步减少计算量和能耗。
此外研究人员进行了实验,实验结果证明了MatMul-free LM在降低能耗的同时,性能毫不逊色:
GPU推理效率提升:在GPU推理阶段,MatMul-free LM的内存占用和延迟显著低于传统Transformer模型,尤其在13B参数规模下优势更为明显,内存占用降低10倍,延迟降低4.5倍。
FPGA实现,功耗骤降:在更利于低功耗设计的FPGA平台上,MatMul-free LM能够以极低的功耗(13W)实现超越人类阅读速度的吞吐量,展现出其在边缘设备和低功耗场景中的巨大应用潜力。
性能比肩Transformer:在多个语言建模任务上,MatMul-free LM的性能表现与最先进的Transformer模型相当,甚至在某些任务上更胜一筹,证明了其架构设计的有效性。
当前,AI能耗问题日益凸显,不仅是技术挑战,还是关乎行业可持续发展的热点议题,AI能耗在计算资源、环境影响和成本效益方面都引起了广泛关注。
1.计算资源消耗,硬件效率瓶颈待突破:
训练大语言模型(如GPT-3)需要大量的计算资源,主要依赖高性能的GPU或TPU。这些设备在处理复杂的矩阵运算时,往往会产生极高的能耗。例如,据OpenAI报告,训练GPT-3耗费了超过1000万kWh的电力。
传统硬件如GPU在处理AI计算时,受限于冯·诺依曼架构的数据搬运瓶颈,能效比难以满足需求。
例如,NVIDIA A100 GPU在进行深度神经网络训练时的峰值性能功耗比也仅为30 GFLOPS/W左右。因此,探索更高效的专用架构,如谷歌TPU、Graphcore IPU等,成为了行业热点。
2.巨量碳排放,环境影响引发担忧:
随着模型规模突破trillions of parameters量级,训练能耗呈指数级增长。以GPT-3为例,其训练过程消耗了的电力,相当于排放了超过550吨二氧化碳。
此外,高能耗对数据中心的冷却系统也提出了严峻挑战,进一步增加了能源消耗和运营成本。
3.模型部署成本高昂,数据中心能耗压力剧增:
在AI模型的开发和部署过程中,能源成本已成为不可忽视的因素。以谷歌为例,其数据中心的能源消耗占其运营成本的相当大一部分。
谷歌通过其“碳智能计算计划”,利用机器学习优化数据中心的能源使用时间和方式,最大限度地利用可再生能源,减少碳排放。这一措施在一定程度上缓解了能源成本压力,但其效果依然有限,难以彻底解决能耗问题。
MatMul-free LM的研究,为AI芯片设计带来了新的可能,看到了模型架构摆脱GPU依赖、迈向低功耗未来的可能性。
该研究若能成功应用,将有望推动轻量级模型的发展部署,降低能耗,减少对环境的影响。