突破内存瓶颈：面向百万token上下文的高效LLM训练系统
计算与语言
2026-03-03 v4
摘要
在长序列上训练大型语言模型（LLM）受限于昂贵的GPU内存开销，而非训练时间。主要问题在于激活值，其内存占用随序列长度线性增长。我们引入OOMB（Out of the Memory Barrier）系统，直接面对这一瓶颈。我们的方案采用块循环训练框架，配合即时激活重计算，保持恒定的激活内存占用（O(1)），将主要瓶颈转移到不断增长的KV缓存。为管理KV缓存，OOMB集成了一套协同的优化措施：为KV缓存及其梯度提供分页内存管理，以消除碎片化；采用异步CPU卸载隐藏数据传输延迟；以及采用块级稀疏注意力机制，降低计算复杂度和通信开销。这些技术的协同作用带来了卓越的效率。我们的实证结果表明，对于每增加10K token的上下文，端到端训练内存开销仅增加约10MB（以Qwen2.5-7B为例）。这使得在单张H200 GPU上训练Qwen2.5-7B模型即可支持400万token的上下文，而若不使用本方法，则需要大型集群进行上下文并行。本工作在长上下文LLM训练资源效率方面代表了一项重大进展。源代码已发布于https://github.com/wenhaoli-xmu/OOMB。
引用
@article{arxiv.2602.02108,
title = {Out of the Memory Barrier: A Highly Memory Efficient Training System for LLMs with Million-Token Contexts},
author = {Wenhao Li and Daohai Yu and Gen Luo and Yuxin Zhang and Fei Chao and Rongrong Ji and Yifan Wu and Jiaxin Liu and Ziyang Gong and Zimu Liao},
journal= {arXiv preprint arXiv:2602.02108},
year = {2026}
}