CPU推理提升4到5倍,蘋果用閃存加速大模型推理,Siri2.0要來了?
近年來,GPT-3、OPT 和 PaLM 等大型語言模型(LLM)在廣泛的 NLP 任務中表現出了強大的性能。不過,這些能力伴隨著大量計算和內存推理需求,畢竟大型語言模型可能包含數千億甚至萬億參數,使得高效加載和運行變得有挑戰性,尤其是在資源有限的設備上。
當前標準的應對方案是將整個模型加載到 DRAM 中進行推理,然而這種做法嚴重限制了可以運行的最大模型尺寸。舉個例子,70 億參數的模型需要 14GB 以上的內存才能加載半精度浮點格式的參數,這超出了大多數邊緣設備的能力。
為了解決這種局限性,蘋果的研究者提出在閃存中存儲模型參數,至少比 DRAM 大了一個數量級。接著在推理中,他們直接并巧妙地從閃存加載所需參數,不再需要將整個模型擬合到 DRAM 中。
這種方法基于最近的工作構建,這些工作表明 LLM 在前饋網絡(FFN)層中表現出高度稀疏性,其中 OPT、Falcon 等模型的稀疏性更是超過 90%。因此,研究者利用這種稀疏性, 有選擇地僅從閃存中加載具有非零輸入或預測具有非零輸出的參數。
具體來講,研究者討論了一種受硬件啟發的成本模型,其中包括閃存、DRAM 和計算核心(CPU 或 GPU)。接著引入兩種互補技術來最小化數據傳輸、最大化閃存吞吐量:
窗口:僅為前幾個 token 加載參數,并重用最近計算的 tokens 的激活。這種滑動窗口方法減少了加載權重的 IO 請求數量;
行列捆綁:存儲上投影和下投影層的串聯行和列,以讀取閃存的更大連續塊。這將通過讀取更大的塊來增加吞吐量。
為了進一步最小化從閃存傳輸到 DRAM 的權重數量,研究者還設法預測 FFN 稀疏性并避免加載歸零參數。結合使用窗口和稀疏性預測可以為每個推理查詢僅加載 2% 的閃存 FFN 層。他們還提出了靜態內存預分配,最大限度減少了 DRAM 內的傳輸并減少了推理延遲。
本文的閃存加載成本模型在加載更好數據與讀取更大塊之間取得了平衡。與 CPU 和 GPU 中的 naive 實現相比,優化該成本模型并有選擇地按需加載參數的閃存策略可以運行兩倍于 DRAM 容量的模型,并將推理速度分別提升 4-5 倍和 20-25 倍。
有人評價稱,這項工作會讓 iOS 開發更加有趣。
閃存和 LLM 推理
帶寬和能量限制
雖然現代 NAND 閃存提供了高帶寬和低延遲,但仍達不到 DRAM 的性能水準,尤其是在內存受限的系統中。下圖 2a 說明了這些差異。
依賴 NAND 閃存的 naive 推理實現可能需要為每個前向傳遞重新加載整個模型,這一過程非常耗時,即使是壓縮模型也需要幾秒時間。此外將數據從 DRAM 傳輸到 CPU 或 GPU 內存需要耗費更多能量。
在 DRAM 充足的場景中,加載數據的成本有所降低,這時模型可以駐留在 DRAM 中。不過,模型的初始加載仍然耗能,尤其是在第一個 token 需要快速響應時間的情況下。本文的方法利用 LLM 中的激活稀疏性,通過有選擇地讀取模型權重來解決這些挑戰,從而減少了時間和耗能成本。
讀取吞吐量
閃存系統在大量連續讀取場景下表現最佳,例如配備 2TB 閃存的 Apple MacBook Pro M2 基準測試表明,對未緩存文件進行 1GiB 線性讀取的速度超過 6GiB/s。然而由于這些讀取固有的多階段性質,包括操作系統、驅動程序、中端處理器和閃存控制器,較小的隨機讀取無法復制這種高帶寬。每個階段都會出現延遲,從而對較小讀取造成較大的影響。
為了規避這些限制,研究者提倡兩種主要策略,它們可以同時使用。
第一種策略是讀取較大的數據塊。雖然吞吐量的增長不是線性的(較大的數據塊需要較長的傳輸時間),但初始字節的延遲在總請求時間中所占的比例較小,從而提高了數據讀取的效率。圖 2b 描述了這一原理。一個與直覺相反但卻有趣的觀察結果是,在某些情況下,讀取比需要更多的數據(但數據塊較大)然后丟棄,比只讀取需要的部分但數據塊較小更快。
第二種策略是利用存儲堆棧和閃存控制器固有的并行性來實現并行讀取。研究結果表明,在標準硬件上使用多線程 32KiB 或更大的隨機讀取,可以實現適合稀疏 LLM 推理的吞吐量。
最大化吞吐量的關鍵在于權重的存儲方式,因為提高平均塊長度的布局可以顯著提高帶寬。在某些情況下,讀取并隨后丟棄多余的數據,而不是將數據分割成更小的、效率更低的數據塊,可能是有益的。
從閃存加載
受上述挑戰的啟發,研究者提出了優化數據傳輸量和提高讀取吞吐量的方法,以顯著提高推理速度。本節將討論在可用計算內存遠遠小于模型大小的設備上進行推理所面臨的挑戰。
分析該挑戰,需要在閃存中存儲完整的模型權重。研究者評估各種閃存加載策略的主要指標是延遲,延遲分為三個不同部分:從閃存加載的 I/O 成本、管理新加載數據的內存開銷以及推理操作的計算成本。
蘋果將在內存限制條件下減少延遲的解決方案分為三個戰略領域,每個領域都針對延遲的特定方面:
1、減少數據負載:旨在通過加載更少的數據來減少與閃存 I/O 操作相關的延遲。
2、優化數據塊大小:通過增加加載數據塊的大小來提高閃存吞吐量,從而減少延遲。
以下是研究者為提高閃存讀取效率而增加數據塊大小所采用的策略:
捆綁列和行
基于 Co-activation 的捆綁
3、有效管理加載的數據:簡化數據加載到內存后的管理,最大限度地減少開銷。
雖然與訪問閃存相比,在 DRAM 中傳輸數據的效率更高,但會產生不可忽略的成本。在為新神經元引入數據時,由于需要重寫 DRAM 中的現有神經元數據,重新分配矩陣和添加新矩陣可能會導致巨大的開銷。當 DRAM 中的前饋網絡(FFN)有很大一部分(約 25%)需要重寫時,這種代價尤其高昂。
為了解決這個問題,研究者采用了另一種內存管理策略。這包括預分配所有必要的內存,并建立相應的數據結構進行有效管理。如圖 6 所示,該數據結構由指針、矩陣、偏置、已用數和 last_k_active 等元素組成:
圖 6:內存管理,首先將最后一個元素復制到刪除神經元,以保持內存塊的連續性,然后將所需元素堆棧到最后,這樣可以避免多次復制整個數據。
值得注意的是,重點并不在過程的計算方面,因為這與本文工作核心無關。這種劃分使研究者能夠專注于優化閃存交互和內存管理,從而在內存受限的設備上實現高效推理。
實驗結果
OPT 6.7B 模型的結果
預測器。如圖 3a 所示,本文的預測器能準確識別大多數激活的神經元,但偶爾也會誤識數值接近于零的非激活神經元。值得注意的是,這些接近零值的假負類神經元被排除后,并不會明顯改變最終輸出結果。此外,如表 1 所示,這樣的預測準確度水平并不會對模型在零樣本任務中的表現產生不利影響。
延遲分析。當窗口大小為 5 ,每個 token 需要訪問 2.4% 的前饋網絡(FFN)神經元。對于 32 位模型,每次讀取的數據塊大小為 2dmodel × 4 字節 = 32 KiB,因為它涉及行和列的連接。在 M1 Max 上,從閃存加載每個 token 的延遲為 125 毫秒,內存管理(包括神經元的刪除和添加)的延遲為 65 毫秒。因此,與內存相關的總延遲不到每個 token 190 毫秒(參見圖 1)。相比之下,基線方法需要以 6.1GB/s 的速度加載 13.4GB 的數據,導致每個 token 的延遲約為 2330 毫秒。因此,與基線方法相比,本文的方法有了很大改進。
對于 GPU 機器上的 16 位模型,閃存加載時間縮短至 40.5 毫秒,內存管理時間為 40 毫秒,由于從 CPU 向 GPU 傳輸數據的額外開銷,時間略有增加。盡管如此,基線方法的 I/O 時間仍然超過 2000 毫秒。
表 2 提供了每種方法對性能影響的詳細比較。
Falcon 7B 模型的結果
延遲分析。在本文的模型中使用大小為 4 的窗口,每個 token 需要訪問 3.1% 的前饋網絡(FFN)神經元。在 32 位模型中,這相當于每次讀取的數據塊大小為 35.5 KiB(按 2dmodel ×4 字節計算)。在 M1 Max 設備上,從閃存加載這些數據所需的時間約為 161 毫秒,內存管理過程又增加了 90 毫秒,因此每個 token 的總延遲時間為 250 毫秒。相比之下,基線延遲時間約為 2330 毫秒,本文的方法大約快 9 到 10 倍。
- 免責聲明
- 本文所包含的觀點僅代表作者個人看法,不代表新火種的觀點。在新火種上獲取的所有信息均不應被視為投資建議。新火種對本文可能提及或鏈接的任何項目不表示認可。 交易和投資涉及高風險,讀者在采取與本文內容相關的任何行動之前,請務必進行充分的盡職調查。最終的決策應該基于您自己的獨立判斷。新火種不對因依賴本文觀點而產生的任何金錢損失負任何責任。