MoE 解码的通道选择不需要校准数据(Channel Selection for Mixture-of-Experts Decoding Needs No Calibration Data)

小批量交互解码时,MoE 模型的瓶颈在于从显存读路由专家的权重。已有做法是预测专家里哪些中间通道重要、只读这些,并用离线数据校准一个阈值。作者指出阈值既不能保证实际读取量、也不稳定,改用每个词取得分最高的 k 个通道(无需校准);并给出一个完全由模型权重构造的读取器:用常驻显存的三值草图粗筛通道,选中的通道按原 BF16 权重精确计算,跳过的通道用一个只由权重得到的低秩映射补回。

小批量交互解码时,MoE 模型的瓶颈在于从显存读路由专家的权重。已有做法是预测专家里哪些中间通道重要、只读这些,并用离线数据校准一个阈值。作者指出阈值既不能保证实际读取量、也不稳定,改用每个词取得分最高的 k 个通道(无需校准);并给出一个完全由模型权重构造的读取器:用常驻显存的三值草图粗筛通道,选中的通道按原 BF16 权重精确计算,跳过的通道用一个只由权重得到的低秩映射补回。

点赞 12 · 代码星数 3

理由

  • 作者称在三个后训练 MoE 模型上,只读 17.8% 的路由专家权重就保留完整模型 98.1% 的准确率;Triton 内核让 MoE 模块在批大小 1 时比优化过的 vLLM 快 1.60–1.90 倍,更大批量时最高 2.69 倍。
  • 拆穿了校准阈值的问题:按 18% 读取预算校准的 Prox 平均实际只读了 12.9%;同一打分下 top-k 版本在 43 组比较里有 40 组持平或更好,生成任务上 23 组全部如此,且避开了校准阈值在 Gemma 4 上引发的重复循环。
  • 不需要任何校准数据:README 称每个模型只需一次性构建恢复映射,单卡几分钟。
  • 模型都是热门 MoE:Qwen3-30B-A3B-Thinking-2507、gemma-4-26B-A4B-it、Qwen3.5-35B-A3B 和 OLMoE,用官方 BF16 权重,不改模型本身。

用法

  • 代码:github.com/AidenGeunGeun/calibration-free-moe(Apache-2.0),Linux、CUDA 12.9+、Python 3.12,基于 vLLM 0.29.0;稀疏解码要把仓库 patches/ 里的两个补丁打在一份 vLLM 副本上,并把它放在 PYTHONPATH 前面。
  • 流程:scripts/build_weight_recovery_bank.py 为模型建恢复映射 → scripts/eval_vllm_generation.py 加 --bank、--retention、--selection-score point、--screen-keep 0.625 解码;各模型约 13% / 18% 读取预算对应的 retention 值 README 里有表。
  • 复现论文数字看 reproduce/README.md,规约脚本不需要 GPU;速度对比用 scripts/speed_block.py 和 speed_harness.py。

风险

  • 研究代码,不是即插即用的服务插件:要维护打过补丁的 vLLM 副本,只有解码阶段稀疏、预填充仍走原版 vLLM。
  • 只为 OLMoE、Qwen3、Qwen3.5、Gemma 4 写了专家绑定;README 称 Qwen3-Thinking 和 Gemma 4 需要一张 96 GB 卡或两张 48 GB 卡。
  • 准确率和加速数字均为作者报告,评测集中在 GPQA-Diamond、IFEval、MMLU-Pro 等基准。

打开 · 代码 · 抱脸论文页

Hugging Face Daily Papers

浏览最新 AI 动态