読んでください: LLMをルーターから切り離されたエキスパートの混合物としてリファクタリングするシステム共同設計

Read-ME: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design

October 24, 2024

著者: Ruisi Cai, Yeonju Ro, Geon-Woo Kim, Peihao Wang, Babak Ehteshami Bejnordi, Aditya Akella, Zhangyang Wang

cs.AI

要旨

大規模言語モデル（LLM）の普及により、専門化されたサブネットワークを動的に活用するMixture-of-Experts（MoE）アーキテクチャが採用され、効率と性能が向上しました。その利点にもかかわらず、MoEモデルは推論中に効率的なメモリ管理やサブ最適なバッチ処理など、重要な課題に直面しています。これは、モデルアーキテクチャとシステムポリシーとの設計の不一致に起因しています。さらに、MoEをゼロからトレーニングする従来のアプローチは、コスト面でますます制約が大きくなっています。本論文では、私たちは、事前トレーニングされた密なLLMを小さなMoEモデルに変換する新しいフレームワークであるRead-MEを提案しています（一般的なMoEを「アップサイクリング」するのとは対照的）。私たちのアプローチは、専門家を抽出するために活性化スパースネスを利用しています。専門家を構成するために、広く採用されている層別のルーターデザインを調査し、その冗長性を示し、MoEバックボーンから分離されたプリゲーティングルーターを導入しています。これにより、システムにやさしい事前計算と先読みスケジューリングが可能となり、専門家を意識したバッチ処理とキャッシングが向上します。したがって、私たちの共同設計は、アルゴリズムとシステムの両面での重要なギャップに対処し、リソース制約のある環境でのLLM推論のためのスケーラブルで効率的な代替手段を確立しています。Read-MEは、同様の規模の他の人気のあるオープンソースの密なモデルを上回り、MMLUで最大10.1％の改善を達成し、平均エンドツーエンドの遅延を最大6.1％改善しています。コードは次で入手可能：https://github.com/VITA-Group/READ-ME.

English

The proliferation of large language models (LLMs) has led to the adoption of Mixture-of-Experts (MoE) architectures that dynamically leverage specialized subnetworks for improved efficiency and performance. Despite their benefits, MoE models face significant challenges during inference, including inefficient memory management and suboptimal batching, due to misaligned design choices between the model architecture and the system policies. Furthermore, the conventional approach of training MoEs from scratch is increasingly prohibitive in terms of cost. In this paper, we propose a novel framework Read-ME that transforms pre-trained dense LLMs into smaller MoE models (in contrast to "upcycling" generalist MoEs), avoiding the high costs of ground-up training. Our approach employs activation sparsity to extract experts. To compose experts, we examine the widely-adopted layer-wise router design and show its redundancy, and thus we introduce the pre-gating router decoupled from the MoE backbone that facilitates system-friendly pre-computing and lookahead scheduling, enhancing expert-aware batching and caching. Our codesign therefore addresses critical gaps on both the algorithmic and system fronts, establishing a scalable and efficient alternative for LLM inference in resource-constrained settings. Read-ME outperforms other popular open-source dense models of similar scales, achieving improvements of up to 10.1% on MMLU, and improving mean end-to-end latency up to 6.1%. Codes are available at: https://github.com/VITA-Group/READ-ME.