字節對MoE模型訓練成本再砍一刀，成本可節省40%！

剛剛，豆包大模型團隊在GitHub上開源了叫做COMET的MoE優化技術。

COMET已應用于字節的萬卡訓練集群，在真實的生產環境中，累計幫助節省了數百萬GPU小時。

早前，豆包團隊發布了新一代稀疏架構UltraMem，將模型推理成本砍掉 83%，此次，又開源了COMET，向模型訓練成本出手。從技術理念上看，兩者還可以結合使用，組成一套“砍價刀法”。

具體來看，COMET主要針對的是MoE模型在分布式訓練中，仍存在大量通信開銷的問題。

COMET內部通過一套細粒度計算-通信折疊技術，并結合GPU資源的動態分配，極致壓榨了MoE專家“摸魚閑置”的時間，在大規模MoE的單個執行層上可提速1.96倍，端到到平均提速1.71倍。

有趣的是，此前DeepSeek也專門針對MoE的通信瓶頸，開源了DualPipe+DeepEP方案，通過排布算子來掩蓋通信開銷。豆包團隊則直接喊話，兩種方案一起開掛，可能帶來更大的提升空間。

不過，COMET這種直接將計算-通信算子融合的方法，可以在MoE訓練框架中像插件一樣直接插拔使用，無需侵入性改動，部署更加方便、靈活，且支持業界絕大部分主流大模型。

因簡潔、高效的設計理念，該工作5/5/5/4高分入選了全球機器學習系統頂級會議 MLSys 2025，并被認為在實際的大規模生產環境中極具應用價值。

接下來，詳細看下COMET的技術解讀。

MoE緊迫的通信開銷問題

混合專家模型（MoE）通過稀疏激活機制突破了傳統稠密模型（Dense Model）的計算瓶頸，然而，MoE的分布式訓練仍面臨一項嚴峻挑戰：跨設備通信開銷巨大。

例如，Mixtral-8x7B模型在Megatron-LM框架中的通信時間占比可高達40%，嚴重制約了訓練效率和成本。

核心問題在于，MoE的專家網絡分布在多個GPU上，每次計算需頻繁執行Token分發與結果聚合，導致GPU計算資源大量閑置。因此，如何將通信隱藏到計算的過程中，提升模型訓練效率、節省計算資源，成為了MoE系統優化的關鍵。

1、通信隱藏到計算里？

一種方案是將流水線調度與通信算子結合起來，即通過定制訓練中流水線并行的調度方式，將不同microbatch的計算和通信進行重疊，例如DeepSeek的DualPipe。但是，這一方式會導致較大的顯存開銷，并需要對現有訓練框架進行復雜的侵入性改動。

其它MoE系統方案則是在microbatch內部采用了粗粒度的計算-通信流水線，將輸入數據分割成「數據塊」進行通信與計算的重疊。然而，這種粗粒度的重疊方式難以高效利用計算資源，且無法實現無縫的通信延遲隱藏，尤其在動態路由、異構硬件環境下，性能損失顯著。

因此，團隊認為現有的系統級MoE解決方案仍面臨兩大困境：

1）難以解決復雜的數據依賴

MoE架構的稀疏特性導致計算和通信間的依賴動態且復雜。MoE會動態地將Token分配給不同專家，而傳統的粗粒度矩陣分塊方式，會導致GPU頻繁等待遠程數據，從而造成計算資源閑置。

如圖1所示，當專家0需要在紫色「數據塊」中進行Tile-level的計算時，必須先通過Token-level的通信接收遠程數據（TokenB），這種由于復雜數據依賴導致的計算-通信粒度上的錯配，使得效率嚴重下滑。

△ 1：單層MoE模型示意圖（專家分布在GPU0和GPU1兩張卡上）

2）難以消除計算-通信流水線氣泡

另一個問題是，現有方法無法精確控制計算任務和通信任務對硬件資源的使用，因而，也無法根據不同的模型結構和動態輸入，來自適應地調整資源分配。這導致計算和通信無法實現無縫重疊，進而產生大量流水線氣泡，增加了系統的延遲。

因此，團隊認為：解決MoE模型中計算與通信的粒度不匹配問題是實現兩者高效重疊的關鍵，同時，還需要根據負載情況自適應調整通信和計算的資源分配，以進一步實現無縫重疊。

2、COMET：最小化整體低延遲，提升性能

COMET是一個針對MoE模型的通信優化系統，通過細粒度計算-通信重疊技術，助力大模型訓練優化。

團隊分析發現，MoE架構包含兩條不同的生產-消費流水線：「計算-通信流水線」和「通信-計算流水線」。如圖2所示，數據在流水線中流動時，各流水線內的操作會通過一個共享緩沖區鏈接，該緩沖區被稱作「共享張量」。

