從DeepSeek MoE專家負載均衡談起
上周中的時候, 同事給我了一份線上DeepSeek-R1推理的Expert激活的數據用來研究一些專家負載均衡的算法, 當然這些線上數據來自于公司內部的請求, 從中觀測出前面10層專家基本上是相對均衡的, 而越到后面不均衡程度差異越大。當時討論到這個問題時, 是懷疑內部的一些請求是否專注于電商領域而帶來的不平衡特性, 于是做了一些研究。恰好搜到Intel的一篇論文《Semantic Specialization in MoE Appears with Scale: A Study of DeepSeek-R1 Expert Specialization》有一些基于語義的MoE分析專家的專業(yè)性相關的問題, 再加上前幾天看到某個公眾號采訪某院長的一個比較有趣的說法:“Dense模型適合toB業(yè)務,MoE模型適合toC業(yè)務”. 因此做了一些分析, 在此記錄下來。
1. 專家Overlap分析
從這篇論文的第一個Word-in-Context的實驗來看, DeepSeek-R1的前面十層專家Overlap的概率相對于較高, 和線上的一些數據分析是一致的。
圖片
比較特別的是在第十層后,不同語義和相似語義之間的區(qū)分度完全顯現出來了, 而模型本身因為細粒度MoE(256選8)而產生的區(qū)分度也顯著降低了, 同時論文還對比了Mistral的兩個MoE模型, 它采用8選2的方式, 看來語義間對不同專家的區(qū)分度有很大的差距. 這個結論也支撐了DeepSeek逐漸向更加細粒度專家的技術路線的正確性。
當然產生這樣的結果有幾方面可能的因素:
- Shared Expert的重要性, 通過Shared Expert消除了一些專家之間的影響, 使得Routed Expert Overlap的概率降低?
- 本質上就是Routed Expert數量的影響?
- R1強化學習的工作流對于Expert Specialization進一步增強了?
但是值得注意的是另一個問題是, 在模型的后面20層內, 層間的Overlap的差異還是很大的, 并且沒有進一步的下降, 這個和我拿到的線上的數據分布也是相似的.
這里引入一個思考, 每一層模型的AlltoAll通信時間實際上是受到分布式部署的帶寬和延遲約束的, 因此模型深度過深后將會影響到TPOT, 雖然可以用一些ScaleUP的辦法來解決, 但是看看GB200的可靠性和成本, 這種取舍是不太恰當的。另一方面, 看到上圖中第40層overlap有明顯的抖動, 一方面是模型在后面的層中還可以更加稀疏來進一步降低Overlap, 是否也會有一個類似的ScalingLaw我們在稍微后面的章節(jié)來分析。
2. SAE分析
從論文中SAE的分析來看, 能夠得出不同的專家在負責不同的推理以及認知專業(yè)化的結論, 這和DeepSeek設計細粒度MoE和專家專業(yè)化的初衷是匹配的。
圖片
其實渣B一直在建議從SAE的角度來分析大模型, 并通過對SAE Activation的約束來作為強化學習工作流的一種手段,SAE對于概念的可視化解釋, Anthropic和OAI都做了相應的可視化展示, 例如Anthropic的多模態(tài)對金門大橋的概念。
圖片
OAI和Claude都在這方面有了蠻長時間的布局, 而國內相對還是落后了一些。
3. 從范疇論的視角看R1
這是一個爛尾很久的專題, 一直想抽一周的時間來好好分析并寫一篇筆記, 但是最近幾個月不停的在各種項目的死線上掙扎. 先簡短的寫一些吧. 其實R1的整個訓練流程從范疇論的意義上來看:
- 首先是V3-Base的模型本質上是通過一系列數據集的Pre-train流程構成了一個預層范疇(Presheaf)。
- R1-Zero是基于V3-Base的Presheaf上來強化了一些Morphism的權重, 而這些權重在MoE模型的底子上使得模型具有了更強的泛化能力。
- 然后在V3-Base的基礎上混合R1-Zero的Coldstart數據和一些General samples來構建最終的R1。
圖片
比較好奇的是在整個后訓練的過程中, 不知道DeepSeek是否記錄了梯度更新的情況, 感覺這個地方配合SAE做一些分析可能會有更多的發(fā)現, 個人覺得雖然ORM取得了很好的結果, 而PRM本身還有一些過程上的缺陷, 是否可以在SAE的視角上來看出更多的原因, 并且某種意義上還可以給ORM訓練輸出一些更加抽象泛化的約束能力。
當然這樣也會面臨一個比較大的算力的挑戰(zhàn), SAE的算力消耗和RL工作流的整體效率上的一個取舍問題。
4. MoE ScalingLaw
本文開頭提到了一個比較有趣的說法:“Dense模型適合toB業(yè)務,MoE模型適合toC業(yè)務”, GPT4是MoE模型吧, 它適合toB還是toC? Llama3是一個Dense模型吧? 它適合toB還是toC? 本質的問題是算力的約束下MoE成為繼續(xù)提高Scaling的一個必然手段. 當然MoE模型本身的Gating數值穩(wěn)定性問題和Reasoning模型本身通常設置的溫度參數相對較低, 使得模型的幻覺程度有所增加而不太適合一些toB的業(yè)務場景。
最近還有一篇《Chain-of-Experts: 釋放MoE專家的溝通潛能》[2]挺有意思的, 即通過在同一層的專家之間的互相處理來得到最后的output hidden. 實際上這里又有了一些RNN的味道. 但是這樣的機制如果迭代次數多了感覺很難去兼顧訓練和推理的效率。
圖片
從本文第一節(jié)的配圖上來看, 似乎某種程度上能夠得出和DeepSpeed-MoE[3]中提出的pyramid-MoE相似的結構, 隨著模型的層數越來越深, 專家專業(yè)化程度越來越高, 相應的專家數量和TopK選擇數量也需要對應的提高?
圖片
其實這也是我最近在考慮的一個問題, MoE的本質是否和HNSW(Hierarchical Navigable Small Word)算法某種程度上有相似性?
圖片
那么借助Grace+Blackwell的架構, 是否還能做出點有趣的東西呢? 大概想到一個增量MoE的算法:
- 首先按照一個相對細粒度的模型進行訓練, 例如256 Routed Experts, TopK=8
- 例如訓練到500B tokens時, 模型逐漸添加一些新的專家在后面若干層
- 反復訓練的過程中把模型逐漸迭代成一個金字塔結構.
- 最后在PostTraining過程中, 基于SAE或者某些層的MoE路由規(guī)則凍結一些Expert的參數或者是在這個基礎上做一些KL散度的約束來降低幻覺?
為什么需要Grace呢, 因為某種程度上還是需要CPU側的更大的內存空間來做一些專家權重的置換. PCIe本身的帶寬還是太小了. 當然這樣的模型部署時在推理階段可能還有更多的挑戰(zhàn). 設計模型架構時兼顧推理性能是必須要考慮的一個因素了, 這部分內容暫時還沒想明白, 隱約覺得在這樣的一個模型下, 順便把Next Few layer的Expert Prediction/Prefetch做了可能是一條路.
目前,阿里云正在GPU加CPU的異構資源池上做優(yōu)化。未來,數據庫要研發(fā)的關鍵能力是將昂貴的GPU盡可能地省下來做最珍貴的計算和緩存,將次要的計算和緩存推到CPU加內存和存儲的三層池化中,讓在線推理變得更低成本。
在基礎設施和分布式系統(tǒng)的視角來看, 和模型的協(xié)同還有更多的工作要做。
參考資料
[1] Semantic Specialization in MoE Appears with Scale: A Study of DeepSeek-R1 Expert Specialization: https://arxiv.org/pdf/2502.10928
[2] Chain-of-Experts: 釋放MoE專家的溝通潛能: https://sandy-server-87f.notion.site/Chain-of-Experts-MoE-1ab9bb750b79801bbfebf01ae9a77b3f
[3] DeepSpeed-MoE: https://arxiv.org/pdf/2201.05596