推理大語言模型通過將復(fù)雜問題分解為一系列較小步驟來解決難題。這些強(qiáng)大的模型在高級編程和多步規(guī)劃等具有挑戰(zhàn)性的任務(wù)上表現(xiàn)出色。

但是，開發(fā)推理模型需要巨大的計算量和能耗，這是由于訓(xùn)練過程中的低效率造成的。當(dāng)少數(shù)高性能處理器持續(xù)處理復(fù)雜查詢時，群組中的其他處理器卻處于閑置狀態(tài)。

來自MIT等機(jī)構(gòu)的研究人員找到了一種利用這些計算空閑時間來有效加速推理模型訓(xùn)練的方法。

他們的新方法能夠自動訓(xùn)練一個更小、更快的模型來預(yù)測較大推理大語言模型的輸出，然后由較大模型進(jìn)行驗證。這減少了推理模型必須完成的工作量，從而加速了訓(xùn)練過程。

該系統(tǒng)的關(guān)鍵在于能夠自適應(yīng)地訓(xùn)練和部署較小的模型，使其僅在某些處理器空閑時才介入工作。通過利用原本會被浪費的計算資源，它在不產(chǎn)生額外開銷的情況下加速了訓(xùn)練。

在多個推理大語言模型上進(jìn)行測試時，該方法在保持準(zhǔn)確性的同時將訓(xùn)練速度提高了一倍。這可以降低成本并提高開發(fā)高級大語言模型的能源效率，這些模型可應(yīng)用于金融趨勢預(yù)測或電網(wǎng)風(fēng)險檢測等領(lǐng)域。

MIT博士后、該技術(shù)論文的共同第一作者胡慶豪表示："人們希望模型能夠處理更復(fù)雜的任務(wù)。但如果這是模型開發(fā)的目標(biāo)，那么我們需要優(yōu)先考慮效率。我們找到了這個問題的無損解決方案，然后開發(fā)了一個能夠在實踐中實現(xiàn)相當(dāng)顯著加速的全棧系統(tǒng)。"

該論文的其他作者包括共同第一作者、電子工程與計算機(jī)科學(xué)研究生楊尚；電子工程與計算機(jī)科學(xué)研究生郭俊賢；資深作者、電子工程與計算機(jī)科學(xué)副教授、電子研究實驗室成員及NVIDIA杰出科學(xué)家韓松，以及來自NVIDIA、蘇黎世聯(lián)邦理工學(xué)院、MIT-IBM Watson人工智能實驗室和馬薩諸塞大學(xué)阿默斯特分校的其他研究人員。該研究將在ACM編程語言和操作系統(tǒng)架構(gòu)支持國際會議上發(fā)表。

訓(xùn)練瓶頸

開發(fā)人員希望推理大語言模型能夠識別并糾正其批判性思維過程中的錯誤。這種能力使它們能夠處理會困擾標(biāo)準(zhǔn)大語言模型的復(fù)雜查詢。

為了教授這項技能，開發(fā)人員使用一種稱為強(qiáng)化學(xué)習(xí)的技術(shù)來訓(xùn)練推理大語言模型。模型為查詢生成多個潛在答案，為最佳候選答案獲得獎勵，并根據(jù)最佳答案進(jìn)行更新。這些步驟重復(fù)數(shù)千次，模型在此過程中學(xué)習(xí)。

但研究人員發(fā)現(xiàn)，生成多個答案的過程（稱為rollout）可能消耗強(qiáng)化學(xué)習(xí)訓(xùn)練所需執(zhí)行時間的多達(dá)85%。

胡慶豪說："相比之下，更新模型——這才是實際的'訓(xùn)練'部分——消耗的時間很少。"

這個瓶頸出現(xiàn)在標(biāo)準(zhǔn)強(qiáng)化學(xué)習(xí)算法中，因為訓(xùn)練組中的所有處理器必須完成其響應(yīng)后才能進(jìn)入下一步。由于某些處理器可能正在處理很長的響應(yīng)，而生成較短響應(yīng)的其他處理器則等待它們完成。

胡慶豪補(bǔ)充說："我們的目標(biāo)是將這種空閑時間轉(zhuǎn)化為加速，而不產(chǎn)生任何浪費的成本。"

他們試圖使用一種稱為推測解碼的現(xiàn)有技術(shù)來加速過程。推測解碼涉及訓(xùn)練一個稱為草案模型的較小模型來快速猜測較大模型的未來輸出。

較大模型驗證草案模型的猜測，其接受的響應(yīng)用于訓(xùn)練。

