arXiv:2505.16710cs.LGcs.AI2025-05ACL被引 7

通过分块优化,让长文本大模型训练更省显存、更快

Training Long-Context LLMs Efficiently via Chunk-wise Optimization

  • 将长文本分块处理,每块独立计算和反向传播,只存一个块的激活值
  • 分块梯度选择性传播,使训练时间随序列变长趋于推理时间,最多快3倍
  • 适合在单张显卡上微调大模型,尤其对长文档处理任务有实用价值

长上下文大语言模型虽具备出色的文档处理能力,但高昂的训练成本常阻碍其定制化应用。为此,我们提出一种内存高效的训练范式——顺序分块优化(SeCO),将长输入划分为可管理的块,每个块独立构建计算图并执行局部反向传播,仅需存储一个块的前向激活值。在此基础上,进一步引入稀疏分块优化(SpaCO),通过选择性地向特定块传播梯度,并加入精心设计的补偿因子,确保梯度估计无偏。SpaCO将反向传播的计算开销与上下文长度解耦,使训练时间随序列增长逐渐趋近于推理时间。二者作为轻量级训练封装实现,显著提升实用性:例如在单张RTX 3090 GPU上微调8B模型(使用LoRA)时,SeCO将最大序列长度从1K扩展至16K;SpaCO在相同条件下训练速度最高可达SeCO的3倍。这些创新为优化长上下文模型提供了新思路,使其更适用于实际场景。代码已开源。

原文摘要 · Abstract (English)

While long-context large language models (LLMs) exhibit remarkable document processing capabilities, their prohibitively high training costs often hinder customized applications. To mitigate this issue, we propose \textit{Sequential Chunk-wise Optimization} (SeCO), a memory-efficient training paradigm that partitions lengthy inputs into manageable chunks. Each chunk independently constructs its computational graph and performs localized backpropagation, ensuring that only one chunk's forward activations are stored in memory. Building on SeCO, we further introduce \textit{Sparse Chunk-wise Optimization} (SpaCO), which reduces computational overhead by selectively propagating gradients to specific chunks and incorporates a carefully designed compensation factor to ensure unbiased gradient estimation. SpaCO decouples the computational cost of backpropagation from the context length, enabling training time to gradually converge to inference time as sequences become longer. Implemented as lightweight training wrappers, both SeCO and SpaCO offer substantial practical benefits. For example, when fine-tuning an 8B model with LoRA on a single RTX 3090 GPU, SeCO expands maximum sequence length from 1K to 16K tokens, while SpaCO demonstrates accelerated training speed -- achieving up to 3x faster than SeCO under the same experimental setup. These innovations provide new insights into optimizing long-context models, making them more accessible for practical applications. We have open-sourced the code at \href{https://github.com/wenhaoli-xmu/seco}{here}.

长文本生成模型训练显存优化

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