普林斯頓大學(xué)團(tuán)隊(duì)如何讓AI專家模型訓(xùn)練速度翻倍

這篇由普林斯頓大學(xué)郭文濤、程新樂，加州大學(xué)伯克利分校米什拉、斯托伊察，以及Together AI的趙天分別在2025年12月17日發(fā)表的最新研究報(bào)告，介紹了一項(xiàng)名為SonicMoE的突破性技術(shù)。這項(xiàng)研究針對(duì)目前人工智能領(lǐng)域最重要的"專家混合模型"(Mixture of Experts，簡(jiǎn)稱MoE)在訓(xùn)練過程中遇到的瓶頸，提出了一套完整的解決方案。有興趣深入了解的讀者可以通過論文編號(hào)arXiv:2512.14080查詢完整論文。

要理解這項(xiàng)研究的重要性，我們首先要明白什么是專家混合模型。這種模型就像一個(gè)擁有眾多專業(yè)老師的學(xué)校，每個(gè)老師只專精某一個(gè)科目，當(dāng)學(xué)生有特定問題時(shí)，會(huì)被引導(dǎo)到最合適的老師那里獲得幫助。這種設(shè)計(jì)讓AI模型在不大幅增加計(jì)算成本的前提下，顯著擴(kuò)展了知識(shí)儲(chǔ)備和處理能力。

然而，隨著AI技術(shù)的發(fā)展，研究人員發(fā)現(xiàn)一個(gè)有趣的現(xiàn)象：讓每個(gè)"專家老師"更加專業(yè)化（也就是處理更細(xì)分的領(lǐng)域），以及增加專家數(shù)量同時(shí)保持激活專家數(shù)量不變（提高稀疏性），確實(shí)能讓模型表現(xiàn)得更好。這就好比把一個(gè)通用數(shù)學(xué)老師拆分成代數(shù)專家、幾何專家、微積分專家等，每個(gè)專家只負(fù)責(zé)自己最擅長(zhǎng)的部分，學(xué)生學(xué)習(xí)效果會(huì)更好。

不過，這種做法也帶來了問題。當(dāng)專家變得越來越專業(yè)化時(shí)，系統(tǒng)需要記住更多信息來協(xié)調(diào)這些專家之間的工作，這就像學(xué)校需要更大的檔案室來存放每個(gè)老師的課程資料。同時(shí)，由于每個(gè)專家處理的內(nèi)容變少了，計(jì)算資源的利用效率也下降了，就好比讓一個(gè)專業(yè)的心臟外科醫(yī)生去處理簡(jiǎn)單的感冒，既浪費(fèi)了專家的時(shí)間，也增加了系統(tǒng)的復(fù)雜性。

一、重新設(shè)計(jì)大腦的記憶系統(tǒng)

研究團(tuán)隊(duì)首先解決的是內(nèi)存占用問題。傳統(tǒng)的專家混合模型在訓(xùn)練過程中需要保存大量中間結(jié)果，這些信息在后續(xù)的學(xué)習(xí)過程中要用到。隨著專家變得更加專業(yè)化，需要保存的信息量呈直線上升，就像一個(gè)圖書館隨著藏書分類越來越細(xì)，需要的目錄卡片也越來越多。

研究團(tuán)隊(duì)通過仔細(xì)分析整個(gè)計(jì)算過程，發(fā)現(xiàn)了一個(gè)巧妙的解決方案。他們重新設(shè)計(jì)了計(jì)算路徑，避免保存某些中間結(jié)果，而是在需要時(shí)重新計(jì)算。這種做法類似于一個(gè)聰明的圖書管理員，不再存儲(chǔ)所有可能用到的索引卡片，而是建立了一套高效的即時(shí)查找系統(tǒng)。

