微軟打破Decoder-Only架構(gòu)！大幅降低GPU內(nèi)存需求，網(wǎng)友：把Llama3 70B弄20GB GPU上運(yùn)行
西風(fēng) 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI
微軟&清華最新研究，打破GPT系列開創(chuàng)的Decoder-Only架構(gòu)——
提出Decoder-Decoder新型架構(gòu)，名為YOCO（You Only Cache Once）。
YOCO僅緩存一次鍵值對，可大幅降低GPU內(nèi)存需求，且保留全局注意力能力。
一張圖來看YOCO和標(biāo)準(zhǔn)Transformer的比較。
在處理512K上下文長度時(shí)，標(biāo)準(zhǔn)Transformer內(nèi)存使用是YOCO的6.4倍，預(yù)填充延遲是YOCO的30.3倍，而YOCO的吞吐量提升到標(biāo)準(zhǔn)Transformer的9.6倍。
去年一張“大語言模型進(jìn)化樹”動圖在學(xué)術(shù)圈瘋轉(zhuǎn)，模型架構(gòu)還只有三大類：Decoder-Only、Encoder-Only、Encoder-Decoder。
那么這個(gè)新出的Decoder-Decoder架構(gòu)到底長啥樣？
嗯，如網(wǎng)友所言，要讀的論文又增加了。
話不多說，一起來看。
打破Decoder-Only
YOCO整體架構(gòu)設(shè)計(jì)如下，分為自解碼器（Self-Decoder）和交叉解碼器（Cross-Decoder）兩部分。
具體來說，YOCO由L個(gè)塊堆疊而成，其中前L/2層是自解碼器，其余模塊是交叉解碼器。
自解碼器利用高效自注意力（efficient self-attention）機(jī)制來獲取鍵值（KV）緩存：
接收輸入序列的嵌入表示，并使用高效自注意力來生成中間向量表示；使用因果掩碼（causal masking）保證解碼的自回歸特性；自解碼器的輸出用于生成全局KV緩存。
而交叉解碼器使用交叉注意力（cross-attention）來重用自解碼器生成的共享KV緩存：
在自解碼器生成的KV緩存基礎(chǔ)上進(jìn)行堆疊，以獲得最終的輸出向量；同樣使用因果掩碼來維持自回歸生成；允許交叉解碼器層間高效地重用KV緩存，減少了對GPU內(nèi)存的需求。
總的來說，自解碼器和交叉解碼器的模塊設(shè)計(jì)與Transformer的解碼器層類似，包含交錯注意力和前饋網(wǎng)絡(luò)子層。不過，研究人員還進(jìn)行了預(yù)RMSNorm、SwiGLU和分組查詢注意力等改進(jìn)。
兩部分之間的區(qū)別在于注意力模塊。
自解碼器使用高效自注意力，如滑動窗口注意力（Sliding-Window Attention）或門控保留（gated retention）。
而交叉解碼器使用標(biāo)準(zhǔn)的多頭交叉注意力，Query向量通過注意力與自解碼器產(chǎn)生的全局鍵值緩存相關(guān)聯(lián)。
推理大幅度省 省 省
實(shí)驗(yàn)階段，研究人員將YOCO模型與同體量的Transformer模型進(jìn)行比較。
分析維度有四個(gè)：語言建模評估、與Transformer比較的可擴(kuò)展性、長上下文評估、推理優(yōu)勢。
語言建模評估
研究人員訓(xùn)練了一個(gè)3B參數(shù)的YOCO語言模型，并根據(jù)訓(xùn)練token數(shù)量（1T和1.6T）進(jìn)行評估。
在LM Eval Harness的多個(gè)下游任務(wù)上，YOCO與Transformer模型OpenLLaMA-3B-v2、StableLM-base-alpha-3B-v2、StableLM-3B-4E1T打得有來有回。
可擴(kuò)展性對比
接著，研究人員在160M到13B參數(shù)規(guī)模范圍內(nèi)，分別訓(xùn)練了YOCO（門控保留和滑動窗口注意力版本）和Transformer語言模型。
對比了它們在驗(yàn)證集上的語言模型損失，YOCO的表現(xiàn)與Transformer基本持平：
結(jié)果證明YOCO在模型大小擴(kuò)展方面具有很強(qiáng)的可擴(kuò)展性。
長上下文評估
將3B的YOCO模型擴(kuò)展到上下文為1M，在“大海撈針”等長序列的needle retrieval任務(wù)上，YOCO-3B-1M的準(zhǔn)確率接近100%。
在多針檢索任務(wù)上，YOCO-3B-1M的性能優(yōu)于一些超3B的Transformer模型：
此外，YOCO模型在長序列上的NLL隨著上下文長度的增加而一致下降，表明YOCO能夠有效地利用長距離依賴信息進(jìn)行語言建模：
綜上，可見YOCO在性能上完全不輸Transformer，關(guān)鍵來看YOCO在推理效率上取得的顯著提升。
推理優(yōu)勢
研究人員評估了YOCO在GPU內(nèi)存占用、prefilling延遲、吞吐量和服務(wù)容量等方面的優(yōu)勢，評估上下文范圍為32K至1M。
如下圖所示，與Transformer相比，YOCO大幅度降低了GPU內(nèi)存占用，且YOCO的內(nèi)存消耗隨上下文長度增加，增長幅度很小。
例如，在1M長度下，整體推理內(nèi)存使用量僅為12.4GB，而傳統(tǒng)的Transformer則占用了9.38倍的GPU內(nèi)存。
下面展示了token的KV緩存對GPU內(nèi)存的占用情況。
YOCO模型只緩存一層全局的鍵值對，因此與Transformer模型相比，它需要的內(nèi)存約少了L（指模型的層數(shù)）倍。
例如，YOCO模型可以使用1GB的GPU內(nèi)存來處理128K token。而具有GQA的Transformer 65B大小模型，僅能支持1.6K token。
也就是說，模型越大，YOCO可以節(jié)省更多。
在預(yù)填充階段，模型并行編碼輸入token。對于512K和1M長度的輸入，Transformer分別需要大約180秒和300秒。Transformer的計(jì)算復(fù)雜度為O(N^2)，處理長上下文需要大量的浮點(diǎn)運(yùn)算操作。
相比之下，YOCO的預(yù)填充時(shí)間為O(N)，隨序列長度線性增長。
YOCO將Transformer的512K上下文預(yù)填充時(shí)間從180秒減少到不到6秒。
預(yù)填充階段可以在進(jìn)入交叉解碼器之前提前退出。因此，即使對于短上下文，預(yù)填充延遲的加速至少是兩倍。例如，對于32K長度，YOCO比Transformer快2.87倍。
吞吐量表示模型每秒可以處理多少個(gè)token，涵蓋了預(yù)填充和生成時(shí)間。如下圖所示，與Transformer相比，YOCO在不同上下文長度下實(shí)現(xiàn)了更高的吞吐量。
以512K查詢?yōu)槔?，Transformer的吞吐量為4.5 token/秒，而YOCO達(dá)到了43.1token/秒，即實(shí)現(xiàn)了9.6倍的加速。
吞吐量提高的原因如前所述，YOCO減少了預(yù)填充所需的時(shí)間。其次，由于內(nèi)存消耗減少，因此可以在推理時(shí)使用更大的批量大小，這也有助于提高吞吐量。
詳細(xì)細(xì)節(jié)，感興趣的家人們可以查看原論文。
論文鏈接：https://arxiv.org/abs/2405.05254
點(diǎn)這里??關(guān)注我，記得標(biāo)星哦～
科技前沿進(jìn)展日日相見 ~?
