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

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

該工作獲 MLSys 2025 會(huì)議 5/5/5/4 高分評(píng)審，核心代碼已開(kāi)源。

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

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

開(kāi)源地址：https://github.com/bytedance/flux

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

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

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

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

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

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

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

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

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

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

圖 1：?jiǎn)螌?MoE 模型示意圖

（專(zhuān)家分布在 GPU0 和 GPU1 兩張卡上）

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

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

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

 2. COMET 核心方案 

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

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

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

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

1）共享張量依賴(lài)解析

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

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

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

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

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

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

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

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

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

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

圖 6：?jiǎn)蝹€(gè) MoE 層使用不同數(shù)量的通信線程塊的時(shí)延結(jié)果

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

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

圖 7：COMET 在多個(gè) MoE 模型中的測(cè)評(píng)結(jié)果

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

圖 8：COMET 在單個(gè) MoE 層不同輸入 Token 長(zhǎng)度下的延遲情況

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

具體表現(xiàn)為：

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

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

 4. 核心代碼開(kāi)源 

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

圖9：COMET 開(kāi)源頁(yè)面

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

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

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