🤖 machine learning

DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism

DisagMoE は、アテンション層とフィードフォワード層をマルチステージパイプラインを用いて個別の GPU グループに分散配置する新しい混合エキスパート学習システムであり、通信と計算を効果的に重畳させることで全対全通信のボトルネックを克服し、最大 1.8 倍の学習速度向上を実現して大規模クラスターにおける効率を向上させる。

原著者： Zhichen Zeng, Chi-Chih Chang, Jiayi Wang, Zezhou Wang, Ningxin Zheng, Zheng Zhong, Cesar A. Stuardo, Dongyang Wang, Mohamed S. Abdelfattah, Haibin Lin, Banghua Zhu, Ang Li, Ziheng Jiang

公開日 2026-05-13

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Zhichen Zeng, Chi-Chih Chang, Jiayi Wang, Zezhou Wang, Ningxin Zheng, Zheng Zhong, Cesar A. Stuardo, Dongyang Wang, Mohamed S. Abdelfattah, Haibin Lin, Banghua Zhu, Ang Li, Ziheng Jiang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大で高速な「思考機械」（AI モデル）を構築する工場を運営している自分を想像してください。この工場には、主に 2 種類の労働者がいます。

「コンテキスト」チーム（アテンション）: これらの労働者は探偵のようです。彼らはこれまでの物語全体を見て、何が起きているかを理解します。彼らは考えることが非常に得意ですが、あまり移動する必要はありません。
「スペシャリスト」チーム（エキスパート/FFN）: これらは実際の建設者です。彼らは何千人もいますが、単一のタスクに対しては、ほんの少数しか必要としません。彼らは重労働をこなすのが得意ですが、工場の床の至る所に散らばっています。

問題：交通渋滞

この工場を運営する従来の方法（エキスパート並列化と呼ばれる）では、コンテキストチームとスペシャリストチームが同じ部屋に閉じ込められていました。

以下のようなことが起こりました。

コンテキストチームが思考を完了します。
彼らは部屋を横切ってスペシャリストたちに叫びます。「おい、お前、お前、そしてお前！ここに来てこの仕事をしろ！」
スペシャリストたちが作業を行います。
彼らは振り返って叫びます。「結果はこちらだ！」

問題は何か？その「叫び声」（コンピュータ間のデータ送信）が遅いことです。特に工場が巨大化し、労働者が異なる階（異なるサーバーノード）にいる場合、顕著でした。スペシャリストたちは、コンテキストチームが叫び終わるのを待ってしばしば待機しており、その逆もまた然りでした。工場は交通渋滞に陥り、作業する時間よりも話す時間の方を費やしていました。

これを修正しようとした以前の試みは、労働者に「働きながら話せ」と言おうとするようなものでした。しかし、コンテキストチームは思考に長い時間がかかること（特に長い物語の場合）と、スペシャリストたちは建設が迅速であることとの間に、タイミングが決して完全に一致しませんでした。常に隠しきれない余分な「叫び声」の時間が存在していました。

解決策：DisagMoE（分離された工場）

この論文の著者であるDisagMoEは、2 つのチームを同じ部屋で働かせようとするのをやめることにしました。代わりに、2 段階のアセンブリラインを構築しました。

1. チームの分離（分離化）
彼らはコンテキストチームをあるセットの機械に移し、スペシャリストチームを全く別のセットの機械に移しました。

チーム A（コンテキスト）: すべての「探偵」レイヤーが 1 つのコンピュータグループ上にあります。
チーム B（スペシャリスト）: すべての「建設者」レイヤーが別のグループ上にあります。

2. 新しいアセンブリライン（AF-Pipe）
混沌とした叫び声の代わりに、滑らかな一方向のコンベアベルトシステムを作成しました。

コンテキストチームが思考のバッチを完了し、作業をベルトの上でスペシャリストたちへ滑らせます。
スペシャリストたちがバッチ#1 を作業している間、コンテキストチームはすでにバッチ#2 を開始しています。
スペシャリストたちがバッチ#1 を完了して結果を戻している間、コンテキストチームはすでにバッチ#3 に取り組んでいます。

これはパイプライン処理と呼ばれます。バトンの受け渡しが滑らかで、どのランナーも走るのを止めることのないリレーレースのようなものです。

3. 賢い管理者（適応的割り当て）
ここが巧妙な部分です。この論文は、コンテキストチームとスペシャリストチームには異なるニーズがあることに気づきました。

コンテキストチームは重厚な思考者です。彼らは強力な脳（計算能力）を必要としますが、高速なインターネット接続はそれほど必要としません。
スペシャリストチームはランナーのようです。彼らはデータを正しい人に迅速に届けるために、超高速のハイウェイ（ネットワーク帯域幅）を必要とします。

