project
COMET - ByteDanceのオープンソース通信最適化システム
ByteDanceが開発したCOMETは、分散トレーニングにおける過剰な通信オーバーヘッドの問題に対処するために設計された、エキスパート混合モデル(MoE)向けの最適化システムです。きめ細かな計算・通信オーバーラップ技術に基づき、COMETは…
COMETとは何ですか?
ByteDanceが開発したCOMETは、分散トレーニングにおける過剰な通信オーバーヘッドの問題を解決する、エキスパート混合モデル(MoE)向けの最適化システムです。COMETは、きめ細かな計算・通信オーバーラップ技術に基づき、通信と計算操作を深く統合することで、従来の手法における粒度の不一致によるリソースの浪費と遅延を排除します。COMETは、共有テンソル依存性解決と適応型負荷分散という2つのメカニズムを導入し、リソース割り当てを動的に調整することで、効率的な通信隠蔽と計算最適化を実現します。大規模な本番環境において、COMETはMoEモデルのトレーニング効率を大幅に向上させ、単一レイヤーで最大1.96倍、エンドツーエンドのトレーニングで最大1.71倍の高速化を実現し、累計で数百万GPU時間を節約します。COMETのコアコードはオープンソースであり、主流の大規模モデルと様々な並列戦略をサポートし、既存のトレーニングフレームワークにシームレスに統合できます。
COMETの主な機能
- きめ細かい計算 - 通信の重複共有テンソルを分解し、計算順序を再スケジュールすることで、計算と通信間のきめ細かな連携を実現します。
- 適応型負荷分散入力サイズと並列処理戦略に基づいて、GPUスレッドブロックのリソースを動的に調整し、通信負荷と計算負荷のバランスを取り、パイプラインのバブルを解消し、全体的な効率を向上させます。
- 高効率なリソース管理通信タスクと計算タスクを独立したスレッドブロックにカプセル化することで、リモートI/Oが計算コアをブロックするのを防ぎ、リソース利用率を向上させます。
- 高い堅牢性と汎化能力低遅延を維持し、不均一な専門家負荷や異なるハードウェア環境下でも、複数の並列処理戦略と大規模なクラスタ展開をサポートします。
- 簡単に統合できます既存のMoEトレーニングフレームワークにプラグインとして直接統合でき、大幅な変更を加えることなく、主流のコンパイルエコシステム(Tritonなど)をサポートします。
COMETの技術原理
- 共有テンソルの依存関係解決:
- テンソル分解MoEレイヤー間で渡される共有テンソルは、通信と計算の最小単位を揃えるために、トークン次元(M)または隠れ層次元(N)に沿って分割されます。例えば、MoEレイヤー1(レイヤー0)はM次元に沿って分割され、レイヤー2(レイヤー1)はN次元に沿って分割されることで、効率的なオーバーラップが実現されます。
- 計算によるスケジュール変更データブロックの計算順序を動的に調整し、ローカルデータブロックの計算を優先し、リモートトークンを非同期的に取得し、リモートデータの待機によって発生する計算遅延を排除します。
- 適応型負荷分散:
- スレッドブロック分離通信タスクと計算タスクは、リモートI/Oが計算コアをブロックするのを避けるため、別々のスレッドブロックにカプセル化されています。計算スレッドブロックは、非同期TMA命令を使用してGEMM演算を実行することに特化しており、通信スレッドブロックはNVSHMEMに基づいてシングルトークンレベルのデータ伝送を実装します。
- 動的負荷分散入力サイズ(トークン長など)と並列処理戦略(EP/TP比)に基づいて、スレッドブロックの割り当て比率がリアルタイムで調整され、計算・通信融合演算子の複数の事前コンパイル済みバージョンに基づいて、実行時に「オーバーヘッドゼロ」の動的な切り替えが実現されます。
COMETのパフォーマンス
- 単層加速大規模なMoEモデルでは、単層あたり1.96倍の高速化が達成される。
- エンドツーエンドのアクセラレーション完全なMoEモデルでは、エンドツーエンドの平均実装速度向上率は1.71倍です。
- 異なるモデルおよび入力サイズにおける安定性:
- いくつかの大規模なMoEモデル(Mixtral-8x7B、Qwen2-MoEなど)において、COMETは他のベースラインシステムと比較して、フォワードレイテンシを31.8%~44.4%削減します。
- 入力トークンの長さが異なる場合でも、COMETの実行時間はベースラインソリューションよりも大幅に短く、平均で1.28倍から2.37倍の速度向上を実現しています。
- 堅牢性専門家の負荷が不均一なシナリオにおいて、COMETは他のベースラインシステムよりも低いレイテンシを維持し、優れた堅牢性を示します。
- 一般化能力COMETは、NVLinkやPCIeなどの様々なネットワーク環境において安定した高速化を実現し、複数の並列処理戦略(EP、TP、EP+TPなど)をサポートしており、大規模なトレーニングフレームワークに適しています。
COMETプロジェクトの住所
- GitHubリポジトリ:https://github.com/bytedance/flux
- arXiv技術論文:https://arxiv.org/pdf/2502.19811
COMETの応用シナリオ
- 大規模MoEモデルの高速トレーニングMixtral-8x7Bなどの大規模モデルの分散学習を最適化し、GPU処理時間を節約し、学習効率を向上させます。
- 異種ハードウェア環境への適応H800やL20など、さまざまなGPUやネットワーク環境(NVLinkやPCIeなど)をサポートし、安定したアクセラレーションを実現します。
- マルチパラレル戦略サポートEP、TP、およびハイブリッド並列処理戦略に対応しており、さまざまなモデルやハードウェア構成に適応します。
- 動的負荷分散エキスパートの負荷が不均一な場合は、低遅延動作を維持するためにリソース割り当てを動的に調整します。
- 既存のフレームワークとのシームレスな統合これは、MoEのトレーニングフレームワーク(Megatron-LMなど)にプラグインとして直接統合できるため、煩雑な変更を加えることなく、迅速な導入が容易になります。