推測的エキスパートが Mixture-of-Experts の推論を加速する

Abstract

Mixture-of-Experts（MoE）モデルは、巨大言語モデル（LLMs）の容量を拡張しつつ、疎な活性化とトークンあたりの計算量を低減させる手段として人気を集めている。 しかし、メモリ制約のある推論設定では、エキスパートの重みをCPUにオフロードする必要があり、デコード時のCPU–GPU間転送が性能のボトルネックとなる。プリフェッチ方式を用いたエキスパートの事前取得を提案し、現在計算済みの内部モデル表現を活用して将来のエキスパートを推測することで、メモリ転送を計算と重ね合わせられる。 複数のMoEアーキテクチャにわたり、これらの内部表現によって将来のエキスパートを信頼性高く予測できることを示す。 また、推測されたエキスパートを実行することで、下流タスクの精度を一般に維持できることを示しており、真のルータ選択エキスパートを再取得する必要をなくすことで、より効果的な計算－メモリの重ね合わせを維持する。 最適化された推論エンジンに統合された本手法は、CPUメモリからエキスパートをオンデマンドでロードする場合と比較して、出力トークンあたりの時間（TPOT）を最大14％削減する。 推測実行だけでは最適な精度を得られないMoEについては、エキスパート予測のヒット率を改善する軽量推定量をさらに検討し、性能低下を低減する。 私たちのコードはオープンソースとして、https://github.com/axonn-ai/yalis/tree/offload_prefetch に公開されています。