COMET – 字節(jié)開(kāi)源的通信優(yōu)化系統(tǒng)

COMET是字節(jié)跳動(dòng)推出的一款專(zhuān)為Mixture-of-Experts（MoE）模型設(shè)計(jì)的優(yōu)化系統(tǒng)，旨在解決分布式訓(xùn)練中的通信開(kāi)銷(xiāo)過(guò)大的問(wèn)題。通過(guò)細(xì)粒度的計(jì)算-通信重疊技術(shù)，COMET將通信與計(jì)算操作深度融合，有效消除傳統(tǒng)方法因粒度不匹配所造成的資源浪費(fèi)與延遲。該系統(tǒng)的引入使得MoE模型的訓(xùn)練效率顯著提升，單層加速可達(dá)1.96倍，端到端加速可達(dá)1.71倍，累計(jì)節(jié)省了數(shù)百萬(wàn)GPU小時(shí)。

COMET是什么

COMET是字節(jié)跳動(dòng)推出的一種針對(duì)Mixture-of-Experts（MoE）模型的優(yōu)化工具，旨在解決分布式訓(xùn)練過(guò)程中的高通信開(kāi)銷(xiāo)問(wèn)題。它依托于細(xì)粒度的計(jì)算-通信重疊技術(shù)，將通信和計(jì)算操作深度融合，消除了傳統(tǒng)方法所引發(fā)的資源浪費(fèi)和延遲。通過(guò)引入共享張量依賴(lài)解析和自適應(yīng)負(fù)載分配機(jī)制，COMET能夠動(dòng)態(tài)調(diào)整資源分配，從而實(shí)現(xiàn)高效的通信隱藏和計(jì)算優(yōu)化。在大規(guī)模的生產(chǎn)環(huán)境中，COMET展現(xiàn)出顯著的訓(xùn)練效率提升，單層加速達(dá)到1.96倍，端到端加速為1.71倍，積累節(jié)省的GPU小時(shí)數(shù)達(dá)到數(shù)百萬(wàn)個(gè)。COMET的核心代碼已經(jīng)開(kāi)源，支持主流大模型和多種并行策略，能夠無(wú)縫集成到現(xiàn)有的訓(xùn)練框架中。

COMET的主要功能

• 細(xì)粒度計(jì)算-通信重疊：通過(guò)分解共享張量和重新調(diào)度計(jì)算順序，實(shí)現(xiàn)計(jì)算與通信的細(xì)致對(duì)齊。
• 自適應(yīng)負(fù)載分配：動(dòng)態(tài)調(diào)整GPU線程塊資源，根據(jù)輸入規(guī)模和并行策略平衡通信與計(jì)算負(fù)載，消除流水線氣泡，提升整體效率。
• 高效資源管理：將通信和計(jì)算任務(wù)封裝在線程塊中，避免遠(yuǎn)程I/O阻塞計(jì)算核心，提升資源利用率。
• 強(qiáng)魯棒性和泛化能力：在專(zhuān)家負(fù)載不均衡或不同硬件環(huán)境下，保持低延遲，支持多種并行策略和大規(guī)模集群部署。
• 易于集成：作為插件直接接入現(xiàn)有的MoE訓(xùn)練框架，無(wú)需侵入性改動(dòng)，支持主流編譯生態(tài)（如Triton）。

COMET的技術(shù)原理

• 共享張量依賴(lài)解析：
  • 張量分解：將MoE層間傳遞的共享張量沿Token維度（M）或隱層維度（N）進(jìn)行切割，使通信與計(jì)算的最小單元對(duì)齊。例如，在MoE第一層（Layer 0）沿M維度分解，在第二層（Layer 1）沿N維度分解，實(shí)現(xiàn)高效重疊。
  • 計(jì)算重調(diào)度：動(dòng)態(tài)調(diào)整數(shù)據(jù)塊的計(jì)算順序，優(yōu)先計(jì)算本地?cái)?shù)據(jù)塊，異步拉取遠(yuǎn)程Token，消除因等待遠(yuǎn)程數(shù)據(jù)導(dǎo)致的計(jì)算延遲。
• 自適應(yīng)負(fù)載分配：
  • 線程塊隔離：將通信與計(jì)算任務(wù)分別封裝在線程塊中，避免遠(yuǎn)程I/O阻塞計(jì)算核心。計(jì)算線程塊專(zhuān)用于執(zhí)行異步TMA指令的GEMM運(yùn)算，而通信線程塊基于NVSHMEM實(shí)現(xiàn)單Token級(jí)數(shù)據(jù)傳輸。
  • 動(dòng)態(tài)負(fù)載平衡：根據(jù)輸入規(guī)模（如Token長(zhǎng)度）和并行策略（EP/TP比例），實(shí)時(shí)調(diào)整線程塊分配比例，基于預(yù)編譯多個(gè)版本的計(jì)算-通信融合算子，實(shí)現(xiàn)運(yùn)行時(shí)的“零開(kāi)銷(xiāo)”動(dòng)態(tài)切換。

COMET的性能表現(xiàn)

• 單層加速：在大規(guī)模MoE模型中，實(shí)現(xiàn)單層加速1.96倍。
• 端到端加速：在完整的MoE模型中，端到端平均實(shí)現(xiàn)加速1.71倍。
• 不同模型和輸入規(guī)模下的穩(wěn)定性：
  • 在多個(gè)大規(guī)模MoE模型（如Mixtral-8x7B、Qwen2-MoE等）中，COMET的前向時(shí)延相比其他基線系統(tǒng)降低31.8%-44.4%。
  • 在不同輸入Token長(zhǎng)度下，COMET的執(zhí)行時(shí)間顯著短于基線方案，平均速度提升1.28倍到2.37倍。
• 魯棒性：在專(zhuān)家負(fù)載不均衡的場(chǎng)景下，COMET能保持低于其他基線系統(tǒng)的延遲，表現(xiàn)出良好的魯棒性。
• 泛化能力：COMET在NVLink和PCIe等不同網(wǎng)絡(luò)環(huán)境下均能提供穩(wěn)定的加速比，支持多種并行策略（如EP、TP、EP+TP），適用于大規(guī)模訓(xùn)練框架。

COMET的項(xiàng)目地址

• GitHub倉(cāng)庫(kù)：https://github.com/bytedance/flux
• arXiv技術(shù)論文：https://arxiv.org/pdf/2502.19811

COMET的應(yīng)用場(chǎng)景

• 大規(guī)模MoE模型訓(xùn)練加速：優(yōu)化如Mixtral-8x7B等大型模型的分布式訓(xùn)練，節(jié)省GPU小時(shí)，提升訓(xùn)練效率。
• 異構(gòu)硬件環(huán)境適配：支持H800、L20等不同GPU及網(wǎng)絡(luò)環(huán)境（如NVLink、PCIe），確保穩(wěn)定加速。
• 多并行策略支持：兼容EP、TP和混合并行策略，適應(yīng)不同模型和硬件配置。
• 動(dòng)態(tài)負(fù)載均衡：在專(zhuān)家負(fù)載不均衡時(shí)，動(dòng)態(tài)調(diào)整資源分配，保持低延遲運(yùn)行。
• 無(wú)縫集成現(xiàn)有框架：作為插件直接接入MoE訓(xùn)練框架（如Megatron-LM），無(wú)需侵入性改動(dòng)，便于快速部署。