让语言模型「轻装上阵」。
一直以来,矩阵乘法(MatMul)稳居神经网络操作的主导地位,其中很大原因归结为 GPU 专门针对 MatMul 操作进行了优化。这种优化使得 AlexNet 在 ILSVRC2012 挑战赛中一举胜出,成为深度学习崛起的历史性标志。
在这当中,有个值得注意的点是,AlexNet 利用 GPU 来提高训练速度,超越了 CPU 的能力,至此,GPU 的加入使得深度学习仿佛赢得了「硬件彩票」。
尽管 MatMul 在深度学习中很流行,但不得不承认的是它占据了计算开销的主要部分,主要表现为 MatMul 在训练和推理阶段消耗大部分执行时间和内存访问。
目前为止,研究者已经开始借助其他更简单的操作替代 MatMul,主要有两种。
第一种策略是使用初等运算代替 MatMul,例如,在卷积神经网络 (CNN) 中,用有符号加法代替乘法;
第二种方法是使用二值或三值化量化,将 MatMul 值在累加之前要么翻转要么清零。比如脉冲神经网络 (SNN) 使用二值激活,而二值化网络 BNN 使用量化权重。
在语言建模方面,BitNet 等技术的出现表明量化的可扩展性, 但这种方式仍然保留了昂贵的矩阵 - 矩阵相乘(MMM)的自注意力机制。研究者尝试过多种努力,但 MatMul 操作在 GPU 上仍然是资源密集型的。
既然 MatMul 占据了 LLM 整体计算成本,且随着 LLM 向更大的嵌入维度和上下文长度扩展时,这种成本只会增加。这引发了一个问题:是否有可能完全从 LLM 中消除 MatMul 操作?
在这项工作中,来自加州大学圣克鲁兹分校等机构的研究者证明了 MatMul 操作可以完全从 LLM 中消除,同时在十亿参数尺度下保持强大的性能。
论文地址:https://arxiv.org/pdf/2406.02528
项目地址:https://github.com/ridgerchu/matmulfreellm
论文标题:Scalable MatMul-free Language Modeling
实验表明,该研究提出的 MatMul-free 模型达到了与最先进的 Transformer 相当的性能,后者在推理期间需要更多的内存,规模至少为 2.7B 参数。
此外,论文还研究了扩展定律,发现随着模型规模的增加,MatMul-free 模型与全精度 Transformer 之间的性能差距逐渐缩小。
研究者还提供了一种高效的 GPU 模型实现方式,在训练期间相比未优化的基线模型减少了多达 61% 的内存使用。通过在推理时利用优化的内核,模型内存消耗可以比未优化的模型减少超过 10 倍。
最后,研究者在 FPGA 上构建了一个自定义硬件解决方案,他们以 13W 的功耗处理了十亿参数规模的模型,超出了人类可读的吞吐量,使 LLM 更接近大脑般的效率。
网友看后不禁感叹道:看来有大事要发生了。
不过,受到计算资源的限制,研究者还没有在非常大的模型(如参数大于100B的模型)上测试MatMul-free模型的有效性,因此其实战效果还有待观察。
方法介绍
该研究构建了首个可扩展的MatMul-free语言模型 (Matmul-free LM),通过在密集层中使用加法操作以及在自注意力类函数中使用元素级 Hadamard 乘积完成。
具体来说,三值权重消除了密集层中的 MatMul,类似于 BNN。为了从自注意力中移除 MatMul,该研究优化了门控循环单元 (GRU) ,使其仅依赖于元素级乘积,并表明该模型可与最先进的 Transformer 相媲美,同时消除了所有 MatMul 运算。
论文详细介绍 MatMul-free 语言模型(LM)的各个组成部分。
具有三值权重的 MatMul-free 密集层
在标准密集层中,输入和权重矩阵之间的 MatMul 可以表示为:
为了避免使用基于 MatMul 的密集层,该研究采用 BitNet 来替换包含 MatMul 的密集层,即使用 BitLinear 模块,这些模块使用三值权重将 MatMul 操作转换为纯加法操作。当使用三值权重时,权重矩阵 W 中的元素被限制在集合 {-1, 0, +1} 中。带有三值权重的 MatMul 可以表达为:
由于三值化权重只能从 {−1, 0, +1} 中取值,因而 MatMul 中的乘法运算可以用简单的加法或减法运算代替:
因此,三值化 MatMul 可以写成如下:
算法 1 为量化融合 RMSNorm 和 BitLinear 算法的流程图:MatMul-free 语言模型架构
研究人员采用了 Metaformer 的观点,该观点认为 Transformer 由两部分组成:token mixer(用于混合时间信息,即自注意力机制)和 channel mixer(用于混合嵌入 / 空间信息,即前馈网络,门控线性单元 GLU )。该架构的高级概览如图 2 所示。
自注意力机制是现代语言模型中最常用的 token mixer,它依赖于三个矩阵 Q、K 和 V 之间的矩阵乘法。为了将这些操作转换为加法,研究人员至少对两个矩阵进行二值化或三值化处理。
假设所有密集层的权重都是三值的,他们将 Q 和 K 量化,得到一个三值的注意力图,从而消除自注意力中的乘法操作。但是,以这种方式训练的模型无法成功收敛。一个可能的解释是,激活值包含对性能至关重要但难以有效量化的异常值。
为了解决这一挑战,研究人员探索了不依赖于矩阵乘法的替代方法来混合 token。
通过采用结合了逐元素操作和累积的三值 RNN,可以构建一个MatMul-free 的 token mixer。在各种 RNN 架构中,GRU 因其简单高效而著称,它在比长短期记忆网络(LSTM)使用更少的门控单元和结构更简单的情况下,实现了类似的性能。因此,研究人员选择 GRU 作为构建 MatMul-free token mixer 的基础。
为了实现 MatMul-free 的通道混合,研究人员采用了门控线性单元(GLU),它在许多现代 LLM 中得到了广泛应用,包括 Llama 、Mistral 和 RWKV。一个适应了 BitLinear 的 GLU 版本可以表达如下:
这里的通道混合器仅由密集层组成,这些层已被三值化累积操作所替代。通过在 BitLinear 模块中使用三值权重,研究人员可以消除对昂贵 MatMul 的需求,这样不仅使通道混合器在计算上更加高效,同时还保持了其在跨通道混合信息方面的有效性。
实验
该研究的重点是在中等规模的语言建模任务上测试 MatMul-free 的语言模型。研究人员将两种变体的 MatMul-free 语言模型与复现的高级 Transformer 架构(基于 Llama-2 的 Transformer++)在三个模型大小上进行比较:370M、13 亿和 27 亿参数。
为了公平比较,所有模型都在 SlimPajama 数据集上进行了预训练,其中 370M 模型训练了 150 亿个 token,而 13 亿和 27 亿模型各训练了 1000 亿个 token。
模型训练使用了 8 个 NVIDIA H100 GPU。370M 模型的训练时间大约为 5 小时,13 亿模型为 84 小时,27 亿模型为 173 小时。
MatMul-free 语言模型的扩展规律
研究团队评估了MatMul-free 语言模型和 Transformer++ 在 370M、13 亿和 27 亿参数模型上的扩展规律,如图 3 所示。
为便于比较,MatMul-free LM 和 Transformer++ 中的每个操作都被同等对待。但请注意,Transformer++ 中的所有权重和激活都是 BF16 格式,而MatMul-free 语言模型中的 BitLinear 层使用三元参数,激活为 BF16。因此,MatMul-free 语言模型的平均运算成本要低于 Transformer++。
有意思的是,与 Transformer++ 相比,MatMul-free 语言模型的扩展投影显示出更陡峭的下降趋势,这表明MatMul-free语言模型在利用额外计算资源提高性能方面更为高效。
因此,MatMul-free 语言模型的扩展曲线预计将在大约 10^23 次浮点运算(FLOPs)处与 Transformer++ 的扩展曲线相交。这个计算规模大致相当于训练 Llama-3 8B(使用 1.5 万亿个 token 训练)和 Llama-2 70B(使用 2 万亿个 token 训练)所需的 FLOPs,这表明MatMul-free 语言模型不仅在效率上胜出,而且扩展时在损失方面也可能表现更好。
下游任务
学习率是语言模型训练中一个关键的超参数,当模型处于三元 / 二元权重状态时,对学习率变得更加敏感。为了确定最优学习率,研究人员使用 370M 模型,批量大小为 50k 个 token,在 1.5e−3 到 3e−2 的范围内进行了搜索。这次搜索的结果如图 4 (c) 所示。
结果显示,当学习率从 1.5e−3 增加到 1e−2 时,最终训练损失单调递减。只有当学习率超过 2e−2 时,模型才表现出不稳定。这一发现表明,以前使用三元权重的作品,如使用 1.5e−3 学习率的 BitNet,可能不是最优的,更高的学习率有可能带来更好的性能。
这些发现与 Deepseek LLM 的观察结果一致,后者发现传统大型语言模型(LLMs)的最佳学习率实际上比大多数 LLM 训练设置中通常报告的值要大。
有趣的是,研究人员还观察到,与使用较小学习率训练的模型相比,训练初期使用较大学习率训练的模型,在训练后期的阶段训练损失下降得更快。
研究人员根据训练时间和内存使用情况评估了他们提出的融合型 BitLinear 和传统型 BitLinear 实现,如图 4 (a-b) 所示。
实验表明,他们的融合操作器在更大的批量大小下,能够带来更快的训练速度,并减少内存消耗。当批量大小为 2 的 8 次方时,1.3B 参数模型的训练速度从每次迭代 1.52 秒提高到 1.21 秒,比 Vanilla 实现快了 25.6%。
此外,内存消耗从 82GB 减少到 32GB,内存使用减少了 61.0%。随着批量大小的增加,融合实现的性能显著提高,允许同时处理更多的样本,并减少总迭代次数。
图 4 (d) 展示了不同模型大小下,所提出的MatMul-free语言模型与 Transformer++ 在 GPU 推理内存消耗和延迟方面的比较。
在MatMul-free语言模型中,研究人员采用 BitBLAS 进行加速,以进一步提高效率。评估是在批量大小为 1,序列长度为 2048 的情况下进行的。
MatMul-free语言模型在所有模型大小上都显示出比 Transformer++ 更低的内存使用和延迟。
对于单层,MatMul-free 语言模型只需要 0.12 GB 的 GPU 内存,并且实现了 3.79 毫秒的延迟,而 Transformer++ 消耗了 0.21 GB 的内存,并且有 13.87 毫秒的延迟。随着模型大小的增加,MatMul-free 语言模型的内存和延迟优势变得更加明显。
值得注意的是,对于大于 2.7B 的模型大小,结果是使用随机初始化的权重进行模拟的。对于最大的 13B 参数模型,MatMul-free 语言模型仅使用 4.19 GB 的 GPU 内存,并且有 695.48 毫秒的延迟,而 Transformer++ 需要 48.50 GB 的内存,并表现出 3183.10 毫秒的延迟。
这些结果突出了MatMul-free语言模型所实现的效率增益,使其成为大规模语言建模任务的有希望的方法,特别是在推理期间。
为了测试MatMul-free语言模型功耗和有效性,研究人员使用 SystemVerilog 创建了一个 FPGA 加速器。概览如图 5 所示。
表 2 显示了单块(N = 1)单核实现的资源利用、功耗和性能。
通过使用完整的 512 位 DDR4 接口,并对占据核心处理时间 99% 的 TMATMUL 功能单元进行并行化,预计在保持相同时钟速率的情况下,无需额外优化或流水线化,可以实现大约 64 倍的进一步加速,如表 3 所示。
注意力研究者gpu 新浪科技公众号“掌”握科技鲜闻 (微信搜索techsina或扫描左侧二维码关注)
相关新闻