古い工場では、全員が同じ量のインターネットと計算能力を割り当てられていました。しかしDisagMoEでは、管理者は賢明です。物語が非常に長い場合、彼らはスペシャリストたちにより多くの「インターネットレーン」を、コンテキストチームにはより多くの「脳力」を割り当てます。どちらのチームも他方を待つことがないように、リソースを絶えず調整します。

結果

チームを分離し、管理者がリソースを調整できるようにすることで、工場はメッセージを待つ時間の浪費を止めました。

速度: 新しいシステムは、従来の標準的な方法よりもトレーニングを1.8 倍高速化することがわかりました。
効率: コンピュータが単に「話す」（データを待つ）ために費やす時間を最大 88% 削減しました。

大きな教訓

この論文は、問題を速くするために単にインターネット帯域幅を投げるだけではダメだと主張しています。AI の各部分がどのような特定の作業を行っているかを見る必要があります。「思考」と「建設」を分離し、各チームに必要な正確な量のリソースを与えることで、これらの巨大な AI モデルをはるかに、はるかに速く構築することができます。

要約すると: 彼らは、労働者を分離し、より良いコンベアベルトを構築し、ワークロードを正確にバランスさせる方法を知っている管理者を雇うことで、四角い杭を丸い穴に無理やり押し込めようとするのをやめました。

技術概要：DisagMoE

1. 問題定義

大規模な混合専門家（MoE）アーキテクチャは、専門家のサブセットをスパースに活性化させることで、兆パラメータ規模の大規模言語モデル（LLM）の学習を可能にします。しかし、これらのモデルの学習は、専門家並列化（EP）に必要とされる全対全通信により深刻なボトルネックに直面しています。モデルが複数の GPU ノードにまたがってスケールするにつれ、ノード間ネットワーク帯域幅（例：InfiniBand）の制限が通信遅延を悪化させ、学習ステップの時間の最大 50% を消費する可能性があります。

既存の最適化戦略は、主にオペレータレベルでの通信と計算のオーバーラップ（例：FFN 計算のチャンキング）やマイクロバッチパイプラインの使用に焦点を当てています。しかし、これらのアプローチには 2 つの根本的な限界があります：

残留ストール（Residual Stalls）: アテンション層と FFN 層との間の計算 - 通信比率の固有の偏りにより、完全なオーバーラップが妨げられます。アテンション計算はシーケンス長に対して二次的にスケールするのに対し、FFN と通信は線形にスケールします。これにより、特にシーケンス長が増加するにつれて隠蔽できない「通信の尾部」が生じます。
リソースの非効率性: 従来のアテンションと FFN を異なるハードウェアに分離する非集約アプローチは、各グループが完全なモデルパラメータを保持する必要があり、スケール時にメモリ不足（OOM）エラーを引き起こすか、あるいは各コンポーネントの特定の計算 - 通信特性にリソース割り当てを適応させることに失敗します。

2. 手法：DisagMoE

DisagMoE は、非集約アーキテクチャを通じてモデル配置とスケジューリングを共同で最適化することで、全対全ボトルネックを緩和するように設計された学習システムです。

2.1 非集約コンポーネント配置

アテンションと FFN 層が同じ GPU グループ上に共存する従来の専門家並列化とは異なり、DisagMoE はコンポーネントタイプによってモデルを分割します：

アテンショングループ（A-Workers）: これらのグループには、複数のトランスフォーマー層からの自己アテンション層が含まれます。これらは密なコンポーネントを複製し、データ並列化（DP）を通じて学習データを分散します。
FFN グループ（F-Workers）: これらのグループには、複数のトランスフォーマー層からの FFN 層（専門家）が含まれます。専門家は、グループ内の GPU 間で専門家並列化（EP）を通じてシャード化されます。
層割り当て: 層は、交互にグループに割り当てられます（例：グループ 0 は層 0、 $p$ 、 $2p$ を処理し、グループ 1 は層 1、 $p+1$ などを処理する）。これにより、メモリ制約下でのデバイス利用率を最大化します。

2.2 AF-Pipe：オーバーラップを伴う多段階パイプライン

DisagMoE は、通信を副次的な効果ではなく第一級の段階として扱うパイプラインスケジューリングであるAF-Pipeを導入します。

