ZeRO-Prefill：MoEプリフィル配信におけるゼロ冗長オーバーヘッド

arXiv cs.LG / 2026/5/6

📰 ニュースDeveloper Stack & InfrastructureModels & Research

共有:

要点

本論文は、分類・レコメンデーション・検証などのプリフィルのみのLLMワークロードにおけるMoEモデル配信のボトルネックが、計算量ではなく分散実行に伴うオーバーヘッドである点を扱っています。
その要因として、エキスパート配置と同期的な活性ルーティングを結び付けた設計が、オートレグレッシブなデコーディング時代から引き継がれていることを指摘しています。
ZeRO-PrefillではAsyncEPを導入し、活性のAllToAllを毎層で行うのではなく、エキスパートの重みを非同期にAllGatherで集めることで、長い計算集約型のプリフィル前向き計算と通信をオーバーラップさせます。
さらに、プレフィックスに基づくルーティングと「真のFLOPs」による負荷計測、物理的に導出した飽和閾値を組み合わせて、ルーティングの偏りを抑えます。
実験ではQwen3-235B-A22Bで、実運用ワークロードにおいて最良の分散ベースライン比で1.35〜1.37×のスループット向上、長いコンテキストの合成テストで最大1.59×を達成し、GPUあたりのモデルFLOPs利用率は29.8〜36.2%となっています。

Abstract

生成LLMのワークロードは、分類、レコメンド、検証といった判別的タスクをますます担うようになっています。これらの回答は、自己回帰的なデコードを行わずに単一のプリフィル（prefill）パスのロジットから読み取られます。混合エキスパート（MoE）モデルにおいて、こうしたプリフィルのみのワークロードを提供する処理は、計算そのものではなく、モデルを収めるために必要な分散実行によってボトルネックになります。既存の並列戦略（テンソル、エキスパート、パイプライン並列）は、冗長な計算、通信、同期と引き換えにメモリ負荷を抑えるため、その結果としてMoEプリフィル提供の効率が著しく低下します。これらのオーバーヘッドが、エキスパート配置を同期的なアクティベーション・ルーティングと結び付けることに起因していることを確認します。これは、デコーディング時代から引き継がれた設計です。大きなバッチのプリフィルにおける長く計算律速のフォワードパスは、レイヤーごとのウィンドウを十分に広くし、エキスパート重みをバックグラウンドでストリーミングできるようにします。これにより、レイヤーごとのアクティベーション AllToAll を、計算と完全にオーバーラップした非同期の重み AllGather で置き換えます。私たちは、プリフィルのみの提供システムである ZeRO-Prefill を提案します。そのバックエンドである AsyncEP（非同期エキスパート並列）は、アクティベーションに基づいてエキスパートをルーティングするのではなく、重みによってエキスパートを収集します。フロントエンドでは、プレフィックス対応のルーティングと true-FLOPs による負荷追跡を通じて、物理的に導出された飽和閾値を共に強制します。Qwen3-235B-A22B を4つのハードウェア／精度構成で評価したところ、ZeRO-Prefill は現実世界のワークロードにおいて最も強力な分散ベースラインを 1.35〜1.37x 上回るスループットを示し、長いコンテキストの合成ワークロードでは最大 1.59x を達成しました。さらに、29.8〜36.2% の per-GPU モデル FLOPs 利用率を維持します。