3 月 10日,字節(jié)豆包大模型團(tuán)隊(duì)官宣開(kāi)源一項(xiàng)針對(duì) MoE 架構(gòu)的關(guān)鍵優(yōu)化技術(shù),可將大模型訓(xùn)練效率提升1.7倍,成本節(jié)省40%。據(jù)悉,該技術(shù)已實(shí)際應(yīng)用于字節(jié)的萬(wàn)卡集群訓(xùn)練,累計(jì)幫助節(jié)省了數(shù)百萬(wàn) GPU 小時(shí)訓(xùn)練算力。
MoE 是當(dāng)前大模型的主流架構(gòu),但其在分布式訓(xùn)練中存在大量跨設(shè)備通信開(kāi)銷,嚴(yán)重制約了大模型訓(xùn)練效率和成本。以海外主流模型Mixtral-8x7B為例, 其訓(xùn)練過(guò)程中通信時(shí)間占比可高達(dá) 40%。針對(duì)這一難題,字節(jié)在內(nèi)部研發(fā)了COMET計(jì)算-通信重疊技術(shù),通過(guò)多項(xiàng)創(chuàng)新,大幅壓縮了MoE專家通信空轉(zhuǎn)時(shí)間。
相較DeepSeek近期開(kāi)源的DualPipe等MoE優(yōu)化方案,COMET可以像插件一樣直接接入已有的MoE訓(xùn)練框架,支持業(yè)界絕大部分主流大模型,無(wú)需對(duì)訓(xùn)練框架進(jìn)行侵入式改動(dòng)。因簡(jiǎn)潔、通用的設(shè)計(jì)理念,該工作以5/5/5/4 的高分入選全球機(jī)器學(xué)習(xí)系統(tǒng)頂級(jí)會(huì)議 MLSys 2025 ,被認(rèn)為“在大規(guī)模生產(chǎn)環(huán)境中極具應(yīng)用潛力”。
具體而言, COMET 從系統(tǒng)層面建立了面向 MoE 的細(xì)粒度流水線編程方式,通過(guò)引入共享張量依賴解析、自適應(yīng)負(fù)載分配兩項(xiàng)關(guān)鍵機(jī)制,來(lái)解決通信與計(jì)算之間的粒度錯(cuò)配問(wèn)題,并精準(zhǔn)平衡通信與計(jì)算負(fù)載,最終大幅提升MoE流水線整體效率。 引入COMET后,單個(gè) MoE 層上可實(shí)現(xiàn) 1.96 倍加速、端到端平均 1.71 倍效率提升,且在不同并行策略、輸入規(guī)模及硬件環(huán)境下均表現(xiàn)穩(wěn)定。

COMET 的設(shè)計(jì)結(jié)構(gòu)
COMET 與Deepseek 研發(fā)的DualPipe方案還可以聯(lián)合使用。在降低MoE通信開(kāi)銷上,COMET 采用了計(jì)算-通信融合算子的優(yōu)化方式, DualPipe則通過(guò)排布算子來(lái)掩蓋通信,兩種方案并不沖突,結(jié)合使用或?qū)⒏蠓葔嚎s模型訓(xùn)練成本。
目前,COMET支持多種MoE并行模式,部署靈活、方便。同時(shí),COMET核心代碼已開(kāi)源,并向開(kāi)發(fā)者提供了一套友好的 Python API,計(jì)劃兼容 Triton 等編譯生態(tài)。
COMET論文鏈接:https://arxiv.org/pdf/2502.19811
開(kāi)源地址:https://github.com/bytedance/flux