△圖2：COMET的設計結構

基于此，COMET引入兩項關鍵機制，以最小化整體延遲并提升流水線性能。

1）共享張量依賴解析

通過分解和重調度共享張量，解決通信與計算之間的粒度錯配問題，實現細至單Token級的重疊。

張量分解：將MoE層間傳遞的共享張量沿Token維度（M）或隱層維度（N）進行切割，使通信與計算的最小單元對齊。例如，在MoE第一層（Layer0，圖3左）沿M維度分解，使通信和計算在M維度進行對齊；在MoE第二層（Layer1，圖3右）沿N維度分解，細粒度傳輸Token結果，保證計算和通信的高效重疊。

△圖3：COMET對共享張量進行依賴解析和分解

計算重調度：為了更好地隱藏計算與通信的延遲，COMET會動態調整數據塊的計算順序。例如，優先計算本地數據塊，同時異步拉取遠程Token。當某個專家需處理Token A（本地）和Token B（遠程）時，系統會優先啟動Token A的計算線程，并與Token B的通信線程并行執行，從而消除等待延遲。

△圖4：COMET在MoE layer0中分解并重新調度共享張量

2）自適應負載分配

動態分配GPU線程塊資源，精準平衡通信與計算負載，消除流水線氣泡。

線程塊隔離：將通信與計算任務分別封裝在獨立線程塊中，避免遠程I/O阻塞計算核心。在Nvidia Hopper架構中，計算線程塊專用于執行異步TMA指令的GEMM運算，通信線程塊通過NVSHMEM實現單Token級數據傳輸，這種設計賦予了系統在算子級別進行資源管理的能力。

△圖5：COMET的計算/通信線程塊隔離設計

動態負載平衡：根據輸入規模（如Token長度M）、并行策略（EP/TP比例）實時調整線程塊分配。如圖6所示，當TP=8、EP=1時，通信線程塊占所有線程塊的比例為19.7%，而在TP=4、EP=2時，該比例需提升至34.8%，系統通過預編譯多個版本的計算-通信融合算子實現在運行時的「零開銷」算子動態切換，并始終提供低延遲的算子。

△圖6：單個MoE層使用不同數量的通信線程塊的時延結果

3、大規模落地驗證

團隊在多個大規模MoE模型中評估了COMET的端到端性能。結果表明，COMET在8卡H800的實驗集群中，端到端MoE模型（Mixtral-8x7B、Qwen2-MoE等）的前向時延較其他基線系統可降低31.8%-44.4%，且在不同并行策略、輸入規模及硬件環境下均表現穩定。

△圖7：COMET在多個端到端MoE模型中的測評結果

在單個MoE層上，當輸入Token數量不同的情況下，COMET的執行時間均顯著短于基線方案，平均實現了1.28倍到2.37倍的速度提升

△圖8：COMET在單個MoE層不同輸入Token長度下的延遲情況

目前，COMET已實際應用于萬卡級生產集群，助力MoE模型高效訓練，并已累計節省數百萬GPU小時。該工作在MLSys 2025會議獲得5/5/5/4的評審高分，并被認為在大規模生產環境中極具應用潛力。

具體表現為：

強魯棒性：COMET采用的細粒度計算-通信重疊方案即使在專家負載不均衡的場景下，也能保持低于其它基線系統的延遲，具有較好的魯棒性；

強泛化能力：COMET在NVLink與PCIe等不同網絡環境下均能提供穩定的加速比；在使用不同并行策略時均能生成低時延算子，以供大規模訓練框架使用。

4、核心代碼開源

COMET包含約1.2萬行C++和CUDA代碼，以及2千行Python代碼，并向開發者提供了一套友好的Python API。

△圖9：COMET 開源頁面

此外，COMET建立了面向MoE的細粒度流水線編程范式，通過深度融合NVSHMEM通信庫與CUTLASS高效計算算子，實現了通信操作與GEMM計算的算子內融合。例如，MoE Layer 1的GEMM計算與Token聚合通信可在單一GPU算子內完成。這與此前提到的Deepseek DualPipe+DeepEP方案并不沖突，兩者結合或將帶來更好的優化空間。

此外，COMET可直接接入已有的MoE訓練框架，支持TP/EP/EP+TP多種并行模式，并提供了靈活的插拔式部署方案。

目前，COMET核心代碼已開源，并計劃兼容Triton等編譯生態。

論文鏈接：
https://arxiv.org/pdf/2502.19811
開源地址：
https://github.com/bytedance/flux