Mixture-of-experts(MoE)推理在消费级硬件上受限于权重存储:一个 35 B 规模的模型在 4 bit 量化下约为 19.5 GB,稀疏化只降低每个 token 的计算量,却不减少必须保留的字节数。直接把权重 offload 到 SSD 并不能解决问题,因为第 N+1 层的专家必须在第 N 层输出产生之前就确定,导致读取无法提前进行以隐藏在计算后面。
Edge0 是一种流式 MoE 推理引擎,通过 prerouter 解决上述瓶颈。prerouter 为每一层生成一个 head,提前预测下一层的路由信息,预测提前一个 token 输出。预测结果直接作为真实路由使用,使得预先调度的专家集合恰好等于实际路由的集合,避免了任何专家被丢弃。
为了弥补 int4 量化和路由替换带来的精度损失,Edge0 采用了 unmerged recovery LoRA,在学生路径上进行额外训练,恢复了与 fp16 教师模型相近的质量。
在一台配备 24 GB 内存的机器上,Edge0 能以约 20 token/s 的吞吐率运行 35 B MoE,峰值活跃内存仅 3 GiB,整体性能在五个公开基准上仅比 fp16 教师模型低几个百分点。相同框架下也支持 8 B 规模模型,全部代码、检查点和适配器均已开源。
点评:Edge0 通过前置路由预测和 LoRA 恢复,实现了在极低显存下高效运行大规模 MoE,展示了 SSD 与模型并行的可行路径,为消费级 AI 推理提供了新思路。