本文經自動駕駛之心公眾號授權轉載，轉載請聯系出處。

端側LLM毫無疑問會成為各手機廠商在2024年的主戰場。從國內各手機廠透露的信息來看，大家幾乎都把希望寄托在了芯片廠身上，自身能做的、會做的工作太少。希望蘋果的工作對國內廠商們有啟發、借鑒意義。

論文鏈接：LLM in a flash: Efficient Large Language Model Inference with Limited Memory

1. Flash Memory and DRAM

在移動端設備中（如手機），DRAM可理解為“運行時內存”，Flash Memory可理解為“存儲空間”。做一個簡單的類比，在PC中，DRAM對應于內存；Flash Memory對應于硬盤存儲（注意：僅僅是對應于，實現方案并不一樣）。

在通常的LLM推理階段，LLM都是直接加載到DRAM中的。一個7B半精度LLM，完全加載進DRAM所需的存儲空間超過14GB?？紤]到目前主流手機的DRAM最高也就16GB的水平，在端側直接使用DRAM來加載7B LLM面臨巨大挑戰。

圖1給出了一個移動端標準的存儲結構示意圖。

圖1 移動端存儲結構示意圖

Flash Memory的特點是大存儲，低帶寬。也就是說，Flash Memory可存儲的內容多（圖中的100GB），但數據傳輸速率低（圖中的1GB/s）。而DRAM的特點是小存儲，高帶寬。

現在的問題是：模型大小 > DRAM，所以無法將模型全部加載進DRAM。

蘋果的解決方案是將LLM放在Flash Memory中，在每次需要進行推理時，僅僅將部分必要參數加載到DRAM中。

蘋果的整個方案重點解決兩個問題：

• 如何快速識別出哪些模型參數是必要的
• 考慮到由Flash memory到DRAM的帶寬較低，如何加快由Flash memory到DRAM的傳輸效率
  論文中從三個不同方面做了嘗試，下面分別介紹。

2. 減少數據傳輸量

這部分介紹論文中采用的三種降低數據傳輸量的方法。

2.1 方法一：Selective Persistence Strategy

對于常見的LLM而言，它的模型參數主要由Attention參數和MLP參數兩部分構成，其中Attention參數占比約為1/3，MLP參數占比約為2/3。除此，還有參數量級可忽略不計的Embedding層的參數。

因為Attention參數量相對較少，所以蘋果的方案是將Attention參數和Embedding層的參數直接加載到DRAM中。

這就是所謂的Selective Persistence Strategy，其意為：有選擇性地把部分參數常駐在DRAM中。而這部分常駐的參數就是Attention參數和Embedding參數。原因是因為它們占比較小。

2.2 方法二：Anticipating ReLU Sparsity

這里主要借鑒了DejaVu的思路：MLP層的輸出只有不到10%的值是激活狀態（不為0）。一般把這種現象稱為稀疏性。稀疏性越強，則非激活狀態的值就越多。

所以我們也可把這句話“MLP層的輸出只有不到10%的值是激活狀態”簡寫作“MLP層的稀疏性超過90%”。

要注意，此處的稀疏性一般稱為“Contextual Sparsity”。也就是說，MLP層的哪些神經元會激活，與當前的輸入相關。

蘋果照搬了DejaVu的方法，使用一個兩層MLP來預測哪些神經元會激活。方法也很簡單，假設神經元個數為4096，只需要將MLP的輸出層的大小設為4096，并為每一個輸出使用sigmoid來做一個二分類即可（“選擇”or“不選擇”）。

注意1：不同Transformer層使用的預測模型不同。

注意2：同一個Transformer層中的MLP一般有兩層。它們的激活神經元始終保持相同。

在能夠準確預測的前提下，每次在推理時動態加載預測為激活神經元對應的參數即可。

