從Transformer誕生起前饋網(wǎng)絡(luò)FFN就是又笨重又繞不開(kāi)的角色它占了模型近三分之二的參數(shù)量但每個(gè)token都得老老實(shí)實(shí)把所有參數(shù)過(guò)一遍。你擴(kuò)大模型容量推理成本就跟著線性往上漲這讓很多想做大規(guī)模模型的人非常頭疼。MoE混合專家模型或者說(shuō)稀疏前饋網(wǎng)絡(luò)正是沖著這個(gè)矛盾來(lái)的——它把單個(gè)FFN拆成一組并行專家再加一道輕量門控路由讓每個(gè)token只激活其中一小部分專家。參數(shù)總量可以擴(kuò)展數(shù)倍到數(shù)十倍但單token的計(jì)算量基本維持原樣容量和成本終于被解耦了。這篇文章我會(huì)從設(shè)計(jì)動(dòng)機(jī)一路講到訓(xùn)練和推理階段的工程坑包括門控網(wǎng)絡(luò)怎么工作、負(fù)載均衡為什么難搞、專家并行怎么通信以及一批可以直接上手的調(diào)參經(jīng)驗(yàn)。不管你是想改造自己的Transformer還是想理解最近一系列MoE大模型的結(jié)構(gòu)這篇都值得認(rèn)真看完。1. 從 FFN 到 MoE改寫前饋網(wǎng)絡(luò)的真實(shí)動(dòng)機(jī)1.1 一個(gè)容易被忽視的事實(shí)Transformer 的容量主要藏在 FFN 里很多人一提到Transformer下意識(shí)關(guān)注的是自注意力Self-Attention——tokens之間的交互全靠它。這話沒(méi)錯(cuò)但你要是統(tǒng)計(jì)一下參數(shù)量注意力層其實(shí)只占很小一部分。一個(gè)標(biāo)準(zhǔn)Transformer Block里FFN通常包含兩個(gè)線性變換先把隱藏維度比如4096升到中間維度比如11008再降回來(lái)。這兩個(gè)矩陣相加參數(shù)量輕輕松松占總模型的2/3甚至更多。為什么設(shè)計(jì)成這個(gè)樣子因?yàn)镕FN負(fù)責(zé)的是逐位置的記憶與抽象每個(gè)token獨(dú)立經(jīng)過(guò)同樣的非線性變換把注意力層已經(jīng)匯集的上下文信息進(jìn)一步加工成更高層的特征?？梢赃@樣理解注意力層負(fù)責(zé)商量FFN負(fù)責(zé)各自做事。商量只需要一個(gè)桌子做事卻需要大量工具工具多了參數(shù)量自然就上去了。問(wèn)題在于推理時(shí)即使工具再多每個(gè)token也得把這些工具全部用一遍。參數(shù)量翻倍單個(gè)前向的計(jì)算量就翻倍延遲也跟著上來(lái)。于是你面臨一個(gè)非此即彼的選擇要么保持計(jì)算量不變放棄擴(kuò)展容量要么接受昂貴的推理成本。這個(gè)矛盾在模型規(guī)模做大之后尤其尖銳我當(dāng)時(shí)在一個(gè)文本模型項(xiàng)目里試著把FFN中間維度從4倍擴(kuò)到8倍效果確實(shí)有提升但顯存和推理延遲的代價(jià)直接讓我放棄了。1.2 稀疏激活的直覺(jué)讓參數(shù)多、計(jì)算少變成現(xiàn)實(shí)MoE的思路一句話就能講清楚既然每個(gè)token用不上所有工具那就把FFN拆成若干個(gè)可以獨(dú)立工作的專家然后裝一個(gè)分診臺(tái)門控網(wǎng)絡(luò)讓每個(gè)token只去它需要的幾個(gè)專家那里其余專家本輪直接跳過(guò)。這就是稀疏激活的含義。傳統(tǒng)FFN是稠密計(jì)算——所有參數(shù)對(duì)所有輸入都參與計(jì)算MoE則是稀疏的——參數(shù)總量很大但單個(gè)token實(shí)際觸發(fā)的參數(shù)很少。用專業(yè)一點(diǎn)的術(shù)語(yǔ)描述假設(shè)總共有E個(gè)專家每個(gè)token從路由分布里挑出K個(gè)專家K遠(yuǎn)小于E通常K取1或2那么這token的FFN計(jì)算量大約是原來(lái)單個(gè)FFN的K/E而不是E倍。直覺(jué)上可以類比全科醫(yī)生和?？崎T診的區(qū)別。普通FFN像一個(gè)能力全面的全科醫(yī)生每個(gè)病人都由他一個(gè)人看診所請(qǐng)不動(dòng)更多醫(yī)生病人再多也是他一個(gè)人忙。MoE則是一個(gè)開(kāi)了很多?？圃\室的醫(yī)院每個(gè)專家只擅長(zhǎng)某一類病分診臺(tái)根據(jù)病人癥狀快速分配科室每個(gè)病人只看一兩個(gè)專家其他科室完全不打擾。醫(yī)院規(guī)模可以很大但單個(gè)病人的看病時(shí)間并沒(méi)有增加多少。模型容量來(lái)自有很多專科醫(yī)生計(jì)算效率來(lái)自每個(gè)病人只看很少的醫(yī)生。1.3 為什么改造的是 FFN而不是注意力層這里有個(gè)值得深思的問(wèn)題為什么大家做MoE都動(dòng)FFN而不直接把注意力層也拆成混合專家原因在于兩種結(jié)構(gòu)的Nature完全不同。注意力層在每層同時(shí)處理所有token通過(guò)Q、K、V矩陣建立跨token的交互每個(gè)token的結(jié)果都與其他token相關(guān)FFN則對(duì)每個(gè)token獨(dú)立操作token之間沒(méi)有任何信息交換。正因?yàn)镕FN的逐token獨(dú)立性我們才可以放心地把不同token送到不同專家專家之間不通信各自處理完再把結(jié)果放回原來(lái)的位置。如果對(duì)注意力層做路由比如讓每個(gè)token只和一部分token交互那被誰(shuí)注意到會(huì)因?yàn)椴煌瑃oken選了不同專家而變得極度碎片化上下文信息容易出現(xiàn)斷裂訓(xùn)練難度也會(huì)直線上升。這不是說(shuō)注意力MoE完全不可行而是在標(biāo)準(zhǔn)Transformer架構(gòu)里FFN是成本收益最劃算的改造點(diǎn)改動(dòng)局部、邏輯清晰、收益最大。還有一點(diǎn)很實(shí)際FFN的權(quán)重是逐token獨(dú)立存儲(chǔ)的天然適合做參數(shù)并行——把不同專家放到不同GPU上互不干擾。注意力層則必須做完整的矩陣乘法和跨卡同步通信開(kāi)銷大得多。所以從工程角度FFN也遠(yuǎn)比注意力層更適合做稀疏化改造。2. 稀疏路由的完整工作鏈路門控網(wǎng)絡(luò)、Top-K 選擇與專家容量2.1 門控網(wǎng)絡(luò)token 怎么決定自己去哪個(gè)專家門控網(wǎng)絡(luò)Router也叫Gating Network是整個(gè)MoE最核心的組件。別看它決定著一大群專家的調(diào)用方式它本身的結(jié)構(gòu)卻非常簡(jiǎn)單通常就是一個(gè)線性層輸入是token的隱藏表示h輸出維度等于專家總數(shù)E。然后過(guò)一個(gè)Softmax得到該token對(duì)各專家的選擇概率分布。用公式表示就是router_logits W_r hrouter_probs softmax(router_logits)這里的W_r是一個(gè)非常小的矩陣假設(shè)隱藏維度是4096專家數(shù)是8那W_r的參數(shù)量只有4096 × 8占整個(gè)模型的比重可以忽略不計(jì)。所以門控網(wǎng)絡(luò)被稱為輕量——它算得快、參數(shù)少卻直接影響模型的容量利用率。一個(gè)關(guān)鍵的設(shè)計(jì)細(xì)節(jié)是Router的輸入到底用哪一層表示。如果是用當(dāng)前層的FFN輸入那路由決策反映的是這一層之前的特征如果是用上一層的輸出或者與attention輸出拼接后的向量則可能包含更多上下文。不同實(shí)現(xiàn)各有取舍我個(gè)人的粗淺經(jīng)驗(yàn)是在層數(shù)較深的地方用本層的輸入就夠如果訓(xùn)練不穩(wěn)定可以嘗試把上一層Router的logits也拼接喂進(jìn)來(lái)給當(dāng)前層更多先驗(yàn)。2.2 Top-K 選擇與專家容量從分診到預(yù)約掛號(hào)拿到全量專家概率后不能真的所有專家都參與計(jì)算那樣又退回稠密了。做法是取概率最高的K個(gè)專家K通常取1或2。取到具體哪K個(gè)專家后對(duì)應(yīng)專家計(jì)算得到輸出再按路由概率或歸一化后的權(quán)重加權(quán)求和得到當(dāng)前token的FFN輸出。當(dāng)K1時(shí)每個(gè)token只走一個(gè)專家計(jì)算效率最高但每個(gè)token的感受域非常受限專家之間也完全沒(méi)有冗余協(xié)作。K2是當(dāng)前很多工作的經(jīng)驗(yàn)?zāi)J(rèn)值既保證了稀疏性又給了token一點(diǎn)選擇的余地。我感覺(jué)K2相當(dāng)于讓分診臺(tái)不僅選最合適的科室還允許備選科室兜底對(duì)訓(xùn)練初期路由還不成熟時(shí)特別友好。接下來(lái)是專家容量Expert Capacity的問(wèn)題。一個(gè)專家在某一層最多能處理多少個(gè)token如果不設(shè)上限某幾個(gè)熱門專家可能被分配大量token其他專家空轉(zhuǎn)計(jì)算負(fù)載嚴(yán)重不均。于是實(shí)現(xiàn)時(shí)會(huì)為每個(gè)專家設(shè)置一個(gè)容量上限capacity capacity_factor * (T / E)。這個(gè)公式的含義是理論平均分配下每個(gè)專家應(yīng)該處理T/E個(gè)tokencapacity_factor則允許一定的彈性冗余。舉個(gè)例子假設(shè)T256個(gè)tokenE8個(gè)專家理論平均是每個(gè)專家32個(gè)tokencapacity_factor取1.25則容量為40。一旦某個(gè)專家被分配的token數(shù)超過(guò)了容量超出的token就只能被丟棄或者走一條殘差連接直接跳過(guò)FFN的通路。很多人第一次看到丟棄token會(huì)嚇一跳訓(xùn)練時(shí)丟信息模型還能學(xué)好嗎實(shí)際上這是工程上的妥協(xié)——如果強(qiáng)行保留所有token計(jì)算圖和同步邏輯會(huì)變得非常復(fù)雜還會(huì)導(dǎo)致不同專家之間的執(zhí)行時(shí)間出現(xiàn)長(zhǎng)尾。后續(xù)有人提出無(wú)丟棄drop-free的負(fù)載均衡方法核心是調(diào)整路由分配而不是限制容量效果也確實(shí)更好。但從工程復(fù)雜度考慮帶容量的實(shí)現(xiàn)依然是目前最穩(wěn)妥的基線方案。2.3 一個(gè)完整的 token 走向示例為了把上面這些概念串起來(lái)我走一遍一個(gè)token在MoE層的完整流程token輸入為形狀[B, H]的向量先復(fù)制一份作為殘差為后續(xù)相加做準(zhǔn)備。門控網(wǎng)絡(luò)根據(jù)當(dāng)前token的表示算出每個(gè)專家的概率得到形狀[E]的分布。挑選概率最高的K2個(gè)專家記為專家3和專家7。token被送往專家3和專家7所在的設(shè)備如果是單卡就是同一設(shè)備上的不同矩陣。專家3和專家7分別對(duì)token做自己的FFN前向計(jì)算得到兩個(gè)輸出。兩個(gè)輸出按路由權(quán)重加權(quán)求和。加權(quán)結(jié)果與第一步的殘差相加得到當(dāng)前層的輸出。這個(gè)流程里其他6個(gè)專家完全沒(méi)有參與計(jì)算對(duì)應(yīng)的參數(shù)也處于空閑狀態(tài)。這就是稀疏兩個(gè)字的真正含義——不是少設(shè)計(jì)了參數(shù)是每次只用到其中一小撮。3. 負(fù)載均衡與路由離散性訓(xùn)練階段最頭疼的兩個(gè)工程問(wèn)題3.1 離散路由不可微梯度怎么回到 RouterMoE聽(tīng)著完美但訓(xùn)練時(shí)第一個(gè)坑就擺在面前Top-K取的是索引是一個(gè)離散操作不可導(dǎo)。你沒(méi)法把我選了專家3這個(gè)行為直接變成梯度信號(hào)來(lái)更新門控網(wǎng)絡(luò)。怎么辦經(jīng)典解法是雖然選擇是離散的但最終輸出是對(duì)所選專家的輸出按概率加權(quán)求和而概率本身是router logits經(jīng)softmax得到的是連續(xù)可導(dǎo)的。于是梯度可以沿著概率權(quán)重這條路回傳到路由網(wǎng)絡(luò)。換句話說(shuō)真正參與加權(quán)求和的權(quán)重是router_probs里選中那個(gè)位置的值而不是一個(gè)one-hot向量即使如此這仍然會(huì)產(chǎn)生一個(gè)什么問(wèn)題呢問(wèn)題在于概率權(quán)重在反向傳播時(shí)數(shù)值通常比較小門控網(wǎng)絡(luò)拿到的梯度信號(hào)偏弱。你可以理解為后端專家們學(xué)得很快但分診臺(tái)學(xué)得有點(diǎn)滯后。這帶來(lái)一系列連鎖問(wèn)題后面會(huì)詳細(xì)講。還有一個(gè)細(xì)節(jié)容易被忽略訓(xùn)練時(shí)如果想用hard gradient估計(jì)算法類似Gumbel-Softmax的Straight-Through情況會(huì)稍有不同——forward時(shí)用one-hotbackward時(shí)用連續(xù)概率的近似梯度。但標(biāo)準(zhǔn)實(shí)現(xiàn)里大多數(shù)框架仍然采用加權(quán)求和式的軟路由簡(jiǎn)單且穩(wěn)定。我之前在某次實(shí)驗(yàn)里試過(guò)hard路由收斂速度明顯變慢后來(lái)果斷退回soft加權(quán)。3.2 負(fù)載坍縮少數(shù)專家過(guò)熱、多數(shù)專家空轉(zhuǎn)MoE最著名的訓(xùn)練困境就是負(fù)載坍縮load collapse。開(kāi)始訓(xùn)練時(shí)Router的參數(shù)幾乎隨機(jī)各專家的分配還算均勻。但很快某些專家因?yàn)槌跏蓟院?、梯度更新方向略微有利就?huì)被分配更多的token形成正反饋熱門專家學(xué)得更快、表現(xiàn)更好Router越來(lái)越傾向于選它其他專家則長(zhǎng)期閑置梯度信號(hào)近乎為零逐漸變成死專家。最終整個(gè)模型退化成其實(shí)只用了兩三個(gè)專家參數(shù)總量虛增容量?jī)?yōu)勢(shì)形同虛設(shè)。解決負(fù)載坍縮的主要手段是加輔助負(fù)載均衡損失Auxiliary Load Balancing Loss讓Router在每個(gè)batch里盡量把token均勻分配給專家。經(jīng)典的實(shí)現(xiàn)大概長(zhǎng)這樣aux_loss E * sum(prob_per_expert * fraction_of_tokens_per_expert)其中prob_per_expert是所有token對(duì)本專家的平均選擇概率fraction_of_tokens_per_expert是實(shí)際被分配到本專家的token占比。二者相乘再求和如果分配完全均勻這個(gè)值會(huì)趨近于1越不均勻值越大相當(dāng)于對(duì)Router施加一個(gè)均勻化的懲罰。E作為縮放因子使得損失數(shù)值與專家數(shù)無(wú)關(guān)方便統(tǒng)一調(diào)系數(shù)。但這又引入新矛盾輔助損失權(quán)重設(shè)太大模型會(huì)傾向于為了均勻而均勻犧牲真正的專業(yè)分化整體表現(xiàn)下降設(shè)太小負(fù)載坍縮壓不住。我在實(shí)際調(diào)參時(shí)觀察到的情況是由于每個(gè)batch的分布都在變固定權(quán)重很難同時(shí)照顧訓(xùn)練前期和后期的需求。比較務(wù)實(shí)的方法是訓(xùn)練前幾百步用小權(quán)重等專家分化初步形成后再逐漸加大或者觀察負(fù)載的variance指標(biāo)——如果某專家占比長(zhǎng)期超過(guò)2倍平均值就臨時(shí)加大權(quán)重拉一把。3.3 初始化與 Router 輸入尺度的穩(wěn)定性除了負(fù)載均衡Router訓(xùn)練的另一個(gè)隱性障礙是logits的數(shù)值波動(dòng)。如果Router的輸出一直非常大softmax出來(lái)的分布就會(huì)非常尖銳幾乎等于固定選Top-1專家別的專家?guī)缀鯖](méi)有機(jī)會(huì)。反過(guò)來(lái)如果logits太小分布太平滑Top-K的選中結(jié)果和未選中結(jié)果差異不大路由失去了區(qū)分能力。一個(gè)常用技巧是給Router加一個(gè)可學(xué)習(xí)的temperature項(xiàng)初始值設(shè)小一點(diǎn)讓logits先溫和一些。入口處做LayerNorm也是常見(jiàn)做法把不同token的表示尺度拉到同一水平避免某些高范數(shù)token支配Router的決策。我還用過(guò)一種更樸素的方法直接從已有Dense模型做熱啟動(dòng)warm start。比如把稠密FFN拆成E個(gè)專家時(shí)克隆同一份參數(shù)作為每個(gè)專家的初始化這樣所有專家在訓(xùn)練初期能力一致Router沒(méi)有偏心的理由。這在多任務(wù)遷移場(chǎng)景下特別有效畢竟從零訓(xùn)練一個(gè)大規(guī)模MoE冷啟動(dòng)穩(wěn)定性的難度是很多人低估的。4. 訓(xùn)練態(tài)與推理態(tài)的差異顯存預(yù)分配、專家并行與通信開(kāi)銷4.1 訓(xùn)練時(shí)怎么擺放專家專家并行與 All-to-All 通信討論完訓(xùn)練穩(wěn)定性還得解決一個(gè)現(xiàn)實(shí)問(wèn)題MoE參數(shù)總量大一張卡根本放不下。傳統(tǒng)模型并行是切分每一層的矩陣MoE則有天然更優(yōu)的方案——專家并行Expert Parallelism把不同專家分布在不同GPU上每張卡只保存一部分專家的權(quán)重Router所在位置則根據(jù)token的分布動(dòng)態(tài)將token散播到對(duì)應(yīng)專家所在的卡。這里有兩個(gè)關(guān)鍵通信階段。第一個(gè)階段叫Dispatch某個(gè)token被路由到專家3而專家3在另一張卡上就必須把token的隱藏向量通過(guò)通信送到那張卡。第二個(gè)階段叫Combine專家3算完輸出后結(jié)果還要送回原來(lái)的卡供后續(xù)的殘差相加和下一層使用。如果K2token實(shí)際上要做兩次dispatch/combine通信量也隨之倍增。這兩個(gè)通信階段使用的通常是All-to-All模式也就是每個(gè)節(jié)點(diǎn)都要向其他所有節(jié)點(diǎn)發(fā)送數(shù)據(jù)。這和傳統(tǒng)All-Reduce所有節(jié)點(diǎn)只匯總一個(gè)結(jié)果完全不同。通信量大約正比于被路由的token數(shù)、隱藏維度和專家跨卡跳數(shù)。所以一個(gè)非常重要的工程教訓(xùn)是MoE的稀疏收益到一定規(guī)模后可能被通信開(kāi)銷吃掉。如果網(wǎng)絡(luò)帶寬不夠你會(huì)發(fā)現(xiàn)訓(xùn)練速度不升反降加專家還不如加計(jì)算密度來(lái)得劃算。我實(shí)際測(cè)試過(guò)一個(gè)典型配置8卡環(huán)境、每層8個(gè)專家、K2。在單機(jī)NVLink互聯(lián)下通信損耗還能接受但一旦跨機(jī)走以太網(wǎng)訓(xùn)練step時(shí)間立刻漲了30%收益大打折扣。所以設(shè)計(jì)MoE規(guī)模時(shí)一定要先盤清楚自己的通信拓?fù)洹?.2 推理態(tài)顯存的真正瓶頸參數(shù)常駐而不是計(jì)算量很多人第一次接觸MoE推理時(shí)都會(huì)有個(gè)錯(cuò)覺(jué)既然每個(gè)token只激活一小部分專家那顯存占用是不是也按比例減小不是。稀疏只減少了計(jì)算量FLOPs參數(shù)依然全部常駐顯存。無(wú)非是一個(gè)token實(shí)際只用到了門控網(wǎng)絡(luò)和少數(shù)專家的權(quán)重但所有專家的參數(shù)都必須在內(nèi)存里待命因?yàn)槟悴恢老乱粋€(gè)token會(huì)訪問(wèn)哪個(gè)專家。這在推理側(cè)會(huì)造成一個(gè)非常有趣的局面MoE模型的推理延遲可以接近同計(jì)算量的Dense模型但顯存占用接近同參數(shù)量的Dense模型。換句話說(shuō)你省的是算力省不了容量。正因如此大規(guī)模MoE推理基本都是多卡分布式方案每張卡放一部分專家。當(dāng)請(qǐng)求到來(lái)時(shí)token先經(jīng)過(guò)路由再通過(guò)通信調(diào)用其他卡的專家。這實(shí)際上把訓(xùn)練時(shí)的專家并行復(fù)用了只是batch更小、通信更多在卡間進(jìn)行。顯存優(yōu)化通常還有兩個(gè)方向一是對(duì)專家權(quán)重做更激進(jìn)的低比特量化因?yàn)閷＜冶旧硎侵餿oken獨(dú)立使用的量化誤差不容易像attention參數(shù)那樣跨token擴(kuò)散二是預(yù)先把所有專家的參數(shù)做統(tǒng)一內(nèi)存池管理配合顯存碎片整理。后一種技巧偏工程化但收益很直觀尤其在長(zhǎng)期多請(qǐng)求服務(wù)的場(chǎng)景里。4.3 padding token 帶來(lái)的路由損耗一個(gè)容易踩的隱藏成本訓(xùn)練時(shí)為了batch對(duì)齊我們會(huì)給短句做padding。在普通Transformer里padding位置雖然參與計(jì)算但因?yàn)樽⒁饬ask的存在它們對(duì)有效位置的影響基本可以被忽略于是人們很少專門去處理padding token的成本。但在MoE里情況變了每個(gè)padding token也要走Router、也要被dispatch到某個(gè)專家、也要占用專家容量。這帶來(lái)兩個(gè)直接后果。第一padding token占用專家容量可能導(dǎo)致有效token被擠掉超過(guò)capacity被丟棄。第二即使padding token沒(méi)有實(shí)際信息它們依然占據(jù)通信通道白白消耗All-to-All帶寬。好在我們找到的解決辦法并不復(fù)雜在padding位置的路由概率上直接加一個(gè)非常大的負(fù)偏置讓Router幾乎不可能選中任何專家或者把padding token在進(jìn)入MoE層之前直接從batch里摘除等所有專家計(jì)算完再放回原位。前者改動(dòng)小后者更徹底但需要處理動(dòng)態(tài)batch工程復(fù)雜度會(huì)增加一些。類似的浪費(fèi)也出現(xiàn)在推理場(chǎng)景現(xiàn)在主流的推理框架普遍采用連續(xù)批處理continuous batching不同請(qǐng)求的長(zhǎng)度不對(duì)齊時(shí)也會(huì)出現(xiàn)大量虛位token。所以做MoE推理服務(wù)時(shí)請(qǐng)求調(diào)度策略對(duì)性能的影響往往比模型結(jié)構(gòu)本身還大。5. 實(shí)測(cè)調(diào)參清單與擴(kuò)展思路從復(fù)現(xiàn)一個(gè)簡(jiǎn)單 MoE 層開(kāi)始5.1 一個(gè)可直接改寫的 MoE 層偽代碼結(jié)構(gòu)看完前面的原理你已經(jīng)可以把一個(gè)標(biāo)準(zhǔn)的MoE層搭出來(lái)了。下面用離實(shí)現(xiàn)很接近的偽代碼展示核心邏輯框架無(wú)關(guān)方便你對(duì)照著改寫自己的代碼。def moe_ffn(x, experts, router, K2, capacity_factor1.25): # x: [batch, hidden] B x.shape[0] E len(experts) T B # 這里假設(shè)每個(gè)token獨(dú)立路由實(shí)際實(shí)現(xiàn)會(huì)考慮padding mask # 1. 計(jì)算路由概率 logits router(x) # [B, E] probs softmax(logits, dim-1) # [B, E] # 2. 取Top-K專家 top_probs, top_idx probs.topk(K, dim-1) # 各[B, K] top_probs top_probs / top_probs.sum(dim-1, keepdimTrue) # 歸一化 # 3. 按專家容量篩選 capacity int(capacity_factor * B / E) dispatch_mask build_dispatch_mask(top_idx, E, capacity) # [B, E, K] 或管理token分配 # 4. 對(duì)每個(gè)選中的專家執(zhí)行 FFN并按權(quán)重聚合 output zeros_like(x) for e in range(E): selected dispatch_mask[:, e] # 哪些token被分到專家e if selected.any(): expert_out experts[e](x[selected]) # [num_selected, hidden] output[selected] (top_probs[selected, e] * expert_out) # 5. 殘差連接 return output residual_x當(dāng)然真實(shí)工程里不會(huì)用for循環(huán)逐專家算而是通過(guò)矩陣乘法一次性完成所有專家的前向但核心邏輯就是上面的流程。為了效率專家計(jì)算通常會(huì)合并成一個(gè)大的張量運(yùn)算把專家參數(shù)沿著一個(gè)軸堆疊起來(lái)然后用grouped GEMM一次算出所有選中專家的輸出。這也是很多訓(xùn)練框架在MoE部分做的主要優(yōu)化。5.2 調(diào)參順序先穩(wěn)定路由再追容量調(diào)MoE和調(diào)普通模型一個(gè)很大的區(qū)別是超參數(shù)之間存在明顯的先后依賴關(guān)系。一上來(lái)就亂試專家數(shù)、路由權(quán)重很容易出現(xiàn)忘記調(diào)了什么的局面因?yàn)槊總€(gè)超參數(shù)都會(huì)間接影響Router的負(fù)載分布。我建議按下面的順序做超參數(shù)常用范圍影響效果調(diào)整優(yōu)先級(jí)專家數(shù) E8~64視參數(shù)量預(yù)算決定容量上限和并行度先定預(yù)算再選Top-K1~2計(jì)算量、冗余度、通信量建議固定為2capacity_factor1.0~1.5丟棄率、負(fù)載余量、性能優(yōu)先調(diào)aux loss 系數(shù)0.001~0.1負(fù)載均衡強(qiáng)度次優(yōu)先Router溫度0.5~1.0可學(xué)習(xí)路由分布尖銳度小步微調(diào)專家Dropout0.0~0.1緩解過(guò)擬合、增強(qiáng)魯棒性最后調(diào)具體來(lái)說(shuō)先固定E和K把capacity_factor從1.25開(kāi)始試觀察訓(xùn)練時(shí)被丟棄token的比例。如果丟棄太多說(shuō)明Router的選擇太集中或容量太小可以適當(dāng)調(diào)大capacity_factor如果幾乎為零則逐步減小壓出更多計(jì)算密度。然后才去調(diào)整aux loss系數(shù)從0.01起步每幾百步看一眼各專家的token占比分布。占比的方差如果持續(xù)縮小到接近均勻說(shuō)明權(quán)重偏大可以降如果某個(gè)專家占比持續(xù)超過(guò)2倍平均則要加大權(quán)重。一個(gè)反直覺(jué)的經(jīng)驗(yàn)是對(duì)于中小規(guī)模的模型參數(shù)量在億級(jí)以下MoE的優(yōu)勢(shì)并不明顯。Router自身也需要數(shù)據(jù)爬坡假設(shè)總數(shù)據(jù)量沒(méi)法讓每個(gè)專家都充分訓(xùn)練那稀疏結(jié)構(gòu)的收益就會(huì)打折扣。所以不要只盯著MoE大模型專利這個(gè)說(shuō)法也要理解它依賴的規(guī)模前提。5.3 什么時(shí)候不應(yīng)該用 MoE以及幾個(gè)值得關(guān)注的擴(kuò)展方向MoE不是銀彈。如果你的使用場(chǎng)景是單機(jī)單卡、推理延遲極敏感、或者模型規(guī)模本來(lái)就很小那么MoE帶來(lái)的通信開(kāi)銷和顯存負(fù)擔(dān)通常不值得。尤其是邊緣端或者移動(dòng)端場(chǎng)景參數(shù)量常駐和動(dòng)態(tài)路由的不確定性會(huì)讓部署變得非常不友好。寧可把精力花在蒸餾、量化或者知識(shí)蒸餾這類壓縮手段上。反過(guò)來(lái)MoE真正發(fā)光的地方是訓(xùn)練語(yǔ)料極其豐富而單卡顯存成為瓶頸、需要在不顯著增加推理計(jì)算量的前提下擴(kuò)展容量、以及多種異構(gòu)任務(wù)希望共用一套底座但各任務(wù)又有自己的特性。多模態(tài)方向也可以大膽嘗試——不同模態(tài)的信息處理模式差異大正好適合交由不同專家專人負(fù)責(zé)門控網(wǎng)絡(luò)學(xué)到的模態(tài)路由有時(shí)會(huì)帶來(lái)意想不到的可解釋性。我在一次改造中把MoE從語(yǔ)言模型搬到了推薦系統(tǒng)的打分模型里不同用戶群體走不同專家冷啟動(dòng)階段用戶走Generalist專家活躍用戶則能訪問(wèn)專屬專家后的組合結(jié)果效果提升比純粹加深網(wǎng)絡(luò)更明顯。這說(shuō)明MoE的分診思想完全值得跨場(chǎng)景復(fù)用關(guān)鍵還是找到那些不同樣本需要不同能力的問(wèn)題。最后分享兩個(gè)小技巧。如果你已經(jīng)用上了MoE建議把Router的決策分布做成可視化日志每訓(xùn)練幾千步看一眼比任何損失曲線都能更快發(fā)現(xiàn)問(wèn)題另一個(gè)是在測(cè)試階段記錄專家熱力分布如果模型收斂后大部分token的Top-1專家高度集中往往說(shuō)明模型只學(xué)到了表面規(guī)律需要回頭檢查訓(xùn)練數(shù)據(jù)本身的多樣性。MoE這條路技術(shù)在往前走但工程的坑說(shuō)到底還是那幾類通信、負(fù)載、穩(wěn)定性。理解了底層邏輯剩下的就是耐心調(diào)參和踩坑。希望這些經(jīng)驗(yàn)?zāi)軒湍闵僮邘撞綇澛贰?