多対多（M2N）通信: 単一のノードグループ内での標準的な全対全交換の代わりに、AF-Pipe はアテンショングループと FFN グループ間の双方向の多対多トークン交換を確立します。
段階の整合: パイプラインはグループ間の境界を整合させ、従来の「ディスパッチ」と「結合」操作を単一の統合された通信段階に融合します。これにより、通信オーバーヘッドは約 $1/k$ （ $k$ は top- $k$ 専門家数）削減され、冗長なデータ転送が排除されます。
非同期実行: システムは、各ワーカー内で 3 つの協調ストリーム（フォワード、バックワード、通信）を利用します。これにより、FFN グループは利用可能なデータ上でフォワードパスを計算しながら、同時に勾配（N2M）を送信することが可能となり、ブロッキングしない実行とほぼ連続的な GPU 利用率を確保します。

2.3 適応的ワーカー割り当て

計算束縛であるアテンションと通信束縛である FFN の間の不均衡に対処するため、DisagMoE は計算 - 通信ルーフラインモデルによって導かれる適応的アロケータを採用します。

モデリング: システムは、コンポーネントがネットワーク束縛から計算束縛へ移行する転換点（ $\hat{I}$ ）をモデル化します。シーケンス長が増加するにつれ、アテンション層は FFN 層よりも早く計算ルーフに到達します。
最適化: システムは、GPU とネットワークインターフェースカード（NIC）のパーティショニングを混合整数線形計画（MILP）問題として定式化します。FFN グループには通信ルーフを上げるためにより多くのネットワーク帯域幅を、アテンショングループにはより多くの計算リソースを動的に割り当て、バブルを最小化するためにパイプライン段階をバランスさせます。
探索戦略: アルゴリズムはルーフラインモデルを使用してシード構成を生成し、その後、モデル化されていないオーバーヘッドを考慮するために測定されたステップごとの時間を用いた局所探索によってこれを洗練させます。

3. 主要な貢献

非集約コンポーネント配置: アテンションと FFN 層を不相交なワーカーグループに分離する新規アーキテクチャ。これにより、独立したスケーリングとリソース割り当てが可能になります。
AF-Pipe スケジューリング: 従来の全対全をグループ間 M2N 交換に置き換え、通信をアテンションおよび FFN 計算の両方と体系的にオーバーラップさせる多段階パイプライン。
適応的リソース割り当て: ルーフラインに導かれたメカニズムにより、アテンションと FFN グループの GPU および NIC リソースを共同でサイズ設定し、変化するシーケンス長やモデル構成全体で計算 - 通信比率を効果的にバランスさせます。

4. 実験結果

著者らは、DeepSeek-MoE、GPT-OSS-120B、Qwen3 の 3 つの MoE モデルを使用し、シーケンス長が 4K から 32K の範囲で 16 ノード（128× H800 GPU）のクラスター上で DisagMoE を評価しました。

スループットの向上: DisagMoE は、Megatron-LM（1F1B 交互）、Tutel、Comet、DualPipe を含む強力なベースラインを一貫して上回りました。
- Megatron-1F1B に対して最大1.81 倍の高速化を達成。
- 最先端の MoE オーバーラップシステムに対して最大1.34 倍の高速化を達成。
通信の削減: ベースラインと比較して、DisagMoE はオーバーラップされない通信時間を最大 88%（Tutel 対比）、75%（Comet 対比）、45%（DualPipe 対比）削減しました。
スケーラビリティ: システムは、異なる top- $k$ 値や専門家並列化（EP）サイズにおいて効果を示しました。
リソース感度: 実験により、アテンションと FFN グループ間の最適な GPU 比率はシーケンス長によって変化する（例：長いシーケンスはより多くのアテンションリソースを好む）ことが確認され、適応的割り当ての必要性が実証されました。

5. 意義と主張

本論文は、DisagMoE が現在の MoE 学習における根本的な限界、すなわちアテンションと FFN 層の構造的な不均衡に起因する通信の完全なオーバーラップの不可能性に対処することを主張しています。これらのコンポーネントを非集約し、ルーフラインモデルに基づいてリソースを適応的に割り当てることで、システムは対称的な帯域幅プロビジョニングが達成できるものを超えて、実効的なオーバーラップウィンドウを拡大します。

著者らは、このアプローチが固定されたシーケンス長を持つ大規模な事前学習ワークロードにとって特に重要であると強調しています。現在の設計は安定した形状（固定されたシーケンス長とマイクロバッチ）を前提としていますが、将来のハードウェアとソフトウェアは、単に生帯域幅をスケールさせるのではなく、コンポーネント間で有利な計算 - 通信比率をターゲットにすべきであるという核心的な洞察は、MoE 学習を効率的にスケールさせるための新たな方向性を示唆しています。このシステムは Megatron-LM 上で実装されており、以前の非集約学習の試みを悩ませたメモリ制約なしに、数百億パラメータ規模のモデルの学習に対する実用的な解決策であることが示されています。