专家推测加速混合专家模型推理
Speculating Experts Accelerates Inference for Mixture-of-Experts
AI总结:
本文提出一种专家预取方案,利用当前计算的内部表示推测未来专家,减少CPU-GPU传输开销,实现14%的TPOT提升,并通过轻量估计器优化预测精度。
摘要:
混合专家(MoE)模型因其在扩展大语言模型容量的同时保持稀疏激活和降低每token计算而受到青睐。然而,在内存受限的推理设置中,专家权重必须卸载到CPU,导致解码过程中CPU-GPU传输成为性能瓶颈。我们提出了一种专家预取方案,利用当前计算的内部模型表示来推测未来专家,使内存传输与计算重叠。在多种MoE架构中,我们证明这些内部表示可以可靠地预测未来专家。我们还证明执行推测专家通常能保持下游任务的准确性,从而通过消除重新获取真实路由选择专家的需要,保留更多的有效计算-内存重叠。在优化的推理引擎中,我们的方法实现了相对于从CPU内存中按需加载专家的14%的TPOT减少。对于那些仅靠推测执行无法达到最佳精度的MoE,我们进一步研究了轻量估计器,以提高专家预测命中率,从而减少性能下降。我们的代码在https://github.com/axonn-ai/yalis/tree/offload_prefetch开源发布。
英文摘要:
Mixture-of-Experts (MoE) models have gained popularity as a means of scaling the capacity of large language models (LLMs) while maintaining sparse activations and reduced per-token compute. However, in memory-constrained inference settings, expert weights must be offloaded to CPU, creating a performance bottleneck from CPU-GPU transfers during decoding. We propose an expert prefetching scheme that leverages currently computed internal model representations to speculate future experts, enabling memory transfers to overlap with computation. Across multiple MoE architectures, we demonstrate that future experts can be reliably predicted by these internal representations. We also demonstrate that executing speculated experts generally maintains downstream task accuracy, thus preserving more effective compute-memory overlap by eliminating the need to re-fetch true router-selected experts. Integrated into an optimized inference engine, our approach achieves up to 14\% reduction in time per output token (TPOT) over on-demand loading of experts from CPU memory. For MoEs where speculative execution alone yields suboptimal accuracy, we further examine lightweight estimators that improve expert prediction hit rates, thereby reducing performance degradation. Our code is released in open-source at https://github.com/axonn-ai/yalis/tree/offload_prefetch.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。