從LLM中完全消除矩陣乘法，效果出奇得好，10億引數跑在FPGA上接近大腦功耗
机器之心發表於2024-06-11
一直以來，矩陣乘法（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 可以表示為：為了避免使用基於 MatMul 的密集層，該研究採用 BitNet 來替換包含 MatMul 的密集層，即使用 BitLinear 模組，這些模組使用三值權重將 MatMul 操作轉換為純加法操作。當使用三值權重時，權重矩陣 W 中的元素被限制在集合 {-1, 0, +1} 中。帶有三值權重的 MatMul 可以表達為： 由於三值化權重
只能從 {−1, 0, +1} 中取值，因而 MatMul 中的乘法運算可以用簡單的加法或減法運算代替：演算法 1 為量化融合 RMSNorm 和 BitLinear 演算法的流程圖：研究人員採用了 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 語言模型和 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 所示。