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

目前，COMET 已實(shí)際應(yīng)用于萬卡級生產(chǎn)集群，助力 MoE 模型高效訓(xùn)練，并已累計節(jié)省了數(shù)百萬 GPU 小時資源。此外，COMET 還可與豆包大模型團(tuán)隊此前發(fā)布的新一代稀疏模型架構(gòu) UltraMem 結(jié)合，實(shí)現(xiàn)協(xié)同優(yōu)化。

該工作獲 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）通過稀疏激活機(jī)制突破了傳統(tǒng)稠密模型（Dense Model）的計算瓶頸，然而，MoE 的分布式訓(xùn)練仍面臨一項(xiàng)嚴(yán)峻挑戰(zhàn)：跨設(shè)備通信開銷巨大。例如，Mixtral-8x7B 模型在 Megatron-LM 框架中的通信時間占比可高達(dá) 40%，嚴(yán)重制約了訓(xùn)練效率和成本。

核心問題在于，MoE 的專家網(wǎng)絡(luò)分布在多個 GPU 上，每次計算需頻繁執(zhí)行 Token 分發(fā)與結(jié)果聚合，導(dǎo)致 GPU 計算資源大量閑置。因此，如何將通信隱藏到計算的過程中，提升模型訓(xùn)練效率、節(jié)省計算資源，成為了 MoE 系統(tǒng)優(yōu)化的關(guān)鍵。

 1. 難點(diǎn):「復(fù)雜的數(shù)據(jù)依賴」與「流水線氣泡」

為了掩蓋巨大的通信開銷，現(xiàn)有方案主要集中在如何對「計算-通信」進(jìn)行高效重疊。

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

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

因此，團(tuán)隊認(rèn)為現(xiàn)有的系統(tǒng)級 MoE 解決方案仍面臨兩大困境：

1）難以解決復(fù)雜的數(shù)據(jù)依賴

MoE 架構(gòu)的稀疏特性導(dǎo)致計算和通信間的依賴動態(tài)且復(fù)雜。MoE 會動態(tài)地將 Token 分配給不同專家，而傳統(tǒng)的粗粒度矩陣分塊方式，會導(dǎo)致 GPU 頻繁等待遠(yuǎn)程數(shù)據(jù)，從而造成計算資源閑置。

如圖 1 所示，當(dāng)專家 0 需要在紫色「數(shù)據(jù)塊」中進(jìn)行 Tile-level 的計算時，必須先通過 Token-level 的通信接收遠(yuǎn)程數(shù)據(jù)（Token B），這種由于復(fù)雜數(shù)據(jù)依賴導(dǎo)致的計算-通信粒度上的錯配，使得效率嚴(yán)重下滑。

圖 1：單層 MoE 模型示意圖

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

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

另一個問題是，現(xiàn)有方法無法精確控制計算任務(wù)和通信任務(wù)對硬件資源的使用，因而，也無法根據(jù)不同的模型結(jié)構(gòu)和動態(tài)輸入，來自適應(yīng)地調(diào)整資源分配。這導(dǎo)致計算和通信無法實(shí)現(xiàn)無縫重疊，進(jìn)而產(chǎn)生大量流水線氣泡，增加了系統(tǒng)的延遲。

因此，團(tuán)隊認(rèn)為：解決 MoE 模型中計算與通信的粒度不匹配問題是實(shí)現(xiàn)兩者高效重疊的關(guān)鍵，同時，還需要根據(jù)負(fù)載情況自適應(yīng)調(diào)整通信和計算的資源分配，以進(jìn)一步實(shí)現(xiàn)無縫重疊。

 2. COMET 核心方案 

COMET 是一個針對 MoE 模型的通信優(yōu)化系統(tǒng)，通過細(xì)粒度計算-通信重疊技術(shù)，助力大模型訓(xùn)練優(yōu)化。

團(tuán)隊分析發(fā)現(xiàn)，MoE 架構(gòu)包含兩條不同的生產(chǎn)-消費(fèi)流水線:「計算-通信流水線」和「通信-計算流水線」。如圖 2 所示，數(shù)據(jù)在流水線中流動時，各流水線內(nèi)的操作會通過一個共享緩沖區(qū)鏈接，該緩沖區(qū)被稱作「共享張量」。

圖 2：COMET 的設(shè)計結(jié)構(gòu)

