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

📄 arXiv: 2410.19123v1 📥 PDF

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

分类: cs.CL, cs.LG

发布日期: 2024-10-24

备注: 38th Conference on Neural Information Processing Systems (NeurIPS 2024)

🔗 代码/项目: GITHUB

💡 一句话要点

Read-ME：通过系统协同设计，将LLM重构为解耦路由的混合专家模型，提升推理效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 模型压缩 推理加速 激活稀疏性 预训练模型 路由算法 系统协同设计

📋 核心要点

  1. 现有MoE模型推理效率低，内存管理和批处理策略与模型架构不匹配，且从头训练成本高昂。
  2. Read-ME框架通过激活稀疏性从预训练的密集LLM中提取专家，避免了昂贵的从头训练。
  3. Read-ME采用解耦的预门控路由器，支持预计算和前瞻调度，优化专家感知的批处理和缓存，提升推理效率。

📝 摘要（中文）

大型语言模型（LLM）的普及推动了混合专家（MoE）架构的应用，该架构动态利用专门的子网络来提高效率和性能。尽管MoE模型具有优势，但在推理过程中面临着严峻的挑战，包括由于模型架构和系统策略之间的设计选择不一致导致的低效内存管理和次优批处理。此外，从头开始训练MoE的传统方法在成本方面越来越令人望而却步。在本文中，我们提出了一个新颖的框架Read-ME，该框架将预训练的密集LLM转换为更小的MoE模型（与“升级改造”通用MoE相反），从而避免了从头开始训练的高昂成本。我们的方法采用激活稀疏性来提取专家。为了组合专家，我们研究了广泛采用的逐层路由设计，并表明其存在冗余，因此我们引入了与MoE骨干解耦的预门控路由器，该路由器有助于系统友好的预计算和前瞻调度，从而增强了专家感知的批处理和缓存。因此，我们的协同设计解决了算法和系统方面的关键差距，为资源受限环境中的LLM推理建立了一种可扩展且高效的替代方案。Read-ME优于其他类似规模的流行开源密集模型，在MMLU上实现了高达10.1%的改进，并将平均端到端延迟提高了高达6.1%。

🔬 方法详解

问题定义：现有混合专家模型（MoE）在推理时存在效率问题，主要体现在内存管理和批处理策略上。这些问题源于模型架构与系统策略的不匹配。此外，从头开始训练大型MoE模型的成本非常高昂，限制了其应用范围。

核心思路：Read-ME的核心思路是将预训练的密集LLM转化为更小、更高效的MoE模型，而不是从头开始训练。通过激活稀疏性来识别和提取专家，并设计一个解耦的预门控路由器来优化专家间的调度和批处理。这种方法旨在降低训练成本，同时提高推理效率。

技术框架：Read-ME框架主要包含两个阶段：专家提取和路由优化。首先，利用激活稀疏性从预训练的密集LLM中提取专家。然后，引入一个与MoE骨干解耦的预门控路由器，该路由器负责将输入路由到相应的专家。该路由器支持预计算和前瞻调度，从而优化专家感知的批处理和缓存。

关键创新：Read-ME的关键创新在于解耦的预门控路由器设计。传统的逐层路由设计存在冗余，而Read-ME的解耦路由器可以独立于MoE骨干进行预计算，从而实现更高效的调度和批处理。此外，Read-ME避免了从头训练MoE模型，大大降低了训练成本。

关键设计：Read-ME的关键设计包括激活稀疏性的阈值选择，预门控路由器的结构和训练方式，以及专家感知的批处理和缓存策略。具体的参数设置和损失函数选择未知，需要参考论文的详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Read-ME在MMLU基准测试中，性能优于其他类似规模的开源密集模型，取得了高达10.1%的提升。同时，平均端到端延迟降低了高达6.1%。这些实验结果表明，Read-ME在提高LLM推理效率方面具有显著优势。

🎯 应用场景

Read-ME框架适用于资源受限环境下的LLM推理，例如边缘设备或低功耗服务器。它可以降低LLM的部署成本，提高推理速度，从而促进LLM在移动应用、智能助手等领域的广泛应用。该方法也为LLM的轻量化和高效部署提供了新的思路。

📄 摘要（原文）

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.