字節(jié)對MoE模型訓(xùn)練成本再砍一刀，成本可節(jié)省40%！

剛剛，豆包大模型團(tuán)隊在GitHub上開源了叫做COMET的MoE優(yōu)化技術(shù)。

COMET已應(yīng)用于字節(jié)的萬卡訓(xùn)練集群，在真實的生產(chǎn)環(huán)境中，累計幫助節(jié)省了數(shù)百萬GPU小時。

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

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

COMET內(nèi)部通過一套細(xì)粒度計算-通信折疊技術(shù)，并結(jié)合GPU資源的動態(tài)分配，極致壓榨了MoE專家“摸魚閑置”的時間，在大規(guī)模MoE的單個執(zhí)行層上可提速1.96倍，端到到平均提速1.71倍。

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

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

因簡潔、高效的設(shè)計理念，該工作5/5/5/4高分入選了全球機(jī)器學(xué)習(xí)系統(tǒng)頂級會議 MLSys 2025，并被認(rèn)為在實際的大規(guī)模生產(chǎn)環(huán)境中極具應(yīng)用價值。

接下來，詳細(xì)看下COMET的技術(shù)解讀。

MoE緊迫的通信開銷問題

混合專家模型（MoE）通過稀疏激活機(jī)制突破了傳統(tǒng)稠密模型（Dense Model）的計算瓶頸，然而，MoE的分布式訓(xùn)練仍面臨一項嚴(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à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)行通信與計算的重疊。然而，這種粗粒度的重疊方式難以高效利用計算資源，且無法實現(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ù)（TokenB），這種由于復(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)致計算和通信無法實現(xiàn)無縫重疊，進(jìn)而產(chǎn)生大量流水線氣泡，增加了系統(tǒng)的延遲。

因此，團(tuán)隊認(rèn)為：解決MoE模型中計算與通信的粒度不匹配問題是實現(xiàn)兩者高效重疊的關(guān)鍵，同時，還需要根據(jù)負(fù)載情況自適應(yīng)調(diào)整通信和計算的資源分配，以進(jìn)一步實現(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引入兩項關(guān)鍵機(jī)制，以最小化整體延遲并提升流水線性能。

1）共享張量依賴解析

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

張量分解：將MoE層間傳遞的共享張量沿Token維度（M）或隱層維度（N）進(jìn)行切割，使通信與計算的最小單元對齊。例如，在MoE第一層（Layer0，圖3左）沿M維度分解，使通信和計算在M維度進(jìn)行對齊；在MoE第二層（Layer1，圖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實現(xiàn)單Token級數(shù)據(jù)傳輸，這種設(shè)計賦予了系統(tǒng)在算子級別進(jìn)行資源管理的能力。

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

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

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

3、大規(guī)模落地驗證

團(tuán)隊在多個大規(guī)模MoE模型中評估了COMET的端到端性能。結(jié)果表明，COMET在8卡H800的實驗集群中，端到端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í)行時間均顯著短于基線方案，平均實現(xiàn)了1.28倍到2.37倍的速度提升

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

目前，COMET已實際應(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高效計算算子，實現(xiàn)了通信操作與GEMM計算的算子內(nèi)融合。例如，MoE Layer 1的GEMM計算與Token聚合通信可在單一GPU算子內(nèi)完成。這與此前提到的Deepseek DualPipe+DeepEP方案并不沖突，兩者結(jié)合或?qū)砀玫膬?yōu)化空間。

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

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

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