當前，MoE 架構是業界拓展模型規模的重要方向，然而，其在分布式訓練中存在的大量通信開銷，仍嚴重制約了訓練效率和成本。為攻克這一瓶頸，豆包大模型團隊提出了一個全新的通信優化系統 COMET，通過更精準、細粒度的計算-通信重疊技術，在大規模 MoE 模型上可達到單層 1.96 倍加速，端到端平均 1.71 倍效率提升，且在不同并行策略、輸入規模及硬件環境下均表現穩定。

目前，COMET 已實際應用于萬卡級生產集群，助力 MoE 模型高效訓練，并已累計節省了數百萬 GPU 小時資源。此外，COMET 還可與豆包大模型團隊此前發布的新一代稀疏模型架構 UltraMem 結合，實現協同優化。

該工作獲 MLSys 2025 會議 5/5/5/4 高分評審，核心代碼已開源。

Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts

論文鏈接：https://arxiv.org/pdf/2502.19811

開源地址：https://github.com/bytedance/flux

混合專家模型（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 的通信接收遠程數據（Token B），這種由于復雜數據依賴導致的計算-通信粒度上的錯配，使得效率嚴重下滑。

圖 1：單層 MoE 模型示意圖

（專家分布在 GPU0 和 GPU1 兩張卡上）

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

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

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

 2. COMET 核心方案 

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

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

圖 2：COMET 的設計結構

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

1）共享張量依賴解析

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

張量分解：將 MoE 層間傳遞的共享張量沿 Token 維度（M）或隱層維度（N）進行切割，使通信與計算的最小單元對齊。例如，在 MoE 第一層（Layer 0，圖 3 左）沿 M 維度分解，使通信和計算在 M 維度進行對齊；在 MoE 第二層（Layer 1，圖 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 方案并不沖突，兩者結合或將帶來更好的優化空間。

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

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