# FarSkip-Collective: Unhobbling Blocking Communication in Mixture of Experts Models > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 3 (May 20 / Wed)、Grand Ballroom 1、16:00 - 16:15 PDT > - **登壇者:** Yonatan Dukler, Guihong Li, Deval Shah, Jiang Liu, Vikram Appia, Emad Barsoum (Advanced Micro Devices Inc. (AMD)) > - **URL:** https://mlsys.org/virtual/2026/oral/3743 > - **OpenReview:** https://openreview.net/forum?id=ruOpvLzsGV > - **コード:** https://github.com/AMD-AGI/FarSkip-Collective > [!abstract] 概要 > MoE モデルの分散実行におけるブロッキング通信を解消する手法 FarSkip-Collective を提案する。モデルアーキテクチャのサブブロック間の接続関係を変更し、通信コレクティブの完了を待たずに利用可能なアクティベーション(部分的または古いもの)で次の計算を開始することで、計算と通信のオーバーラップを実現する。自己蒸留レシピ FCSD により 16B から 109B 規模の最先端 MoE を FarSkip-Collective モデルへ変換し、元モデルの精度を平均 1% 以内で維持できることを示した。推論では DeepSeek-V3 アーキテクチャの変換モデルを SGLang で提供し、プリフィル時に最大 32.6% の TTFT 高速化と 97.3% の通信オーバーラップを達成した。訓練では Megatron-LM 上で all-to-all 通信の 88.4% をオーバーラップさせた。 ## テーゼ MoE モデルの分散実行で不可避とされてきたブロッキング通信(Expert Parallelism の Dispatch/Combine に伴う all-to-all)は、アーキテクチャレベルでサブブロック間の残差接続の依存関係を解きほぐすことで、計算と完全にオーバーラップさせることができる。この変更はモデルのパラメータ配置や密な計算カーネルを一切変えずに接続のみを修正するため、既存チェックポイントから自己蒸留で効率的に変換可能であり、100B 超規模でも精度劣化は 1% 以内に収まる。 ## 問題設定 MoE モデルは条件付き計算によりトークンあたりの FLOPS を削減し、総パラメータ数 500B 超への拡張を可能にしている。しかし分散実行ではアクティベーションと重みの通信が必要であり、特に Expert Parallelism における Dispatch(トークンをエキスパートへ送信する all-to-all)と Combine(結果を集約する all-to-all)はブロッキングで実行される。通信の完了を待つ間アクセラレータは遊休状態(exposed idle time)となり、ハードウェア演算性能の向上に伴い通信がエンドツーエンドワークロードに占める割合は拡大の一途である。 既存の通信オーバーラップ手法はテンソル並列の dense モデル向けに設計されており、MoE 固有の Dispatch/Combine パターンには対応していない。また小規模モデルや部分的なレイヤー変換にとどまっており、100B 超の最先端 MoE に全レイヤー適用して精度を維持できるかは未解明であった。 ## 提案手法 ### FarSkip-Collective アーキテクチャ 通常のトランスフォーマーの残差接続 $o_{k+1} = o_k + f_{k+1}(o_k)$ では、$f_k$ の出力に通信を含む場合 $o_k$ が確定するまで $f_{k+1}$ を開始できない。FarSkip-Collective はこれを $o_{k+1} = o_k + f_{k+1}(o_k^*)$ に変更する。$o_k^*$ は通信完了前に利用可能なアクティベーションであり、2 種類の定式化がある。 - **outdated(古い)**: $o_k^* = o_{k-1}$(前のレイヤーの残差をそのまま使用) - **partial(部分的)**: $o_k^* = o_{k-1} + f_k^*(o_{k-1}^*)$(通信前に確定済みの計算結果を加算) MoE レイヤーへの具体的な適用では、アテンションサブブロックには partial(アテンション出力と共有エキスパート出力を含むが、ルーティングエキスパート出力を含まない)、MLP サブブロックには outdated(前レイヤーの残差)を入力として使用する。これにより Dispatch、Combine、アテンション後通信の 3 つのブロッキング通信すべてをオーバーラップ可能にしつつ、各入力は最大 1 レイヤー分だけ古い状態に留まる。 ### FarSkip-Collective Self-Distill (FCSD) アーキテクチャ変更はパラメータ配置を変えないため、既存チェックポイントをそのまま読み込んで接続のみ変更できる。しかし再訓練なしでは全レイヤー変換時に精度が大幅に劣化する。FCSD は元モデルを教師とする自己蒸留で精度を回復するレシピである。 - KL ベースのロジット蒸留を使用し、元モデルのロジットに合わせる - 中間表現の L2 損失も併用可能($\mathcal{L}_{L2}$) - 10B トークン未満のインストラクションチューニングデータで訓練(元モデルのスクラッチ事前訓練の約 100--1000 倍安価) - MBPP+ を用いた早期停止で訓練不安定性に対処 ### 明示的オーバーラップ実装 PyTorch レイヤーでの汎用的な実装を設計方針とし、低レベルカーネル最適化を避けている。 **訓練(Megatron-LM / Primus):** `torch.dist` の `async_op=True` と `torch.cuda.Stream()` を活用して非ブロッキングスケジューリングを実現する。後方パスでは PyTorch の `torch.autograd` の Sequence Number 機構をハイジャックし、勾配ノードの実行順序を制御して非同期の後方通信を計算と十分にオーバーラップさせる。 **推論(vLLM / SGLang):** vLLM と SGLang では Dispatch/Combine の代わりに all-reduce を使用するため、アテンション出力と共有エキスパート出力の通信を統合した単一の all-reduce に遅延・結合し、ルーティングエキスパート計算とオーバーラップさせる。HIP/CUDA グラフとの統合には PyNCCL の直接バインディングを使用する。 ## 実験・評価 ### モデル精度(蒸留) FCSD により 3 つの最先端オープンソース MoE の全レイヤーを変換した結果(スライド p.26 の正規化精度より): | モデル | 規模 | 正規化平均精度 | |---|---|---| | DeepSeek-V2-Lite | 16B-A3B | 96% | | Qwen-3-30B-MoE | 30B-A3B | 97% | | Llama 4 Scout | 109B-A17B | 99% | Llama 4 Scout ではインストラクションチューニング版と比較して平均精度の低下が 1% 以内である。SFT による変換と比較して FCSD は一貫して優れた性能を示し、特に生成タスク(HumanEval+ など)での破局的忘却を回避する。 事前訓練からのアブレーション(DeepSeek-V2-Lite アーキテクチャ、200B トークン)では、通常アーキテクチャとの平均精度差が 0.3%(51.5% 対 51.2%)であり、アーキテクチャ自体の表現力が維持されることを示した。 ### 推論性能(DeepSeek-V3 アーキテクチャ、SGLang) スライド p.28 のベンチマーク結果: - **プリフィル時間(BS=32、EP=8/TP=8):** コンテキストサイズ 256 で 1.18 倍、512 で 1.24 倍、1024 で 1.32 倍、2048 で 1.31 倍の高速化 - **プリフィルスループット(Context x BS = 65536、EP=8/TP=8):** 1.31--1.34 倍の向上 - **デコード時間(BS=1024、WideEP: EP=16/TP=16):** コンテキストサイズ 256 で 1.25 倍、512 で 1.24 倍、1024 で 1.22 倍の高速化 - **通信オーバーラップ率:** プリフィルステージで 97.3% 論文の報告値では、変換済み Llama 4 Scout モデルの推論で TTFT が 18.5% 高速化、DeepSeek-V3 変換モデルでは TTFT が 32.6% 高速化した。 ### 訓練性能(DeepSeek-V2-Lite、Megatron-LM、EP8) - all-to-all 通信の 88.4% を計算とオーバーラップ(フォワード 87.6%、バックワード 89.0%) - フォワードの最初の all-to-all とバックワードの最後の all-to-all はオーバーラップ候補が存在しないため除外 ### スパース化への展望 解析モデルによる理論的プリフィル高速化(コンテキスト 2048)は、スパーシティ係数が増すほど大きくなる。DeepSeek-V3 アーキテクチャ(約 32 倍スパース)では約 1.5 倍の高速化が理論値として示され、64 倍スパースでは約 1.65 倍に達する。より疎な MoE の実用化を後押しする技術である。 ## 結論・オープン課題 FarSkip-Collective は MoE モデルのブロッキング通信問題に対し、アーキテクチャの接続修正という原理的に明快なアプローチで、訓練と推論の双方で大幅な高速化を実現した。FCSD レシピにより既存の大規模 MoE チェックポイントを低コストで変換でき、100B 超規模で精度劣化 1% 以内を達成した点は実用上の意義が大きい。 残された課題として以下が挙げられる: - **精度ギャップの完全解消:** Llama 4 Scout の MMLU では FCSD で 4.1 ポイントの差が残る(SFT では 10.3 ポイント)。データ混合やモデルマージング等の追加手法による改善が今後の課題である - **マルチブロックスキッピング:** 通信時間が 1 サブブロックの計算時間を超える極端にスパースな MoE や異なるハードウェアパラダイムへの対応として、複数ブロックをまたぐスキップは将来の発展方向である - **訓練不安定性:** FCSD の KL 蒸留では教師と生徒の微小な不一致が大きな勾配を生じ、訓練が不安定化する場合がある。早期停止で対処しているが、根本的な解決策は未確立である