突破长文本训练内存瓶颈,单卡实现百万级上下文训练
Out of the Memory Barrier: A Highly Memory Efficient Training System for LLMs with Million-Token Contexts
- 采用分块循环训练+即时重计算,激活内存恒定不变
- 每增加1万词元仅增10MB内存,400万词元上下文可单卡训练
- 适合需要超长上下文的模型训练,如文档理解、代码生成
在长上下文上训练大语言模型受限于高昂的显存开销,而非训练时间。主要瓶颈在于激活值随序列长度线性增长。我们提出OOMB系统,通过分块循环训练框架与即时激活重计算,将激活内存占用控制在常数级别(O(1)),使主瓶颈转向不断增长的键值缓存。为管理键值缓存,OOMB集成多项协同优化:页式内存管理器减少碎片,异步CPU卸载隐藏数据传输延迟,页级稀疏注意力降低计算复杂度与通信开销。这些技术协同带来显著效率提升:对于Qwen2.5-7B模型,每增加1万词元上下文,端到端训练内存仅增加10MB。这使得单张H200 GPU即可训练400万词元上下文的Qwen2.5-7B,而此前需大规模集群配合上下文并行。该工作大幅提升了长上下文大模型训练的资源效率。源码见https://github.com/wenhaoli-xmu/OOMB。
原文摘要 · Abstract (English)
Training Large Language Models (LLMs) on long contexts is severely constrained by prohibitive GPU memory overhead, not training time. The primary culprits are the activations, whose memory footprints scale linearly with sequence length. We introduce OOMB, a highly memory-efficient training system that directly confronts this barrier. Our approach employs a chunk-recurrent training framework with on-the-fly activation recomputation, which maintains a constant activation memory footprint (O(1)) and shifts the primary bottleneck to the growing KV cache. To manage the KV cache, OOMB integrates a suite of synergistic optimizations: a paged memory manager for both the KV cache and its gradients to eliminate fragmentation, asynchronous CPU offloading to hide data transfer latency, and page-level sparse attention to reduce both computational complexity and communication overhead. The synergy of these techniques yields exceptional efficiency. Our empirical results show that for every additional 10K tokens of context, the end-to-end training memory overhead increases by a mere 10MB for Qwen2.5-7B. This allows training Qwen2.5-7B with a 4M-token context on a single H200 GPU, a feat that would otherwise require a large cluster using context parallelism. This work represents a substantial advance in resource efficiency for long-context LLM training. The source code is available at https://github.com/wenhaoli-xmu/OOMB.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。