论文

FlexMoE:MoE 语言模型的一对一嵌套专家内部剪枝

FlexMoE: One-for-All Nested Intra-Expert Pruning for MoE Language Models

摘要

专家混合 (MoE) 语言模型通过稀疏激活的专家扩展模型能力,使该架构成为现代大型模型的标准配方。然而,稀疏激活并不能消除存储和服务所有专家的部署负担,并且可用的部署预算可能因设备、用户和工作负载的不同而有很大差异。现有的 MoE 压缩方法在很大程度上仍然是固定预算,通常在每个选定的目标预算下优化一个压缩端点。我们研究了一种不同的设置:将大型预训练 MoE LLM 转换为跨预算的可部署子网络的嵌套系列。我们的方法首先根据专家 FFN 通道的重要性对其进行排名,然后让每个专家学习一个离散的操作来修剪其通道。通过逐渐增加成本压力,单个动作训练运行会导出一系列从高预算到低预算的动作掩码,每个动作掩码都标识嵌套在排序基础模型中的可靠的较小子网络。此外,我们在中间修剪预算(40%)处使用单个恢复微调来恢复退化的模型质量并将恢复的模型转移到其他看不见的预算。总的来说,我们的框架超越了最近的 MoE 压缩基准。具体来说,在 Qwen2-57B-A14B 上,即使没有 微调,我们的方法也保留了约 99.8% 的基本性能,同时修剪了 50% 的路由专家参数。对于部署而言,我们修剪的子网络可实现真正的内存减少和吞吐量增益,并通过内核级协同设计进一步支持实时在线预算切换。