内存墙的另一半：用训练好的路由预测从 SSD 服务 35B MoE

TL;DR Edge0 用预路由器提前预测下一层 MoE 路由，让 SSD 专家读取脱离关键路径，在 24GB 机器上以 3GiB 峰值内存跑 35B MoE 达 20 tok/s，质量接近 fp16 教师模型。

问题

问题背景

MoE 模型（如 Mixtral、DeepSeek-V3）以稀疏专家结构在增大参数容量的同时控制每 token 计算量，已成为大模型扩展的主流路线之一。但其权重内存需求远超消费级 GPU 显存，如何在单卡或边缘设备上高效推理成为热点。

现有方法局限

仅靠量化不够：35B 级 MoE 在 4-bit 下仍有 19.5GB，加上激活与 KV-cache，24GB GPU 难以完全驻留。简单 SSD offloading 存在根本缺陷——每层 top-k 专家必须在上一层输出计算完成后才能由 router 确定，SSD 读取无法提前启动，读延迟完全暴露在关键路径上，吞吐骤降至个位数 tok/s。若采用“提前随机/启发式预取”或“丢弃部分专家”，要么大量无效读取浪费带宽，要么改变路由导致模型质量下降；批量推理虽能摊薄延迟，却不适用于低延迟单流场景。

为什么这个问题难 / 重要

难点在于：路由预测必须在看到当前层输出之前完成，预测结果若直接替代真实 router 会产生分布漂移；同时 4-bit 量化误差与路由替换误差相互叠加，需要额外补偿。业界高度关注 MoE 的端侧/消费级部署，内存墙是比算力更硬的约束，软件层面对专家加载时机和路由可控性的创新直接影响单位显存的吞吐与质量下限。

行业类比

类似流媒体播放中基于用户行为预测提前缓存下一片段：预测足够可靠时，就可以用低带宽换取流畅体验；Edge0 将 MoE 参数按需从 SSD 流式加载，用训练好的 prerouter 精准预取下一层专家。

核心洞察

• Edge0 的核心创新在于 prerouter 的预测结果被直接当作路由决策，而不是传统预取机制中的辅助信号。这消除了“预测-验证-回滚”的复杂路径，使系统结构大幅简化。与投机解码或分支预测等需要检查正确性并可能回退的方案不同，Edge0 接受预测路由的轻微偏差，通过训练使预测头逼近真实路由分布，并配合 recovery LoRA 补偿误差，从而在保持推理质量的同时实现高效的 SSD 流式加载。
• 未合并的 recovery LoRA 是另一个有工程启示的设计：它在“学生路径”上训练，即量化+路由替换后的模型上，专门补偿由 int4 量化和路由预测引入的误差。与常规 LoRA 合并到主权重不同，Edge0 保持 LoRA 未合并，这允许在推理时以极低成本叠加修正，同时不影响 SSD 中权重布局和预取逻辑。这种分离设计使得模型可以灵活地在不同部署场景中切换精度与内存开销，且 LoRA 参数可以驻留在有限内存中。

方法

Edge0 以 int4 量化 MoE 权重（存储于 SSD）为输入，运行时仅驻留少量活跃参数，目标输出为高质量 token 序列。

核心模块

• SSD 流式专家加载：按层顺序从 SSD 读取专家权重到内存。由于 MoE 每层仅激活少数专家，计算量小但权重大，读取延迟成为瓶颈。
• Prerouter：每个 transformer 层引入一个轻量预测头，接收当前 token 的隐藏状态，输出下一层路由决策（top-k 专家选择）。预测被直接用作实际路由，因此系统可在处理当前 token 时就提前启动下一层所需专家的 SSD 读取，将 I/O 隐藏在计算之后。预测头通过蒸馏教师路由训练（Phase 1）。
• Recovery LoRA：在量化后的 student 路径上训练未合并的 LoRA 适配器，补偿 int4 量化和路由替换带来的质量损失。训练采用 SFT 与 on-policy distillation 相结合（Phase 2/3）。

输出与效果

在单张 24GB GPU 上，Edge0 以约 20 tok/s 速度服务 35B MoE，峰值活跃内存仅 3GiB，平均质量与 fp16 教师相差几个百分点。

与同类差异：传统 offloading 方案依赖被动预取或缓存重放，且预测路由与真实路由不一致时会丢弃专家；Edge0 将预测直接作为路由，主动驱动 SSD 预取且无丢弃，因此能完全隐藏 I/O 并保持质量。

实验

实验设计

