训练路由把35B MoE专家留在SSD上,峰值内存3GiB跑到20 tok/s

The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin, Yiming Wang, Runyuan Cai, Hanze Liu, Xiaodong Zeng

cs.AI

2026-09-16

Edge0用逐层prerouter提前一个token预测专家路由。35B MoE在24GB机器上跑到20.4 tok/s、峰值2.9GiB,五项公开基准平均落后fp16教师3.9分,比常驻权重快五倍。

这篇在解决什么

MoE 推理在消费级机器上卡的是权重体积,不是计算。35B 级模型 int4 仍有 19.5 GB。稀疏只缩小每个 token 算多少,不缩小必须放在某处的字节。量化到 4 bit 已经是能用的底线,再低误差会把模型毁掉。

把专家卸到 SSD 也救不了 decode:第 N+1 层选哪些专家,要等第 N 层输出出来,磁盘读藏不到计算后面。Mac mini M4 Pro 24 GB 上,整份 4-bit 权重常驻的 mlx-lm 只跑到 3.9 tok/s,还占 18.2 GiB,操作系统和 KV 几乎没地方。

方法

Edge0 把专家权重留在 SSD,按需 mmap 进有界流式池。核心有三处。

预测即路由。 第 N 层的小头用当前 token 的 post-attention 隐状态,预测第 N+1 层在下一个 token 上的路由。decode 时 MoE 层直接吃这份预测,不再跑原 router。预取集合和真正激活集合按构造相等,没有 drop,也没有 fallback 加载。提前量必须是整 token:同一 token 内、attention 后立刻选下一层,会在他们的引擎里每步同步 30–100 ms,比能藏住的加载还贵。

头是 fc1 → erf-gelu → fc2,加一条从下一层 router 权重热启动的线性残差。输入拼上本层本步和上一步的 top-k one-hot。35B 档 33 个头(owner 6–38),32/40 层走预测。

不合并的 recovery LoRA。 int4 基座冻结,低秩适配器在学生路径上蒸馏 fp16 教师,同时补偿量化和路由替换。合并后再量化回 4-bit 会把适配器效果抹掉:注意力投影只剩 34%,稠密投影剩 2%,logits 层只剩 18%。并行 delta 只要 42 MB,decode 时间测不出差别。

训练三阶段,基座始终冻在 4-bit 反量化的 bf16 上:先蒸馏路由头,再在学生路径上对约 200 万行教师文本做 SFT,最后用学生自己的生成做 on-policy 蒸馏(reverse KL)。只蒸头、不做 SFT 的检查点会重复崩溃。

35B 发布档把原模型的 K=8 收到 K=4,换适配器文件即可,不必重训基座。

结果

吞吐在 Mac mini M4 Pro 24 GB 上测。质量用 OpenCompass,和 fp16 基座同设置对照。

方法吞吐峰值内存
mlx-lm 常驻 4-bit3.9 tok/s18.2 GiB
Edge0 35B K=420.4 tok/s2.9 GiB
Edge0 8B28 tok/s1.5 GiB

24 GB 这台机器上关掉 prerouter 仍有 19.9 tok/s,页缓存已经能把热专家留下。真正拉开的是权重放不下的 16 GB MacBook M2:相对纯按需加载,K=2/4/8 分别从 4.8、3.5、1.8 tok/s 提到 8.6、6.4、3.3 tok/s,增益 +80%、+82%、+84%。K=4 时主线程堵在专家加载上的时间从 244.0 ms 降到 101.9 ms。

基准edge0-35bQwen3.6 fp16edge0-8bLing 3.0 tiny fp16
AIME 202686.692.763.373.3
HumanEval90.995.191.592.7
GPQA-Diamond79.881.870.771.2
MMLU-Pro81.084.670.165.8
IFBench57.961.753.960.6
平均79.283.269.972.7

逐项平均差距 35B 为 3.9 分、8B 为 2.8 分。8B 的 MMLU-Pro 比教师高 4.3 分。

为什么重要

消费级机器跑 35B MoE 的常规答案是买内存或上云。Edge0 给出第三条:专家放 SSD,用训练好的提前路由把磁盘延迟藏进 decode,质量账单用一份不合并的 LoRA 付。框架、检查点、适配器都开源。

这是 Apple Silicon 上单机、单请求的方案。MLX 是唯一实现,CUDA 槽位还停在架构层。宽度 K 同时动速度和内存:同一缓存预算下 K=8 收到 K=4,decode 从 3.3 提到 6.4 tok/s。

页缓存热、盘快时,别指望 prerouter 再给 80%。24 GB 机器上它几乎不加吞吐,它赚钱的场景是权重根本塞不进内存。

局限与存疑

单请求 FIFO,没有批处理。批会改变专家工作集,现有 profile 覆盖不了。decode 的地板在 CPU 构图,40 层前向每步 44 ms,优化存储搬不走这块。相邻 token 只共享约四分之一专家,预取经常读一次付两次。

质量缺口集中在长链推理:AIME 上 35B 掉 6.1 分,8B 掉 10.0 分。其余 8B 基准都在 6.7 分以内。单次基准在这套硬件上有 ±40% 的 run-to-run 波动,跨 session 能到 2.3×,所以论文只用同 session 轮换 A/B 的中位数。

术语

原文与代码

相关论文

全部论文解读