因為較大模型可以一次驗證草案模型的所有猜測，而不是按順序生成每個輸出，所以它加速了過程。

自適應(yīng)解決方案

但在推測解碼中，草案模型通常只訓(xùn)練一次并保持靜態(tài)。這使得該技術(shù)對強(qiáng)化學(xué)習(xí)不可行，因為推理模型在訓(xùn)練期間會更新數(shù)千次。

靜態(tài)草案模型在幾步之后很快就會變得過時和無用。

為了克服這個問題，研究人員創(chuàng)建了一個名為"馴服長尾"（TLT）的靈活系統(tǒng)。

TLT的第一部分是自適應(yīng)草案訓(xùn)練器，它利用空閑處理器的空閑時間來動態(tài)訓(xùn)練草案模型，使其與目標(biāo)模型保持良好對齊，而不使用額外的計算資源。

第二個組件是自適應(yīng)rollout引擎，管理推測解碼以自動為每批新輸入選擇最優(yōu)策略。該機(jī)制根據(jù)訓(xùn)練工作負(fù)載特征改變推測解碼配置，如草案模型處理的輸入數(shù)量和驗證期間目標(biāo)模型接受的輸入數(shù)量。

此外，研究人員將草案模型設(shè)計得輕量化，以便可以快速訓(xùn)練。TLT重用推理模型訓(xùn)練過程的一些組件來訓(xùn)練草案模型，從而獲得額外的加速收益。

胡慶豪說："一旦某些處理器完成了短查詢并變得空閑，我們立即切換它們使用相同的數(shù)據(jù)進(jìn)行草案模型訓(xùn)練，這些數(shù)據(jù)也用于rollout過程。關(guān)鍵機(jī)制是我們的自適應(yīng)推測解碼——沒有它這些收益是不可能的。"

他們在使用真實世界數(shù)據(jù)集訓(xùn)練的多個推理大語言模型上測試了TLT。該系統(tǒng)在保持每個模型準(zhǔn)確性的同時，將訓(xùn)練速度提高了70%到210%。

作為額外的好處，小型草案模型可以作為免費副產(chǎn)品用于高效部署。

未來，研究人員希望將TLT集成到更多類型的訓(xùn)練和推理框架中，并找到可以使用這種方法加速的新強(qiáng)化學(xué)習(xí)應(yīng)用。

韓松說："隨著推理繼續(xù)成為驅(qū)動推理需求的主要工作負(fù)載，慶豪的TLT是應(yīng)對訓(xùn)練這些推理模型計算瓶頸的出色工作。我認(rèn)為這種方法在高效AI計算的背景下將非常有用。"

這項工作由MIT-IBM Watson人工智能實驗室、MIT AI硬件項目、MIT亞馬遜科學(xué)中心、現(xiàn)代汽車公司和美國國家科學(xué)基金會資助。

Q&A

Q1：TLT方法是什么？它如何提升大語言模型訓(xùn)練效率？

A：TLT（馴服長尾）是MIT研究人員開發(fā)的新方法，它通過訓(xùn)練一個較小的草案模型來預(yù)測大語言模型的輸出，然后由大模型驗證。關(guān)鍵是利用空閑處理器的計算時間來自適應(yīng)訓(xùn)練草案模型，將原本浪費的計算資源轉(zhuǎn)化為訓(xùn)練加速，在不增加額外開銷的情況下將訓(xùn)練速度提高70%到210%。

Q2：為什么推理大語言模型訓(xùn)練會出現(xiàn)效率瓶頸？

A：瓶頸主要出現(xiàn)在強(qiáng)化學(xué)習(xí)的rollout過程中，這個過程需要生成多個潛在答案，可能消耗高達(dá)85%的執(zhí)行時間。由于所有處理器必須等待最慢的處理器完成長響應(yīng)后才能進(jìn)入下一步，導(dǎo)致處理短響應(yīng)的處理器長時間閑置，造成計算資源浪費。

Q3：TLT方法與傳統(tǒng)推測解碼有什么不同？

A：傳統(tǒng)推測解碼中的草案模型是靜態(tài)的，只訓(xùn)練一次，在強(qiáng)化學(xué)習(xí)中會快速變得過時無用。TLT的創(chuàng)新在于自適應(yīng)特性：它包含自適應(yīng)草案訓(xùn)練器，能在訓(xùn)練過程中動態(tài)更新草案模型；還有自適應(yīng)rollout引擎，能根據(jù)工作負(fù)載特征自動選擇最優(yōu)策略配置。