Edge0 在单张 24 GB 显存机器上评估 35B 级 MoE 模型（4-bit 量化），采用 SSD 流式加载专家权重。核心对比路径：naive SSD offloading（无预路由）与 Edge0（prerouter 预测 + recovery LoRA）。五个公开基准（具体名称未在摘要中列出）用于衡量质量。评估指标包括生成吞吐量、峰值主动内存和相对 fp16 教师模型的平均分差。

关键发现

• 吞吐量：约 20 tok/s，在消费级硬件上首次实现 35B MoE 的可用速度。
• 峰值主动内存：仅 3 GiB，大幅低于模型静态大小（19.5 GB @4-bit）。
• 质量：平均分差在 fp16 教师几个点以内，表明 prerouter 替换路由和 int4 量化未造成显著能力损失。

与基线对比解读

naive SSD offloading 无法隐藏 SSD 读取延迟，因为下一层专家必须在当前层输出后确定，读取无法提前开始。Edge0 的 prerouter 逐层预测下一层路由，并将预测直接作为路由使用，使预取专家集合与实际路由一致，从而将加载移出关键路径。这一设计不是优化预测精度，而是利用“预测即路由”消除丢弃和重读开销。同时，recovery LoRA 未合并，保持适配器独立，便于在不同量化/路由替身上复用。

行业影响

落地场景

Edge0 使 35B MoE 模型在单块 24GB GPU 上达到约 20 tok/s，峰值活跃内存仅 3GiB。这让大规模稀疏模型得以运行在边缘设备、个人工作站和成本敏感型云实例上。具体用例：电商平台的实时客服机器人可利用 35B 模型提供专家级回答，无需多卡集群；内容平台的个性化推荐排序可部署 MoE 模型于单卡处理低延迟请求；医疗领域的辅助诊断对话可在医院本地服务器运行，保障数据隐私。

商业价值

• 降本：硬件门槛从 A100/H100 多卡或大内存机器降至消费级 24GB GPU，SSD 存储比 HBM 便宜一个数量级，总体拥有成本下降 80% 以上。
• 体验提升：在边缘设备上实现大模型实时交互，响应延迟约 50ms/token，适合对话、代码补全等场景。
• 增收：SaaS 厂商能以更低价格提供大模型 API，扩大中小客户群；私有化部署能力成为新卖点。

与现有工作流集成

Edge0 提供 框架 + 检查点 + adapter，可直接接入现有推理栈：

1. 替换 Hugging Face Transformers 中的 AutoModelForCausalLM 加载逻辑，使用 Edge0 流式执行器。
2. 预路由器与恢复 LoRA 作为附加权重，与量化基础模型一同发布，不修改原模型架构。
3. 兼容 PyTorch 与 ONNX Runtime 生态，可作为 vLLM 或 TGI 的插件，或独立服务通过 gRPC 暴露。
4. 对现有 MoE（如 Mixtral、DeepSeek MoE）只需训练轻量 prerouter head 和 LoRA adapter，即可启用 SSD 流式推理，迁移成本极低。

局限

• - **训练开销与泛化成本**：Edge0 需要为每个 MoE 模型单独训练 prerouter 头，并经历蒸馏、学生路径 SFT 和 on-policy 蒸馏三个阶段。这种非即插即用的设计对比零样本 SSD offloading 或单纯添加 LoRA 的方案显著提高了使用门槛；当需要服务多种模型族或频繁迭代版本时，训练成本与工程复杂度会成为瓶颈。论文未给出完整的训练算力与时间数据，使得实际可行性存疑。
• - **硬件假设与页缓存敏感**：系统依赖 SSD 顺序读带宽和 OS 页缓存，冷页粒度问题被明确提及。一旦工作集超过缓存或访问序列转为随机，读延迟会暴露在关键路径上，吞吐可能大幅下降。实验仅在单张 24GB 显卡搭配高性能 NVMe 的环境下验证，未涵盖速度较慢的存储介质、共享云盘或带宽受限场景；3GiB 峰值活跃内存的宣称也依赖于特定的 staging 与复用策略，缺乏跨硬件鲁棒性证据。
• - **质量换取与评测范围**：尽管 recovery LoRA 补偿了 int4 量化与路由替换损失，但预测路由在分布偏移下仍可能偏离真实 top-k 选择，论文仅在五个公开基准上报告平均分数“within a few points”，未能覆盖长尾生成、代码、数学推理或安全敏感任务。实验对象限于 35B 和 8B 两类 MoE，对更宽专家数、不同路由函数或混合架构的泛化性未经验证；相比无损 offloading 或 fp16 服务，这一 inherent 质量-成本权衡可能在严苛场景中不被接受。