Apple最近發(fā)表了一篇文章，可以在iphone, MAC 上運(yùn)行大模型：【LLM in a flash: Efficient Large Language Model Inference with Limited Memory】。

主要解決的問(wèn)題是在DRAM中無(wú)法存放完整的模型和計(jì)算，但是Flash Memory可以存放完整的模型。但是Flash帶寬較低，LLM in Flash通過(guò)盡量減少?gòu)腇lash中加載參數(shù)的數(shù)量，優(yōu)化在DRAM中的內(nèi)存管理，實(shí)現(xiàn)在Flash帶寬有限的條件下提高計(jì)算速度的目的。

這篇文章很多都是工程上的細(xì)節(jié)，很少理論。下面是這篇論文的總結(jié)，如有不對(duì)的地方，歡迎私信。

利用FeedForward 層的稀疏度，只加載FeedForward層輸入非0和預(yù)測(cè)輸出非0的參數(shù)

通過(guò)Window Sliding 只加載增量的參數(shù)，復(fù)用之前的計(jì)算，減少需要加載的參數(shù)。

將up-projection的row和down-projection的column放在一起存放，這樣在flash中可以一次讀取比較大的chunk，提高flash的帶寬利用效率。

如下圖所示，chunk越大，帶寬也就越大，初始加載chunk的latency可以被平攤。

1. 大模型在FeedForward layer有超過(guò)90%的稀疏度。將embedding和attention 矩陣一直保存在DRAM中。Attention 的權(quán)重占據(jù)了model總量的1/3。對(duì)于FeedForward Layer，只有非稀疏的部分被動(dòng)態(tài)的加載進(jìn)去DRAM。

2. 預(yù)測(cè)Relu層的稀疏性。在attention層的輸出后面增加low-rank predictor，預(yù)測(cè)在relu層之后可能是0的元素。

經(jīng)過(guò)優(yōu)化后，最終只需要加載2%的FeedForward層的參數(shù)到DRAM中。

3. Sliding Window

每次滑動(dòng)窗口，在生成新的token后，刪掉不在window內(nèi)的neuron，增加新的neuron。

上圖右側(cè)為一個(gè)window size為5的示意圖，粉色的是要?jiǎng)h除的元素，藍(lán)色的是新加入的元素。

上圖左側(cè)是如何在aggregated usage和incremental transfer中保持平衡，window設(shè)置的越大，每次新需要加載neruon也就越少，但是需要在memory中累計(jì)保存的空間占用的也就越大。

上圖左側(cè)的目標(biāo)就是如何讓aggregated usage和incremental transfer都比較小。

譯者疑問(wèn)：這個(gè)window就是Longformer: The Long-Document Transformer中的sliding window嗎？歡迎私信。

4.內(nèi)存管理

內(nèi)存管理也是因?yàn)閟liding window引入的。

譯者注：

就像c++中vector的維護(hù)一樣，如果每次刪除vector中間的一個(gè)元素，都需要導(dǎo)致該元素后面所有元素的移動(dòng)。

下圖描述的就是刪除和加入新neuron的內(nèi)容。

最后文章還提出了比較有意思的一點(diǎn)，他們主要進(jìn)行了稀疏化的優(yōu)化，在計(jì)算和加載參數(shù)方面。他們也嘗試了通過(guò)和當(dāng)前neuron關(guān)系緊密的 “closest friend”綁定，每次加載neuron時(shí)，也都加載他的closest friend。

作者說(shuō)但是這樣帶來(lái)了負(fù)面作用，因?yàn)榇嬖谝恍ヽlosest friend是很多neuron的closest friend (譯者注：類似于大眾之友)，這些neuron被頻繁的加載到DRAM中，反而降低了性能。

審核編輯：湯梓紅