小米大模型提效新框架：訓(xùn)練最高提速34%，推理最高提速52%！Kaldi之父合作出品

西風(fēng) 2024-06-24 19:16:25 來源：量子位

小米AI實(shí)驗(yàn)室 投稿

量子位 | 公眾號(hào) QbitAI

大模型推理速度提升50%以上，還能保證少樣本學(xué)習(xí)性能！

小米大模型團(tuán)隊(duì)提出SUBLLM（Subsampling-Upsampling-Bypass Large Language Model），國際AI語音大牛、開源語音識(shí)別工具Kaldi之父Daniel Povey也參與指導(dǎo)。

與Llama等模型相比，SUBLLM在訓(xùn)練和推理速度以及降低內(nèi)存方面都有了顯著提升。

在大模型訓(xùn)練中，SUBLLM的速度提高了26%，每個(gè)GPU的內(nèi)存減少了10GB。在推理中，它的速度提高了37%，每個(gè)GPU的內(nèi)存減少了1GB。

訓(xùn)練和推理速度分別最高可以提高至34%和52%。

SUBLLM通過智能地選擇和處理數(shù)據(jù)，使得模型在訓(xùn)練和推理時(shí)更加高效：子采樣模塊剔除不必要的信息，上采樣模塊恢復(fù)數(shù)據(jù)的完整性，而繞過模塊則加快了學(xué)習(xí)過程。

在一萬字中挑選最關(guān)鍵的五百字

目前，云端的大模型處理超長文本任務(wù)，通常需要?jiǎng)佑枚噙_(dá)8個(gè)GPU，這個(gè)過程不僅耗時(shí)，而且成本昂貴。如果將大模型類比于人腦，那么當(dāng)前大模型的運(yùn)行功率相比于人腦運(yùn)行功率的100倍以上。

此前，Daniel Povey在語音識(shí)別領(lǐng)域提出了Zipformer，Zipformer可以用最低壓縮16倍的幀率，達(dá)到與更大模型一致甚至更高的語音識(shí)別率，完成了語音識(shí)別領(lǐng)域的“四兩撥千斤”。

小米集團(tuán)大模型團(tuán)隊(duì)嘗試將這一思路擴(kuò)展至大型語言模型中，在性能不受損害的前提下，實(shí)現(xiàn)了更高效率的大模型運(yùn)算。

總的來說，SUBLLM的工作原理通過引入子采樣、上采樣和旁路模塊等方式，對(duì)計(jì)算資源動(dòng)態(tài)分配，從而減少了冗余的token計(jì)算負(fù)擔(dān)，加速了模型的訓(xùn)練和推理過程。

能做到就像在一萬字中挑選最關(guān)鍵的五百字一樣，保留文本中必需的部分，刪減其中的冗余，從而讓大模型所需處理的文本更短。

就實(shí)現(xiàn)路徑而言，會(huì)將子采樣模塊根據(jù)token的重要性分?jǐn)?shù)對(duì)其進(jìn)行篩選，保留重要的token并丟棄不重要的部分。

隨后，上采樣模塊將子采樣后的序列恢復(fù)到原始長度，確保語言模型在生成token時(shí)的順序一致性。

同時(shí)，旁路模塊通過結(jié)合子采樣前后的序列，進(jìn)一步提高了模型的收斂速度。這種設(shè)計(jì)不僅顯著減少了計(jì)算成本，還保持了輸入序列的語義完整性。

如果將SUBLLM理解為一個(gè)聰明的編輯，就像我們的大腦會(huì)識(shí)別要點(diǎn)一樣，它可以在閱讀一大段文字時(shí)快速識(shí)別出哪些詞是關(guān)鍵的，哪些詞不那么重要。SUBLLM會(huì)保留那些重要的詞匯，而忽略那些不太重要的部分，這就大大減少了需要處理的信息量。

隨后，就像我們能通過只言片語補(bǔ)充完整故事的來龍去脈，SUBLLM也能將精簡(jiǎn)后的信息恢復(fù)到原有的完整度，確保整個(gè)文本在表達(dá)時(shí)的連貫與完整。在處理信息時(shí)，SUBLLM還能更加迅速地找到最佳的表達(dá)方式。

接下來具體看SUBLLM的模型結(jié)構(gòu)。

SUBLLM具體長啥樣？

前不久，谷歌Deepmind提出了mixture of depths（MoD）模型結(jié)構(gòu)，MoD使用靜態(tài)計(jì)算預(yù)算，使用每個(gè)塊的路由器選擇token進(jìn)行計(jì)算，并通過對(duì)自注意力和MLP塊或殘差連接的選擇來優(yōu)化FLOP使用。

更早以前，經(jīng)典論文CoLT5使用條件路由來決定給定token是通過輕量分支還是重量分支在前饋和注意力層中傳遞，以便將更多資源分配給重要token。

與這些模型結(jié)構(gòu)類似，SUBLLM采用的原理接近于人腦對(duì)于信息的處理機(jī)制。

人腦有兩種思維模式，一種低功耗的快模式，一種高功耗的慢模式，分工明確，且兩種模式恰恰用的是同一個(gè)腦部區(qū)域。

因此，SUBLLM作者也從這一信息處理模式的角度思考了如何將大模型的算力進(jìn)行合理地分配：重要的token用全部算力，相對(duì)不重要的token使用更少算力。

具體來說，SUBLLM的模型結(jié)構(gòu)是基于decoder-only的大語言模型架構(gòu)，在不改變?cè)心Ｐ徒Y(jié)構(gòu)的基礎(chǔ)上，在一些特殊的層上進(jìn)行了結(jié)構(gòu)升級(jí)。

為了管理要處理的token數(shù)量，子采樣和上采樣模塊被集成到Transformer塊之間。

首先，模型使用幾個(gè)Transformer塊處理完整序列，捕獲全面的token序列表示。

引入子采樣模塊后，這些模塊暫時(shí)去除不關(guān)鍵的token，從而減少處理所需的序列長度。

然后對(duì)縮減后的序列進(jìn)行更多次的子采樣過程，也就是序列的縮減是嵌套的。序列壓縮的最高級(jí)別發(fā)生在網(wǎng)絡(luò)的最中間的Transformer塊中。

隨后，使用上采樣模塊逐步恢復(fù)序列長度。這些模塊將較短的處理序列與子采樣前的原始序列合并，將它們恢復(fù)到完整長度。

這種機(jī)制允許僅解碼器模型作為語言模型操作，按順序生成token，保證輸入和輸出序列長度相同。

此外，上采樣過程后集成了繞過連接模塊，以利用每個(gè)子采樣前的嵌入，幫助改進(jìn)從子采樣到上采樣的學(xué)習(xí)過程。

隨后的實(shí)驗(yàn)證實(shí)，這種方法顯著提高了收斂效率。

與LLaMA模型相比，SUBLLM在訓(xùn)練和推理方面分別實(shí)現(xiàn)了26%和37%的速度提升，同時(shí)顯著降低了內(nèi)存成本，同時(shí)保持了性能。

預(yù)訓(xùn)練階段、推理階段計(jì)算效率的詳細(xì)分析：

論文鏈接：https://arxiv.org/abs/2406.06571