从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

让谈话模型「轻装上阵」。一直以来,矩阵乘法(MatMul)稳居神经网络操纵的主导地位,其中很大原因归结为 GPU 专门针对 MatMul 操纵进行了优化。这种优化使得 AlexNet 在 ILSVRC2012 挑战赛中一举胜出,成为深度进修崛起的历史性标志。在这当中,有个值得注意的点是,AlexNet 利用 GPU 来提高训练速度,超越了 CPU 的能力,至此,GPU 的加入使得深度进修仿佛赢得了「硬件彩票」。尽管 MatMul 在深度进修中很流行,但不得不承认的是它占据了计算开销的主要部分,主要表现为 MatMu

让谈话模型「轻装上阵」。

一直以来,矩阵乘法(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 操纵? 

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

在这项工作中,来自加州大学圣克鲁兹分校等机构的钻研者证明了 MatMul 操纵可以完全从 LLM 中消除,同时在十亿参数尺度下保持强大的功能。

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

论文地址: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 更接近大脑般的效率。

网友看后不禁感叹道:看来有大事要发生了。

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

不过,受到计算资源的限制,钻研者还没有在非常大的模型(如参数大于100B的模型)上测试MatMul-free模型的有效性,因此其实战效果还有待观察。

方法介绍

该钻研构建了首个可扩张的MatMul-free谈话模型 (Matmul-free LM),通过在麋集层中利用加法操纵以及在自注意力类函数中利用元素级 Hadamard 乘积完成。 

具体来说,三值权重消除了麋集层中的 MatMul,类似于 BNN。为了从自注意力中移除 MatMul,该钻研优化了门控循环单元 (GRU) ,使其仅依赖于元素级乘积,并表明该模型可与最先进的 Transformer 相媲美,同时消除了所有 MatMul 运算。

论文详细介绍 MatMul-free 谈话模型(LM)的各个组成部分。

具有三值权重的 MatMul-free 麋集层

在标准麋集层中,输入和权重矩阵之间的 MatMul 可以表示为:

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

为了避免利用基于 MatMul 的麋集层,该钻研采用 BitNet 来替换包含 MatMul 的麋集层,即利用 BitLinear 模块,这些模块利用三值权重将 MatMul 操纵转换为纯加法操纵。当利用三值权重时,权重矩阵 W 中的元素被限制在集合 {-1, 0, +1} 中。带有三值权重的 MatMul 可以表达为: 

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

由于三值化权重从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗只能从 {−1, 0, +1} 中取值,因而 MatMul 中的乘法运算可以用简单的加法或减法运算代替:

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

因此,三值化 MatMul 可以写成如下:

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

算法 1 为量化融合 RMSNorm 和 BitLinear 算法的流程图:

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

MatMul-free 谈话模型架构

钻研职员采用了 Metaformer 的观点,该观点认为 Transformer 由两部分组成:token mixer(用于混合时间信息,即自注意力机制)和 channel mixer(用于混合嵌入 / 空间信息,即前馈网络,门控线性单元 GLU )。该架构的高级概览如图 2 所示。

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

自注意力机制是现代谈话模型中最常用的 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 版本可以表达如下:

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

这里的通道混合器仅由麋集层组成,这些层已被三值化累积操纵所替代。通过在 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++。

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

有意思的是,与 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 训练设置中通常报告的值要大。

有趣的是,钻研职员还观察到,与利用较小进修率训练的模型相比,训练初期利用较大进修率训练的模型,在训练后期的阶段训练损失下降得更快。

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

钻研职员根据训练时间和内存利用情况评估了他们提出的融合型 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 所示。

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

表 2 显示了单块(N = 1)单核实现的资源利用、功耗和功能。

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

通过利用完整的 512 位 DDR4 接口,并对占据核心处理时间 99% 的 TMATMUL 功能单元进行并行化,预计在保持相同时钟速率的情况下,无需额外优化或流水线化,可以实现大约 64 倍的进一步加速,如表 3 所示。

从LLM中完全消除矩阵乘法,效果出奇得好,10亿参数跑在FPGA上接近大脑功耗

给TA打赏
共{{data.count}}人
人已打赏
工程

轻松建立聊天机器人、准确性新SOTA,RAG有了更强大的AI检索器

2024-6-9 0:42:00

工程

i人小助手:Meta推出多模态对话图,帮你轻松鉴别交际状态

2024-6-11 0:49:00

0 条回复 A文章作者 M管理员
    暂无讨论,说说你的看法吧
个人中心
今日签到
搜索