這里有對DejaVu詳細介紹的文章：[ICML'23] DejaVu：LLM中的動態剪枝

2.3 方法三：Sliding Window

根據2.2小節中介紹的稀疏性可知，在每一次LLM進行前向推理時，它都需要使用模型預測每一個MLP層中激活神經元的編號，并將所需的神經元所對應的權重由Flash memory加載到DRAM中。

因為LLM的推理階段是逐token進行的，這意味著在生成不同token的時候，需要加載到DRAM中的MLP的參數也不同。

用一個簡單的例子來說明這個基礎概念，只考慮第  層的Transformer模塊。在處理當前token 時，該層使用模型預測MLP會激活的神經元編號，假設為{0, 1, 3, 5}，并將其對應的參數從Flash memory加載到DRAM中，然后進行推理。

在處理下一個token 時，將 從DRAM中刪除，再使用模型預測MLP會激活的神經元編號，假設為{0, 2, 3, 6}，并將其對應的參數從Flash memory加載到DRAM中，然后進行推理。

注意到在我們的例子中，兩次前向推理時，第  層的Transformer結構中MLP有部分被預測為激活的神經元是重疊的：{0, 3}。所以實際上在進行第二次前向推理時，沒有必要把完全從DRAM中刪除，而是將其中編號為{1, 5}神經元對應的參數刪除，再將編號為{2, 6}的神經元對應的參數讀入即可。這樣可以減少I/O的總開銷。

這就是Sliding Window的核心思想：保留處理過去k個token時的激活神經元所對應的參數在DRAM中，并在處理當前token時只對：1）部分多余的參數進行刪除；2）缺少的參數進行加載。圖2是原文中的示意圖。

圖2 Sliding Window示意圖

圖中上圖表示在處理當前token “Was”之前，前5個token（k=5）的激活神經元（淡藍色偏綠部分）。圖中下圖表示在處理當前token “Was”之時，需要新加入的神經元（藍色部分）和需要刪除的神經元（分紅部分）。

Sliding Window的核心假設是LLM在處理相鄰token時產生的稀疏性具有相似性。原文沒有仔細分析和論證這個假設。

3 提高傳輸吞吐量

3.1 Bundling Columns and Rows

通常LLM中的MLP層包含兩個全連層。在忽略激活函數的情況下，這兩個全連層可以寫為：

在2.2小節中曾經提到，稀疏性預測是對MLP中兩個全連層同時進行的。也就是說，如果我們預測結果是第一個全連層中0號神經元不會被激活，那么該預測結果同樣也適用于第二個全連層：第二個全連層的0號神經元也不會被激活。

對于第一個全連層的參數矩陣，第i個神經元對應于它的第i列；對于第二個全連層的參數矩陣，第i個神經元對應于它的第i行。

當第i個神經元被預測為激活時，需要同時讀取的第i列和的第i行。所以為了提高讀取速度，可以將的每一列和的對應行拼接起來存儲，如下圖所示：

圖3 將兩個全連層的列與行拼接存儲

圖3中的Up Proj指的是第一個全連層，對應于上文參數矩陣；Down Proj指第二個全連層，對應于上文參數矩陣。

這樣做的好處是原本需要兩次I/O，現在只需要一次了。雖然總的數據讀取量并沒有變，但讀取大塊、連續的數據會比讀取大量小塊、非連續數據更加高效，因此整體傳輸吞吐量提升了。

3.2 Bundling Based on Co-activation

這是一個原文嘗試過，但被驗證為無效的策略。既然原文提到了，所以這里也羅列出來。

原文中猜測某些神經元之間可能存在一些緊密聯系。比如對于兩個神經元a和b，當a激活時，b也會激活（或者當b激活時，a也會激活）。

因此可以通過分析來找到每個神經元的“closest friend”（與該神經元同時激活頻率最高的其它某個神經元）。然后在存儲Flash memory中存儲時，也將它們的參數拼接存在一起。這樣的讀取效率更高。

但該方法之所以無效，主要原因是可能會存在某個神經元i，它是其它很多神經元的“closest friend”。這樣導致的問題則是神經元i被額外傳輸了太多次，導致實際的I/O成本增加了。

4 Optimized Data Management in DRAM

雖然DRAM的數據讀取速度比Flash memory快很多，但當需要對其中數據進行大量、高頻讀寫時，它的時耗仍然不可忽略。在本文介紹的內容中，對DRAM的讀寫主要發生在對MLP層中所需神經元對應參數的刪除與新增（參考圖2）。

為此，論文中設計了一種特殊的數據結構來對DRAM中的數據做精細化管理。該數據結構的核心變量如下：

• Matrix：按照“Bundling Columns and Rows”的方法存儲激活神經元所對應的參數
• bias：激活神經元所對應的bias參數
• num_used：激活神經元的個數
• last_k_active：過去k個token所對應的激活神經元編號
• Pointer：當前行參數對應的神經元編號

圖4 Optimized Data Management in DRAM

通過預分配一個足夠大的空間，可以避免因反復分配而導致的額外開銷。下面來說明基于該數據結構的一些操作的高效實現方法。

該矩陣中的行對應的是當前存儲在DRAM中激活神經元的參數。前文提到（2.3小節），當處理新的token時，需要將不會被激活的神經元刪除，并添加新的會被激活的神經元。所以最重要的兩個操作是“刪除”和“新增”。

當需要刪除某個神經元時（如圖4中左圖標紅部分，對應的是編號為10的神經元），只需將num_rows的數量減1，并將最后一行Copy至被刪除行，結果如圖4的中圖所示。虛線框表示當前有效的數據。

當需要“新增”時，直接將其對應的參數由Flash memory中copy至該矩陣中即可，無需額外分配存儲空間。

5. 實驗結果

蘋果這篇paper的主要關注點在于：讓LLM在運行時內存受限的情況下能高效地跑起來。所以論文的實驗主要對比了各種情況下I/O導致的時耗，如下圖所示。

圖5 實驗結果

圖5中的實驗使用的是OPT 6.7B模型，半精度。表中第一行和第二行都是基準baseline。第一行假設初始模型全部在Flash memory中，那么為了完成一次完整的推理，需要將模型全部加載一遍，整個I/O耗時2130ms。

第二行對應于假設模型有一半參數提前在DRAM中的整個加載耗時。

第三行到第五行對應于應用了Predictor（2.2小節）、Windowing（2.3小節）和Bundling（3.1小節）后對應的耗時。

效率提升非常明顯。