具體來說，傳統(tǒng)方法需要保存每個(gè)專家的輸出結(jié)果Y和相應(yīng)的梯度信息dY，用于后續(xù)的參數(shù)更新。SonicMoE通過數(shù)學(xué)推導(dǎo)，找到了一條不需要這些信息的計(jì)算路徑。他們將路由器評(píng)分的梯度計(jì)算重新表述為dS = ?dA', A?的形式，而不是傳統(tǒng)的dS = ?dO, Y?。這種改變看似微小，實(shí)際上消除了對(duì)大量中間結(jié)果的依賴。

這種優(yōu)化的效果非常明顯。對(duì)于一個(gè)7B參數(shù)的精細(xì)化專家混合模型，SonicMoE將每層的激活內(nèi)存使用量減少了45%。更重要的是，無論專家變得多么專業(yè)化，內(nèi)存使用量都保持恒定，這為訓(xùn)練更大更復(fù)雜的模型提供了可能。

二、讓計(jì)算和數(shù)據(jù)傳輸并行進(jìn)行

第二個(gè)創(chuàng)新點(diǎn)是充分利用現(xiàn)代GPU硬件的異步處理能力。這就好比一個(gè)餐廳廚房，傳統(tǒng)做法是廚師完成一道菜的全部步驟后再開始下一道，而SonicMoE的方法是讓廚師在炒菜的同時(shí)，助手已經(jīng)開始準(zhǔn)備下一道菜的食材，實(shí)現(xiàn)真正的流水線作業(yè)。

在現(xiàn)代GPU（特別是NVIDIA的Hopper和Blackwell架構(gòu)）中，矩陣乘法運(yùn)算和內(nèi)存讀寫操作可以同時(shí)進(jìn)行。SonicMoE巧妙地利用了這一特性，設(shè)計(jì)了一種被稱為"乒乓調(diào)度"的機(jī)制。在這種機(jī)制下，當(dāng)一個(gè)計(jì)算單元在執(zhí)行矩陣乘法時(shí)，另一個(gè)單元同時(shí)在加載下一批數(shù)據(jù)，兩者交替進(jìn)行，充分利用了硬件資源。

這種設(shè)計(jì)特別適合精細(xì)化的專家混合模型，因?yàn)檫@類模型的計(jì)算強(qiáng)度相對(duì)較低，傳統(tǒng)方法往往被內(nèi)存訪問速度限制。通過重疊計(jì)算和數(shù)據(jù)傳輸，SonicMoE能夠在相同硬件上實(shí)現(xiàn)更高的整體吞吐量。

研究團(tuán)隊(duì)還在內(nèi)存管理上做了進(jìn)一步優(yōu)化。他們避免了傳統(tǒng)方法中的同步寫入操作，而是使用異步的張量?jī)?nèi)存訪問（TMA）指令。這類似于快遞員不再等待每個(gè)包裹的簽收確認(rèn)，而是將包裹放在指定位置后立即前往下一個(gè)目的地，大大提高了整體配送效率。

三、消除計(jì)算資源浪費(fèi)的智能路由

第三個(gè)重要?jiǎng)?chuàng)新是"令牌舍入路由"方法。要理解這個(gè)概念，我們需要知道GPU在處理矩陣運(yùn)算時(shí)，為了效率考慮，通常將數(shù)據(jù)分成固定大小的"瓦片"進(jìn)行處理，比如每次處理128個(gè)數(shù)據(jù)點(diǎn)。但在實(shí)際應(yīng)用中，分配給某個(gè)專家的數(shù)據(jù)量可能是任意數(shù)字，比如145個(gè)，這就需要補(bǔ)齊到256個(gè)才能充分利用兩個(gè)瓦片，剩余的111個(gè)位置就被浪費(fèi)了。

SonicMoE的令牌舍入方法就像一個(gè)聰明的班車調(diào)度員，會(huì)調(diào)整每趟班車的乘客數(shù)量，確保每輛車都能坐滿或接近坐滿。具體做法是，對(duì)于每個(gè)專家接收到的令牌數(shù)量，自動(dòng)調(diào)整到最接近的瓦片大小倍數(shù)。如果某個(gè)專家原本應(yīng)該處理145個(gè)令牌，系統(tǒng)會(huì)智能地決定是調(diào)整到128個(gè)（舍棄17個(gè)）還是256個(gè)（增加111個(gè)），選擇偏差最小的方案。

