Turn Waste into Worth: Rectifying Top-$k$ Router of MoE

要約

Sparse Mixture of Experts (MoE) モデルは、計算効率が高いため、大規模な言語モデルのトレーニングに人気があります。
ただし、一般的に使用される上位 $k$ ルーティング メカニズムでは、不均衡なルーティングにより冗長計算とメモリ コストが発生します。
一部のエキスパートはオーバーフローし、超過したトークンがドロップされます。
一部の専門家は空白であり、ゼロが埋め込まれており、モデルのパフォーマンスに悪影響を及ぼします。
ドロップされたトークンとパディングに対処するために、GPU 内調整とフィルイン調整で構成される Rectify-Router を提案します。
GPU 内調整はドロップされたトークンを処理し、GPU 間の通信を回避するためにトークンが配置されている GPU 内のエキスパートにトークンを効率的にルーティングします。
フィルイン修正は、パディング トークンをルーティング スコアの高いトークンに置き換えることでパディングに対処します。
私たちの実験結果は、GPU 内調整とフィルイン調整がそれぞれドロップされたトークンとパディングを効果的に処理することを示しています。
さらに、これらを組み合わせることで、バニラのトップ 1 ルーターの精度を 4.7% 上回る優れたパフォーマンスを実現します。

要約(オリジナル)

Sparse Mixture of Experts (MoE) models are popular for training large language models due to their computational efficiency. However, the commonly used top-$k$ routing mechanism suffers from redundancy computation and memory costs due to the unbalanced routing. Some experts are overflow, where the exceeding tokens are dropped. While some experts are vacant, which are padded with zeros, negatively impacting model performance. To address the dropped tokens and padding, we propose the Rectify-Router, comprising the Intra-GPU Rectification and the Fill-in Rectification. The Intra-GPU Rectification handles dropped tokens, efficiently routing them to experts within the GPU where they are located to avoid inter-GPU communication. The Fill-in Rectification addresses padding by replacing padding tokens with the tokens that have high routing scores. Our experimental results demonstrate that the Intra-GPU Rectification and the Fill-in Rectification effectively handle dropped tokens and padding, respectively. Furthermore, the combination of them achieves superior performance, surpassing the accuracy of the vanilla top-1 router by 4.7%.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google