MOM：面向长上下文语言模型的内存高效卸载最小序列推理
机器学习
2025-04-18 v1 人工智能
计算与语言
摘要
长上下文语言模型虽在推理阶段因高额 GPU 内存需求而难以部署，但展现出惊人的性能。我们提出一种名为 Memory-efficient Offloaded Mini-sequence Inference（MOM）的内存高效卸载最小序列推理方法，该方法将关键层划分为较小的“最小序列”，并无缝集成于 KV 缓存卸载机制。实验在各种 Llama、Qwen 和 Mistral 模型上均表明，MOM 在平均情况下将峰值内存使用降低超过 50%。在单张 A100 80GB GPU 上，MOM 将 Meta-Llama-3.2-8B 模型的最大上下文长度从 155k 扩展至 455k token，同时保持输出完全一致且不牺牲准确性。MOM 由于最小计算开销和高效的最后一层处理，仍保持高度竞争的吞吐量。相较于传统的分块预填充方法，MOM 实现的上下文长度扩展幅度提高 35%。更重要的是，我们的方法显著降低了预填充阶段的内存消耗，彻底消除了长时间推理中长期以来的主要内存瓶颈。这一突破性成果从根本上改变了研究优先级，将未来的研究重心从预填充阶段的优化，转向改进解码阶段残余 KV 缓存的效率。
引用
@article{arxiv.2504.12526,
title = {MOM: Memory-Efficient Offloaded Mini-Sequence Inference for Long Context Language Models},
author = {Junyang Zhang and Tianyi Zhu and Cheng Luo and Anima Anandkumar},
journal= {arXiv preprint arXiv:2504.12526},
year = {2025}
}
备注
Submitted to COLM