這種方法的巧妙之處在于，它在保持原始令牌選擇邏輯基本不變的前提下，消除了GPU計(jì)算中的填充浪費(fèi)。實(shí)驗(yàn)結(jié)果顯示，在高度稀疏的專家混合模型中，這種方法能夠帶來額外16%的速度提升，而且對(duì)模型的最終性能沒有負(fù)面影響。

研究團(tuán)隊(duì)通過大量實(shí)驗(yàn)驗(yàn)證了令牌舍入方法的有效性。他們訓(xùn)練了多個(gè)不同規(guī)模的模型（從5億到14億參數(shù)），測(cè)試結(jié)果表明，使用令牌舍入訓(xùn)練的模型在各種下游任務(wù)上的表現(xiàn)與傳統(tǒng)方法相當(dāng)，有時(shí)甚至略好。這證明了這種優(yōu)化方法不會(huì)損害模型的學(xué)習(xí)能力。

四、全面的性能評(píng)估與驗(yàn)證

為了驗(yàn)證SonicMoE的整體效果，研究團(tuán)隊(duì)進(jìn)行了全面的性能測(cè)試。他們使用了從1.4B到120B參數(shù)的各種模型配置，在NVIDIA H100 GPU上進(jìn)行了詳細(xì)的基準(zhǔn)測(cè)試。結(jié)果顯示，SonicMoE在各種配置下都能顯著超越現(xiàn)有的最先進(jìn)方法。

在實(shí)際應(yīng)用測(cè)試中，SonicMoE在64張H100 GPU上訓(xùn)練7B專家混合模型的速度達(dá)到每天2130億令牌，這個(gè)速度相當(dāng)于使用96張H100 GPU運(yùn)行傳統(tǒng)ScatterMoE方法的每天2250億令牌。換句話說，SonicMoE用更少的硬件資源實(shí)現(xiàn)了幾乎相同的訓(xùn)練速度，硬件利用效率提升了約50%。

更令人印象深刻的是，在極端稀疏的配置下（比如DeepSeek-V3.2-Exp這樣的685B參數(shù)模型），SonicMoE是唯一能夠在單張H100 GPU上成功運(yùn)行的方法，其他基準(zhǔn)方法都因?yàn)閮?nèi)存不足或其他限制而無法工作。

研究團(tuán)隊(duì)還測(cè)試了SonicMoE在不同硬件配置下的表現(xiàn)。除了H100，他們還驗(yàn)證了在最新的Blackwell架構(gòu)GPU上的性能。SonicMoE充分利用了新硬件的特性，比如張量?jī)?nèi)存（TMEM）和統(tǒng)一矩陣乘法累加（UMMA）指令，在新硬件上表現(xiàn)出更好的性能。

五、對(duì)業(yè)界的深遠(yuǎn)影響

這項(xiàng)研究的意義遠(yuǎn)遠(yuǎn)超出了技術(shù)優(yōu)化本身。隨著AI模型規(guī)模的不斷增長(zhǎng)，訓(xùn)練效率已經(jīng)成為制約技術(shù)發(fā)展的關(guān)鍵瓶頸。SonicMoE提供的解決方案不僅能夠降低訓(xùn)練成本，還能讓研究機(jī)構(gòu)用更少的資源訓(xùn)練出更強(qiáng)大的模型。

從行業(yè)發(fā)展角度看，這種效率提升對(duì)于推動(dòng)AI技術(shù)的普及具有重要意義。降低訓(xùn)練成本意味著更多的研究團(tuán)隊(duì)和公司能夠參與到大規(guī)模AI模型的開發(fā)中，這將加速整個(gè)領(lǐng)域的創(chuàng)新步伐。同時(shí)，更高效的訓(xùn)練方法也減少了能源消耗，符合當(dāng)前對(duì)環(huán)境友好技術(shù)的需求。

