Capacity-Aware Inference: Mitigating the Straggler Effect in Mixture of Experts

要約

専門家（MOE）の混合は、まばらな専門家の活性化を活用してパフォーマンスと効率のバランスをとることにより、大規模な言語モデルをスケーリングするための効果的なアーキテクチャです。
ただし、専門家の並列性の下では、MoEは不均衡なトークンから専門家への割り当てのために推論の非効率性に苦しんでいます。
この現象は、最も負担のない専門家が全体的な推論潜時を決定するため、この現象を\ textbf {\ textit {straggler effect}}と定義します。
これに対処するために、最初に\ textIT {\ textBf {容量とアウェアトークンドロップ}}を提案します。これは、過負荷の専門家から余分なトークンを破棄し、最小限のパフォーマンスへの影響を伴う負荷の不均衡を効果的に削減することにより、専門家の容量制限を強制します（$ 30 \％$ $ speedupのみ$ 0.9 \％$ degraging）。
次に、低負荷の専門家の存在が容量のしきい値をはるかに下回っていることを考えると、\ textIT {\ textBf {capational-awareの拡張ドロップ}}を導入します。これにより、トークンは厳格なローカル容量制約を強制し、負荷バランスを改善し、乱用の専門家の実用化を強化する前に、候補セットに追加のローカルエキスパートを含めることができます。
言語とマルチモーダルMOEモデルの両方での広範な実験は、私たちのアプローチの有効性を示し、専門家の利用、モデルのパフォーマンス、および推論効率の大幅な利益をもたらします。

要約(オリジナル)

The Mixture of Experts (MoE) is an effective architecture for scaling large language models by leveraging sparse expert activation to balance performance and efficiency. However, under expert parallelism, MoE suffers from inference inefficiencies due to imbalanced token-to-expert assignment, where underloaded experts complete computations early but must wait for overloaded experts, leading to global delays. We define this phenomenon as the \textbf{\textit{Straggler Effect}}, as the most burdened experts dictate the overall inference latency. To address this, we first propose \textit{\textbf{Capacity-Aware Token Drop}}, which enforces expert capacity limits by discarding excess tokens from overloaded experts, effectively reducing load imbalance with minimal performance impact (e.g., $30\%$ speedup with only $0.9\%$ degradation on OLMoE). Next, given the presence of low-load experts remaining well below the capacity threshold, we introduce \textit{\textbf{Capacity-Aware Expanded Drop}}, which allows tokens to include additional local experts in their candidate set before enforcing strict local capacity constraints, thereby improving load balance and enhancing the utilization of underused experts. Extensive experiments on both language and multimodal MoE models demonstrate the effectiveness of our approach, yielding substantial gains in expert utilization, model performance, and inference efficiency, e.g., applying Expanded Drop to Mixtral-8$\times$7B-Instruct yields a {0.2\%} average performance improvement and a {1.85$\times$} inference speedup.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google