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

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

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

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

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

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

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

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

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

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

MoE緊迫的通信開(kāi)銷問(wèn)題

混合專家模型（MoE）通過(guò)稀疏激活機(jī)制突破了傳統(tǒng)稠密模型（Dense Model）的計(jì)算瓶頸，然而，MoE的分布式訓(xùn)練仍面臨一項(xiàng)嚴(yán)峻挑戰(zhàn)：跨設(shè)備通信開(kāi)銷巨大。

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

核心問(wèn)題在于，MoE的專家網(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、通信隱藏到計(jì)算里？

一種方案是將流水線調(diào)度與通信算子結(jié)合起來(lái)，即通過(guò)定制訓(xùn)練中流水線并行的調(diào)度方式，將不同microbatch的計(jì)算和通信進(jìn)行重疊，例如DeepSeek的DualPipe。但是，這一方式會(huì)導(dǎo)致較大的顯存開(kāi)銷，并需要對(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ù)依賴

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

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

△ 1：?jiǎn)螌覯oE模型示意圖（專家分布在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ū)被稱作「共享張量」。

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

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

1）共享張量依賴解析

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

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

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

計(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è)專家需處理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ì)算線程塊專用于執(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%，而在TP=4、EP=2時(shí)，該比例需提升至34.8%，系統(tǒng)通過(guò)預(yù)編譯多個(gè)版本的計(jì)算-通信融合算子實(shí)現(xiàn)在運(yùn)行時(shí)的「零開(kāi)銷」算子動(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ì)算-通信重疊方案即使在專家負(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+DeepEP方案并不沖突，兩者結(jié)合或?qū)?lái)更好的優(yōu)化空間。

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

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

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