Рефакторинг 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 с нуля становится все более запретительным с точки зрения затрат. В данной статье мы предлагаем новую структуру Read-ME, которая преобразует предварительно обученные плотные LLM в более компактные модели MoE (в отличие от "переработки" общих MoE), избегая высоких затрат на обучение с нуля. Наш подход использует разреженность активации для извлечения экспертов. Для формирования экспертов мы исследуем широко принятый дизайн маршрутизатора на уровне слоя и показываем его избыточность, и вводим предварительный маршрутизатор, отделенный от основы MoE, который облегчает дружественное к системе предварительное вычисление и планирование на основе просмотра вперед, улучшая пакетизацию и кэширование, осведомленные об экспертах. Наше совместное проектирование поэтому решает критические пробелы как на алгоритмическом, так и на системном фронтах, устанавливая масштабируемую и эффективную альтернативу для вывода LLM в ресурсоограниченных средах. Read-ME превосходит другие популярные открытые плотные модели аналогичных масштабов, достигая улучшений до 10,1% по MMLU и улучшая среднюю задержку от начала до конца до 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.