大模型生成提速2倍！單GPU幾小時(shí)搞定微調(diào)，北大數(shù)院校友共同一作丨開源

　　本文來(lái)自于微信公眾號(hào) 量子位 (ID:QbitAI)，作者:蕭蕭 。

　　只需給大模型“加點(diǎn)小零件”，推理速度立刻提升2倍!

　　不需要額外訓(xùn)練一個(gè)模型，也不需要對(duì)計(jì)算硬件做優(yōu)化，單張A100最快幾小時(shí)就能微調(diào)完成。

　　這項(xiàng)新研究名叫Medusa(美杜莎)，來(lái)自普林斯頓、UIUC、CMU和康涅狄格大學(xué)，F(xiàn)lashAttention作者Tri Dao也在其中。

　　目前，它已經(jīng)成功部署到伯克利70億參數(shù)的“駱馬”Vicuna中，后續(xù)還會(huì)支持其他大模型，已經(jīng)登上GitHub熱榜:

　　但其實(shí)，在這種方法推出之前，業(yè)界并非沒(méi)有大模型推理加速方法，主流的就是DeepMind推出的投機(jī)采樣(speculative decoding)。

　　相比這種方法，Medusa有什么不一樣的地方?投機(jī)采樣的2個(gè)“bug”

　　要想加速大模型推理，需要先知道究竟是什么“限制”了它的速度。

　　相比計(jì)算量的增加，大模型推理速度更容易受到內(nèi)存帶寬的影響(memory bound)。

　　這是因?yàn)�，大模型由于參�?shù)量巨大、遠(yuǎn)超緩存容量，因此推理時(shí)需要先把權(quán)重從外部?jī)?nèi)存(顯存)讀取一次到緩存中，這個(gè)過(guò)程受內(nèi)存帶寬限制，速度通常很慢。

　　因此，模型做批量推理(batch inference)時(shí)，一次處理100個(gè)tokens和一個(gè)tokens時(shí)間上區(qū)別不大。

　　基于這個(gè)特點(diǎn)，DeepMind去年11月想出了一個(gè)名叫投機(jī)采樣的神奇操作——

　　訓(xùn)練一個(gè)更小的模型(draft模型)，給大模型提前生成一批“候選詞”，相比于讓大模型自己“思考”生成，直接做“選擇”就好。

　　由于小模型生成速度比大模型快好幾倍，一旦大模型覺(jué)得小模型已有的詞“可用”，就直接拿來(lái)，不用自己再緩慢生成一遍。

　　這個(gè)過(guò)程，有點(diǎn)像是輸入法的聯(lián)想詞候選，在我們(大模型)想好下一個(gè)詞用什么之前，輸入法(小模型)先給列出一些備選項(xiàng):

　　要是看到覺(jué)得不錯(cuò)，就從中選一個(gè)用;要是覺(jué)得生成的都不行，就pass掉自己重新打。

　　這種投機(jī)采樣方法確實(shí)取得了顯著成效，甚至能輕輕松松在M2Ultra上以高精度跑340億參數(shù)LLaMA大模型。

　　BUT，這種方法存在兩個(gè)問(wèn)題。

　　一方面，給大模型找個(gè)生成“候選詞”的draft小模型，沒(méi)那么容易。

　　這個(gè)小模型可不是隨便抓個(gè)生成模型就能用，除了接口統(tǒng)一、概率分布接近等要求，生成質(zhì)量也不能比大模型差太多。

　　對(duì)于Meta發(fā)布的LLaMA這種模型可能還好，既有幾百億參數(shù)的大模型版本，又有幾十億參數(shù)的小模型版本，可以把參數(shù)量更小的版本拿來(lái)當(dāng)draft模型使用。

　　但對(duì)于其他開源大模型，這種方法就不太適用了，自己去搭建訓(xùn)練一個(gè)小模型，不僅時(shí)間成本更高，生成效果可能還不達(dá)預(yù)期。

　　另一方面，雙模型的組合，使得后續(xù)要想做系統(tǒng)調(diào)優(yōu)變得更復(fù)雜。

　　這是因?yàn)椋�相比于大模型自身是一個(gè)系統(tǒng)，新增加的draft模型相當(dāng)于又引入了一個(gè)系統(tǒng)。

　　這樣會(huì)導(dǎo)致模型部署起來(lái)更復(fù)雜，包括額外的網(wǎng)絡(luò)傳輸、不同的硬件條件都需要考慮到，在做計(jì)算優(yōu)化時(shí)難度也會(huì)進(jìn)一步提升。

　　為了解決這些問(wèn)題，Medusa出現(xiàn)了。 不用小模型，加幾個(gè)“頭”就行

　　Medusa(美杜莎，一種長(zhǎng)有多個(gè)頭的妖怪)是一種新的大模型推理加速方法。

　　相比投機(jī)采樣，它選擇直接給Transformer大模型多加幾個(gè)解碼頭(decoding heads)，每個(gè)頭都是一個(gè)單層前饋網(wǎng)絡(luò)。

　　這幾個(gè)多出來(lái)的解碼頭，可以讓大模型直接一次多生成幾個(gè)詞，而不是“擠牙膏式”一個(gè)一個(gè)生成。

　　生成準(zhǔn)確率也還可以，在預(yù)測(cè)“下一個(gè)詞的下一個(gè)詞”時(shí)，Medusa準(zhǔn)確率達(dá)到了60%，還在不斷優(yōu)化中。

　　隨后，結(jié)合樹狀注意力機(jī)制(tree-based attention mechanism)并行驗(yàn)證這些詞，從而實(shí)現(xiàn)推理加速。

　　基于Medusa，Vicuna的70億、130億和330億參數(shù)大模型推理速度，均有了1.9倍以上的效率提升:

　　針對(duì)70億參數(shù)的模型，研究者們還在不同任務(wù)上測(cè)試了一下加速效果，顯示最高在代碼生成上有2.15倍的速度提升。

　　最關(guān)鍵的是，用上Medusa后，并不需要將整個(gè)大模型重新訓(xùn)練一遍。

　　相比之下，它可以和大模型一起訓(xùn)練，只需要凍結(jié)大模型的參數(shù)就行，甚至單個(gè)GPU就能搞定。

　　由于不增加額外的模型，對(duì)于分布式推理也很友好。 作者介紹

　　這項(xiàng)研究有兩位共同一作。

　　共同一作蔡天樂(lè)，普林斯頓大學(xué)博士生，研究方向包括優(yōu)化、表示學(xué)習(xí)、架構(gòu)設(shè)計(jì)等，本科畢業(yè)于北京大學(xué)數(shù)學(xué)科學(xué)學(xué)院，獲得應(yīng)用數(shù)學(xué)和計(jì)算機(jī)科學(xué)雙學(xué)位。

　　共同一作Yuhong(Jesse)Li，伊利諾伊大學(xué)香檳分校(UIUC)博士生，研究方向是高效機(jī)器學(xué)習(xí)，本科畢業(yè)于北京郵電大學(xué)。

　　此外，這項(xiàng)研究也有FlashAttention作者、斯坦福博士Tri Dao的參與。

　　FlashAttention是一種能加快注意力并減少內(nèi)存占用的方法，相比PyTorch標(biāo)準(zhǔn)注意力實(shí)現(xiàn)，最高能提速9倍。