MaskCoFT：面向内存高效 MoE 推理的掩码协同自适应微调

arXiv:2609.34077v2 公告类型：替换 摘要：混合专家（MoE）语言模型通常会超出单个 GPU 的内存容量。专家卸载将大多数专家保存在主机内存中，并按需加载，因此解码速度取决于每个 token 需要获取多少个专家。缓存和预取只能在路由机制允许的范围内降低这一开销。仅对路由器进行微调可以重塑路由，使其复用专家，但专家保持冻结状态，因而无法适应新路由发送给它们的 token。我们提出 MaskCoFT，这是一种掩码协同自适应微调方法，仅使用交叉熵损失共同训练路由器和专家。在微调过程中，一个可学习的二值掩码将每一层的 Top-K 路由限制在专家子集内，同时专家适应被重新定向给它们的 token。在推理时，学习到的掩码会成为一种软先验，用于重新排序专家，因此每个专家仍然可以被选中。我们分别针对 Mixtral-8x7B 的每层 4 个专家和 DeepSeek-V2-Lite 的每层 12 个专家模拟 GPU 缓存。相较于基础模型，MaskCoFT 将每个 token 的专家获取次数分别减少了 23.7% 和 10.1%。在真实的卸载系统服务中，它将每个输出 token 的耗时分别最多降低了 16.4% 和 5.5%。在九项基准测试中的平均准确率分别比基础模型高出 0.92 和 0.53 个百分点。