「DeepSeek-V3 技術解析」:無輔助損失函數的負載均衡 原創
編者按: 在混合專家模型(MoE)的實踐中,負載不均衡儼然已成為制約模型性能提升的關鍵瓶頸之一。傳統的均衡策略往往需要引入復雜的輔助損失函數,不僅增加了訓練的復雜度,還可能干擾模型的核心學習目標。工程師們在提升模型效率的道路上,一直苦苦追尋著一個優雅而高效的平衡解決方案。
DeepSeek 團隊的這項研究,為這一長期困擾業界的技術難題提供了令人耳目一新的解決思路:通過在門控分數中直接添加專家層面的偏置項,在絕大部分不引入額外損失函數的情況下,實現了模型訓練過程中的自適應負載均衡。更令人驚嘆的是,這一方法不僅保持了模型的因果關系,還顯著提升了訓練的穩定性和最終性能。
作者 | Shirley Li
編譯 | 岳揚
這是 DeepSeek-V3 系列文章的第三篇,我們將探討 DeepSeek 模型[1, 2, 3]中與混合專家模型(MoE)相關的另一項關鍵架構突破:無輔助損失函數的負載均衡(Auxiliary-Loss-Free Load Balancing)[5]。
在本文,我們將深入解析 DeepSeek 如何解決 MoE 的隱藏瓶頸——負載均衡,同時還通過消除梯度干擾和嚴格遵循因果關系約束,提升了訓練和推理效率,為后續專家模型的優化方向提供了標桿。
本文目錄:
- 技術背景:介紹混合專家模型(MoE)的基本原理,解釋負載均衡的重要性,回顧之前的工作,包括輔助損失函數(auxiliary loss)方法和專家選擇(Expert Choice)策略。
- DeepSeek 的無輔助損失函數的負載均衡:解析其運作原理
- 性能評估:討論無輔助損失函數的負載均衡的性能表現
- 總結
- 參考文獻
「DeepSeek-V3 技術解析」系列其他文章:
01 技術背景
1.1 MoE (Mixture-of-Experts) in Transformers
MoE(Mixture-of-Experts,混合專家模型)在 Transformer 中的實現方式通常為:每隔若干 Transformer 層,將其中的 FFN(前饋網絡)替換為多個 FFN(每個 FFN 充當一個"專家")。當 input token 進入該層時,通過門控操作(Gating)選擇 Top-K 個專家,并將 input token 只路由至這些被選中的 FFN,從而僅激活對應的專家網絡。
下圖展示了這一過程:左側標準 Transformer 層中的 FFN 子層被替換為右側的 MoE 層。
如需更詳細的 MoE 技術解析,可參考??「DeepSeek-V3 技術解析」:DeepSeekMoE??(文中通過餐廳類比直觀解釋了 MoE 原理)。
圖 1. Transformer 中的 MoE 層(紅框內)示意圖。圖片改編自文獻 [6]。
1.2 負載均衡及其重要性
在??本系列上一篇文章??進行的餐廳類比中,我們通過一個能夠提供多種菜系菜品的餐廳解釋了 MoE 的概念:餐廳的每位廚師都是專家,主廚(Head Chef)的工作類似于門控操作,將每道菜品分配給具備對應技能的特定廚師。
為確保該系統高效運作,需滿足以下條件:
- 每位專業廚師必須精通自身菜系所需技能(例如餃子廚師必須會包餃子),同時所有廚師需能共同處理所有菜品。
- 主廚需充分了解所有專業廚師的專長,并能高效分配訂單。
在 MoE 中,前者對應 expert specialization 與 knowledge sharing 的權衡(我們已在??介紹 DeepSeekMoE 的這篇文章中??進行了詳細討論),后者則體現了負載均衡的重要性 —— 這也是本文的核心主題。
為何負載均衡如此重要?
原因在于,當負載不均衡發生時,MoE 無法有效運作。 最常見的問題是路由崩潰(Route Collapse):少數專家接收大部分 input token,而其他專家的利用率極低。
因此,大部分計算由超負荷工作的專家承擔,而這些專家通常分布在多個 GPU 核心上,因此會導致硬件資源浪費。
由于梯度沖突(gradient conflict),路由奔潰也會導致訓練不穩定。超負荷工作的專家接收更多 input token,他們積累的梯度也會更大,學習速度也會比工作負荷不足的專家快得多;因此,兩者的梯度在幅值和方向上均可能發生偏離,導致訓練難以收斂。
最后,MoE 中的負載不均衡也會導致性能低下和泛化效果不佳,因為工作負荷不足的專家會因為訓練 tokens 不足,而難以學習有效知識。
由于負載均衡技術對 MoE 至關重要,因此研究者們針對這一問題提出了多種解決方案。其中,最常用的策略是為負載均衡添加輔助損失函數(Auxiliary Loss)和專家選擇(Expert Choice)。
1.3 帶輔助損失函數的負載均衡
一種常見的改善負載均衡的策略是在模型訓練的目標函數基礎上引入輔助損失函數。
圖 2. 用于強化負載均衡的輔助損失函數示例。圖片編輯自文獻[5]。
上圖展示了一個輔助損失函數的示例,其中
- N 是專家數量,T 是 token 數量,K 是每個 input token 激活的專家數量。
- s_{i, t} 是門控機制的輸出,通過 Softmax 歸一化到 [0, 1] 區間,表示第 t 個 token 選擇第 i 個專家的概率。向量 u_t 是第 t 個 token 的輸入隱藏狀態,而 e_i 是第 i 個專家的“質心”,可以看作歷史上路由到第 i 個專家的 token 嵌入平均值。因此,s_{i, t} 度量的是當前輸入與第 i 位專家歷史接收 token 的平均值的接近程度。
- 因此,P_i 可視為整個輸入序列選擇第 i 個專家的平均概率。
- f_i 表示被路由到第 i 個專家的 token 比例。
需要注意的是,f_i 是不可微分的,因此最小化上述損失函數實際上轉化為了最小化 s_{i, t}。同時由于 f_i 依賴于 s_{i, t},對 s_{i, t} 的調整也會影響 f_i,從而實現對各專家負載分配的調節。
然而,用這種輔助損失函數來均衡負載需要付出一定代價,因為其梯度可能會干擾語言建模目標(language modeling objective)的梯度,導致模型性能下降,在極端不平衡情況下(工作負荷過大的專家的 f_i 和 P_i 都變得極大時)尤其明顯。
因此,采用這種方法進行負載均衡需要謹慎設置輔助損失函數的權重。為更清晰地說明這一點,文獻[5]的作者進行了一個實驗,用不同 alpha 值訓練模型,結果如下圖所示,其中縱軸表示困惑度指標下的模型性能,橫軸表示 MaxVio(衡量負載不平衡程度的指標,MaxVio 值越高表示負載越不平衡,i 表示第 i 個專家):
圖 3. 輔助損失函數控制訓練中負載均衡與模型性能的權衡困境。圖片引自文獻[5]。
如圖所示,當 alpha 過小時(alpha=0),MaxVio 保持高位,說明輔助損失函數未能有效實現負載均衡目標。另一方面,當 alpha 過大時(alpha=0.01),模型最終會產生更高的困惑度。
綜上,輔助損失函數控制的負載均衡是把雙刃劍:若 alpha 未經仔細調校,可能損害模型性能。實際 LLM 訓練中,由于資源限制,alpha 的調校過程充滿挑戰,這進一步增加了優化難度。
上圖同時展示了本文提出的無損失函數方法在相同 Perplexity-MaxVio 坐標系下的表現,該方法同時實現了低困惑度和低 MaxVio,證明了無損失函數方法的有效性。
1.4 專家選擇(Expert Choice)策略
在此需提及的另一項前期工作是 Expert Choice [7],它提出了一種簡單高效的負載均衡方法,將路由策略從“token choice”切換為“expert choice”。
具體而言,MoE 路由中的門控分數通常通過對 affinity matrix(譯者注:二維矩陣,用于量化 input token 與各個專家之間的匹配程度。)應用 Softmax 計算得出,如圖 2 所示。傳統路由方法從 token 維度應用 Softmax 為每個 token 選擇專家,因此這些方法被稱為“token choice”。問題在于,該機制下我們無法控制每個專家接收的 token 數量,最終導致負載不均衡問題。
而 Expert Choice 方法則從專家維度應用 Softmax,為每個專家選擇被路由的 token。通過這種設計,每個專家接收的 token 數量能夠天然達到完美均衡,因此無需依賴輔助損失函數實現負載均衡。在文獻[7]中,這種方法同時展現出更優的模型性能和更快的訓練速度。
然而,Expert Choice 這種方法也存在局限 —— 未來 token 泄露問題。由于每個專家需要先查看所有 token 的路由分數才能決定處理哪些 token,這違反了因果關系(causality),在文本生成、機器翻譯等自回歸任務中可能引發嚴重問題。
02 DeepSeek 的無輔助損失函數的負載均衡
為在不引入 gradient inference(譯者注:此處或為作者筆誤?應當為“gradient interference(梯度干擾)”,多個損失函數(或多個優化目標)的梯度方向發生沖突。) 的情況下解決負載均衡問題,DeepSeek 提出了一種名為" Loss-Free Balancing"的新技術,通過直接調整門控分數 s_{i,t} 實現。
如前文所述,當我們最小化圖 2 所示的輔助損失函數時,最終會通過調整 s_{i,t} 來實現最小化 P_i。
因此,若能直接調整 s_{i,t},理論上應能達到與施加輔助損失函數相似的效果。
為此,我們在每個專家的門控分數上添加了專家層面的偏置項,如下圖所示。需注意 b_i 并不用于最終門控分數的計算(后文也將說明該偏置項也是不可微分的),而是用于 TopK 選擇專家時:
圖 4. 在門控分數中引入偏置項 b_i。圖片引自文獻[5]
該偏置項 b_i 的計算方式非常直觀,如下圖所示:首先獲取分配給各專家的 token 數量平均值及所有專家的理論全局均值,然后計算給各專家分配的 token 數量與理論全局均值的差值,偏置項由該差值(或誤差)的符號乘以固定更新率(fixed update rate)決定(該更新率為可調超參數)。后續章節我們將對該超參數的影響進行更多實驗分析。
圖 5. DeepSeek 無損失函數的負載均衡算法。圖片引自文獻[5]
現可通過下表總結不同負載均衡方法的優勢與局限:
圖 6. 不同負載均衡方法對比。圖片引自文獻[5]
圖 3 已展示該方法在模型性能與負載均衡間取得了更好的權衡,但仍有多方面需要驗證。下一章節我們將深入分析實驗結果。
03 Evaluation
有三個關鍵問題需要回答:
- DeepSeek 所提出的方法能否在性能和負載均衡之間實現更好的權衡?
- 圖 5 中更新率 u 有什么影響?
- 我們能否進一步優化偏置項更新規則(鑒于其如此之簡單)?
3.1 性能 vs. 負載均衡
為回答第一個問題,作者在 1B 和 3B 模型上進行了實驗,比較 loss-controlled 負載均衡和 loss-free 負載均衡的困惑度(Perplexity)和 MaxVio,結果如下圖所示:
圖 7. loss-controlled 負載均衡和 loss-free 負載均衡的對比。圖片來自文獻[5]。
上述結果與我們在圖 3 中看到的結果類似:所提方法同時實現了更低的困惑度和更低的MaxVio。
除了評估最終的 checkpoint 外,作者還展示了訓練過程中的 MaxVio 曲線,以便更全面地理解該方法在整個訓練過程中的表現,如下圖所示:
圖 8. 訓練過程中的 MaxVio 曲線。圖片來自文獻[5]。
如圖中所示,在 1B 和 3B 的模型配置下,loss-free 方法在整個訓練過程中都展現出更優的負載均衡能力,體現了該方法的穩定性。
3.2 超參數的影響(更新率)
如圖 5 所示,所提方法引入了一個新的超參數 u(稱為更新率(update rate)),該超參數如何影響 loss-free 方法的有效性?具體而言,我們需要理解 loss-free 方法對超參數 u 的取值是敏感還是不敏感,以及如何選擇一個最優值來最大化該方法的效果。
如前文所述,在門控分數中添加偏置項的概念類似于繞過損失函數的反向傳播,直接對門控分數進行梯度更新。在這種情況下,更新率 u 的作用與梯度更新中的步長(step size)類似。由此可推測其影響也相似:過小的更新率可能會導致收斂速度緩慢。過大的更新率可能導致不穩定和引發波動。
在原文中,作者對更新率進行了實驗測試(取值從 1e-4 到 1e-2),結果如下圖所示:
圖 9. 更新率(update rate)對訓練負載均衡的影響。圖片來自文獻[5]。
與預期一致,當 u 過小時(如 1e-4),MaxVio 下降速度較慢;而過大的 u(如 1e-2)則因波動性增強,導致訓練過程中 MaxVio 持續偏高。
3.3 其他偏置項更新規則
為回答第三個問題,研究者嘗試了多種備選策略,并將它們與 DeepSeek 提出的版本進行對比:
- 策略變體 1:使用 e_i 的數值(而不僅僅是符號)計算偏置項,即從 b_i = b_i +u?sign(e_i) 改為b_i = b_i +u?e_i。
- 策略變體 2:使用乘法偏置項而非加法偏置項。
其中策略變體 2 可以更正式地描述如下:
實驗表明,策略變體 1 能帶來略優的負載均衡效果,但未提升模型性能:
圖 10. 策略變體 1 的性能表現。圖片來自文獻[5]。
而策略變體 2 甚至顯示出略差的模型性能:
圖 11. 策略變體 2 的性能表現。圖片來自文獻[5]。
以上所有結果均表明,最簡單的策略反而是最佳選擇。
04 Summary
在本文中,我們解釋了 DeepSeek 模型的核心架構創新之一 —— DeepSeekMoE 中使用的無輔助損失函數的負載均衡方法。
本文首先介紹了混合專家模型(MoE)的基本原理,強調了負載均衡的重要性,并回顧了先前的解決方案(包括 auxiliary loss 方法和 Expert Choice 機制)。接著,本文闡釋了 DeepSeek 的無損失函數的負載均衡方法及其性能表現。
DeepSeek 的無損失函數方法在保持因果關系的同時避免了引入梯度干擾,其有效性已通過原論文的實證結果得到驗證。
感謝您花時間閱讀本文!
參考文獻
[1] DeepSeek(??https://www.deepseek.com/)??
[2] DeepSeek-V3 Technical Report(??https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf)??
[3] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model(??https://arxiv.org/abs/2405.04434)??
[4] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models(??https://arxiv.org/abs/2401.06066)??
[5] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts(??https://arxiv.org/abs/2408.15664)??
[6] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding(??https://arxiv.org/abs/2006.16668)??
[7] Mixture-of-Experts with Expert Choice Routing(??https://arxiv.org/abs/2202.09368)??
Thanks for reading!
Hope you have enjoyed and learned new things from this blog!
About the author
Shirley Li
I am a Machine Learning Engineer working on building multi-modality models to solve real-world problems.
END
本期互動內容 ??
?在實際工程中,您認為負載均衡對模型性能的影響有多大?除了本文提到的技術路徑,您還了解哪些有效的負載均衡方案?
原文鏈接:
??https://ai.gopubby.com/deepseek-v3-explained-3-auxiliary-loss-free-load-balancing-4beeb734ab1f??
