MaskCoFT:掩码协同微调降低 MoE 推理显存开销
MaskCoFT: Masked Co-Adaptive Fine-Tuning for Memory-Efficient MoE Inference
跑 MoE 显存不够的人可以看看:微调路由加专家,Mixtral-8x7B 每层只缓存 4 个专家,取回次数降了 23.7%,速度更快精度还涨。
MoE 模型常超出单张 GPU 显存,专家卸载方案下解码速度取决于每个 token 需要取多少专家。MaskCoFT 用可学习的二值掩码限制每层 Top-K 路由到专家子集,让 router 和专家在交叉熵损失下共同微调。推理时掩码变成软先验对专家重新排序,所有专家仍可被选中。在 Mixtral-8x7B 模拟每层缓存 4 个专家、DeepSeek-V2-Lite 缓存 12 个专家的设定下,每 token 专家取回次数分别下降 23.7% 和 10.1%,实际卸载系统中单 token 输出耗时最多降低 16.4% 和 5.5%,九个基准的平均准确率仍比基线高 0.92 和 0.53 分。
MaskCoFT: Masked Co-Adaptive Fine-Tuning for Memory-Efficient MoE Inference
Mixture-of-experts (MoE) language models often exceed the memory of a single GPU. Expert offloading keeps most experts in host memory and loads them on demand, so decoding speed depends on how many experts each token must fetch. Caching and prefetching reduce this cost only as far as the routing allows. Router-only fine-tuning can reshape the routing to reuse experts, but it keeps the experts frozen, so they cannot adapt to the tokens the new routing sends them. We propose MaskCoFT, a masked co-adaptive fine-tuning method that trains routers and experts together with the cross-entropy loss alone. During fine-tuning, a learnable binary mask restricts the Top-K routing of each layer to a subset of experts, and the experts adapt to the tokens redirected to them. At inference, the learned mask becomes a soft prior that re-ranks experts, so every expert remains selectable. We simulate a GPU cache of 4 experts per layer for Mixtral-8x7B and 12 for DeepSeek-V2-Lite. MaskCoFT cuts expert fetches per token by 23.7% and 10.1% relative to the base model. In real offloading system serving, it lowers the time per output token by up to 16.4% and 5.5%, respectively. Its average accuracy over nine benchmarks stays above the base model by 0.92 and 0.53 points.