arXiv:2410.01485cs.CL2024-10ICLR被引 11

用部分上下文训练模型,实现长文本高效处理。

A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts

  • 混合全注意力与稀疏注意力结构,兼顾效率与性能。
  • 在50亿数据上微调即可将上下文扩展至128K。
  • 训练速度提升1.55倍,推理内存减少62%。

训练和部署长上下文大语言模型(LLMs)存在显著开销。传统方法需分两步:先在长上下文数据上微调以扩展上下文长度,再通过架构修改降低推理时的键值缓存(KV cache)开销。本文提出将上下文扩展与高效的KV缓存架构整合,减少训练开销并提升长上下文表现。为此,我们设计了LongGen:在扩展上下文长度的同时,将预训练模型微调为高效架构。核心洞察包括:(1)稀疏注意力模式(如窗口注意力、注意力记忆点、分块稀疏注意力)具有友好的GPU内存访问模式,可实际提升效率;(2)模型需直接访问所有令牌,采用1/3全注意力+2/3高效注意力的混合结构可平衡效率与性能;(3)仅需在50亿条长上下文数据上轻量级训练,即可将模型上下文从4K扩展至128K。我们在Llama-2 7B和70B上验证,训练时达到1.55倍加速,壁钟时间减少36%;推理时KV缓存内存减少62%,预填充速度提升1.67倍,解码速度提升1.41倍。

原文摘要 · Abstract (English)

Training and serving long-context large language models (LLMs) incurs substantial overhead. To address this, two critical steps are often required: a pretrained LLM typically undergoes a separate stage for context length extension by training on long-context data, followed by architectural modifications to reduce the overhead of KV cache during serving. This paper argues that integrating length extension with a GPU-friendly KV cache reduction architecture not only reduces training overhead during length extension, but also achieves better long-context performance. This leads to our proposed LongGen, which finetunes a pretrained LLM into an efficient architecture during length extension. LongGen builds on three key insights: (1) Sparse attention patterns, such as window attention (attending to recent tokens), attention sink (initial ones), and blockwise sparse attention (strided token blocks) are well-suited for building efficient long-context models, primarily due to their GPU-friendly memory access patterns, enabling efficiency gains not just theoretically but in practice as well. (2) It is essential for the model to have direct access to all tokens. A hybrid architecture with 1/3 full attention layers and 2/3 efficient ones achieves a balanced trade-off between efficiency and long-context performance. (3) Lightweight training on 5B long-context data is sufficient to extend the hybrid model's context length from 4K to 128K. We evaluate LongGen on both Llama-2 7B and Llama-2 70B, demonstrating its effectiveness across different scales. During training with 128K-long contexts, LongGen achieves 1.55x training speedup and reduces wall-clock time by 36%, compared to a full-attention baseline. During inference, LongGen reduces KV cache memory by 62%, achieving 1.67x prefilling speedup and 1.41x decoding speedup.

长上下文稀疏注意力高效推理微调

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。