基于此，COMET 引入兩項(xiàng)關(guān)鍵機(jī)制，以最小化整體延遲并提升流水線性能。

1）共享張量依賴解析

通過分解和重調(diào)度共享張量，解決通信與計算之間的粒度錯配問題，實(shí)現(xiàn)細(xì)至單 Token 級的重疊。

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

圖 3：COMET 對共享張量進(jìn)行依賴解析和分解

計算重調(diào)度：為了更好地隱藏計算與通信的延遲，COMET 會動態(tài)調(diào)整數(shù)據(jù)塊的計算順序。例如，優(yōu)先計算本地數(shù)據(jù)塊，同時異步拉取遠(yuǎn)程 Token。當(dāng)某個專家需處理 Token A（本地）和 Token B（遠(yuǎn)程）時，系統(tǒng)會優(yōu)先啟動 Token A 的計算線程，并與 Token B 的通信線程并行執(zhí)行，從而消除等待延遲。

圖 4：COMET 在 MoE layer0 中分解并重新調(diào)度共享張量

2）自適應(yīng)負(fù)載分配

動態(tài)分配 GPU 線程塊資源，精準(zhǔn)平衡通信與計算負(fù)載，消除流水線氣泡。

線程塊隔離：將通信與計算任務(wù)分別封裝在獨(dú)立線程塊中，避免遠(yuǎn)程 I/O 阻塞計算核心。在 Nvidia Hopper 架構(gòu)中，計算線程塊專用于執(zhí)行異步 TMA 指令的 GEMM 運(yùn)算，通信線程塊通過 NVSHMEM 實(shí)現(xiàn)單 Token 級數(shù)據(jù)傳輸，這種設(shè)計賦予了系統(tǒng)在算子級別進(jìn)行資源管理的能力。

圖 5：COMET 的計算/通信線程塊隔離設(shè)計

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

圖 6：單個 MoE 層使用不同數(shù)量的通信線程塊的時延結(jié)果

 3. 大規(guī)模落地驗(yàn)證 

團(tuán)隊在多個大規(guī)模 MoE 模型中評估了 COMET 的端到端性能。結(jié)果表明，COMET 在 8 卡 H800 的實(shí)驗(yàn)集群中，端到端 MoE 模型（Mixtral-8x7B、Qwen2-MoE 等）的前向時延較其他基線系統(tǒng)可降低 31.8%-44.4%，且在不同并行策略、輸入規(guī)模及硬件環(huán)境下均表現(xiàn)穩(wěn)定。

圖 7：COMET 在多個 MoE 模型中的測評結(jié)果

在單個 MoE 層上，當(dāng)輸入 Token 數(shù)量不同的情況下，COMET 的執(zhí)行時間均顯著短于基線方案，平均實(shí)現(xiàn)了 1.28 倍到 2.37 倍的速度提升。

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

目前，COMET 已實(shí)際應(yīng)用于萬卡級生產(chǎn)集群，助力 MoE 模型高效訓(xùn)練，并已累計節(jié)省數(shù)百萬 GPU 小時。該工作在 MLSys 2025 會議獲得 5/5/5/4 的評審高分，并被認(rèn)為在大規(guī)模生產(chǎn)環(huán)境中極具應(yīng)用潛力。

具體表現(xiàn)為：

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

強(qiáng)泛化能力：COMET 在 NVLink 與 PCIe 等不同網(wǎng)絡(luò)環(huán)境下均能提供穩(wěn)定的加速比；在使用不同并行策略時均能生成低時延算子，以供大規(guī)模訓(xùn)練框架使用。

 4. 核心代碼開源 

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

圖9：COMET 開源頁面

此外，COMET 建立了面向 MoE 的細(xì)粒度流水線編程范式，通過深度融合 NVSHMEM 通信庫與 CUTLASS 高效計算算子，實(shí)現(xiàn)了通信操作與 GEMM 計算的算子內(nèi)融合。例如，MoE Layer 1 的 GEMM 計算與 Token 聚合通信可在單一 GPU 算子內(nèi)完成。這與此前提到的 Deepseek DualPipe 方案并不沖突，兩者結(jié)合或?qū)砀玫膬?yōu)化空間。

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

核心代碼現(xiàn)已開源，并計劃兼容 Triton 等編譯生態(tài)。