SonicMoE的另一個(gè)重要貢獻(xiàn)是它的開源特性。研究團(tuán)隊(duì)將所有核心組件以開放許可證發(fā)布，這意味著整個(gè)AI社區(qū)都能從這些優(yōu)化中受益。這種開放合作的精神對(duì)于推動(dòng)技術(shù)進(jìn)步至關(guān)重要，也體現(xiàn)了學(xué)術(shù)研究回饋社會(huì)的價(jià)值。

展望未來，SonicMoE的設(shè)計(jì)理念和優(yōu)化策略很可能會(huì)被廣泛采用，成為訓(xùn)練大規(guī)模專家混合模型的標(biāo)準(zhǔn)方法。研究團(tuán)隊(duì)也指出了一些未來的發(fā)展方向，包括支持更多的數(shù)值精度格式（如FP8、MXFP8）以及在分布式訓(xùn)練環(huán)境中進(jìn)一步優(yōu)化通信與計(jì)算的重疊。

說到底，這項(xiàng)研究展示了在AI快速發(fā)展的今天，系統(tǒng)優(yōu)化和算法創(chuàng)新同樣重要。SonicMoE不是通過改變模型結(jié)構(gòu)或訓(xùn)練算法來提升性能，而是通過深度理解硬件特性和精心設(shè)計(jì)計(jì)算流程，在不損失任何功能的前提下大幅提升了效率。這種"軟硬結(jié)合"的優(yōu)化思路為未來的AI系統(tǒng)設(shè)計(jì)提供了重要參考，也證明了在追求更強(qiáng)大AI能力的同時(shí)，我們同樣需要關(guān)注如何更好地利用現(xiàn)有資源。

這項(xiàng)研究的成功還說明了跨領(lǐng)域合作的價(jià)值。普林斯頓大學(xué)、加州大學(xué)伯克利分校和Together AI的聯(lián)合團(tuán)隊(duì)結(jié)合了理論研究、系統(tǒng)設(shè)計(jì)和工程實(shí)現(xiàn)的專長(zhǎng)，這種多元化的合作模式正是解決復(fù)雜技術(shù)挑戰(zhàn)所需要的。對(duì)于有志于AI研究的讀者來說，這也提示了掌握多方面技能的重要性，從算法設(shè)計(jì)到系統(tǒng)優(yōu)化，每個(gè)環(huán)節(jié)都可能成為突破的關(guān)鍵點(diǎn)。

Q&A

Q1：SonicMoE是什么？

A：SonicMoE是普林斯頓大學(xué)等機(jī)構(gòu)聯(lián)合開發(fā)的專家混合模型訓(xùn)練優(yōu)化技術(shù)，主要解決精細(xì)化和稀疏化MoE模型訓(xùn)練中的內(nèi)存占用過大和計(jì)算效率低下問題，能夠?qū)⒂?xùn)練速度提升近一倍。

Q2：令牌舍入路由方法如何提升訓(xùn)練效率？

A：令牌舍入路由通過調(diào)整每個(gè)專家接收的令牌數(shù)量到GPU瓦片大小的整數(shù)倍，消除了計(jì)算中的填充浪費(fèi)。這種方法在保持模型性能不變的前提下，能夠額外提升16%的訓(xùn)練速度。

Q3：為什么專家混合模型訓(xùn)練會(huì)遇到內(nèi)存瓶頸？

A：隨著專家變得更加專業(yè)化，模型需要保存更多中間結(jié)果用于后續(xù)計(jì)算，內(nèi)存使用量呈線性增長(zhǎng)。SonicMoE通過重新設(shè)計(jì)計(jì)算路徑，避免保存某些中間結(jié)果，將內(nèi)存使用量減少45%且保持恒定。