论文

内存墙的另一半:以训练后的路由预测从SSD服务35B MoE

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

AI 基础设施模型架构MoE模型服务框架

摘要

消费级硬件上的混合专家推理受权重内存限制:一个35B级模型在4位量化下占19.5GB,稀疏性减少的是每token计算,而非必须存放的字节数。简单卸载到SSD本身无济于事,因为在第N层输出产生前必须选定第N+1层专家,读取无法足够早启动以隐藏在计算之后。我们提出流式MoE推理引擎Edge0,通过预路由器弥合这一缺口:每层预测头提前一个token预测下一层路由,预测本身直接作为路由,因此预备专家集合与实际路由集合相同,不会丢弃任何专家。在学生路径上训练且未合并的恢复LoRA,补偿int4量化与路由替换造成的质量损失。在单台24GB机器上,Edge0以3GiB峰值活跃内存、每秒20token运行35B MoE,在五个公开基准上的平均表现与fp16教师相差仅数分。同一框架也运行8B档模型,框架、检查点及适配器均已开源。