Training Long-Context LLMs Efficiently via Chunk-wise Optimization

TL;DR

This paper tackles the high memory and compute costs of fine-tuning long-context LLMs by introducing Sequential Chunk-wise Optimization (SeCO), which partitions long inputs into chunks and backprops each chunk with a localized graph to keep forward activations memory-footprint constant. Building on SeCO, Sparse Chunk-wise Optimization (SpaCO) sparsifies backpropagation across chunks while applying a compensation factor to preserve unbiased gradient estimates, decoupling backpropagation cost from context length. Empirical results show substantial practical gains: SeCO and SpaCO dramatically reduce memory usage (SpaCO up to >4x vs gradient checkpointing) and can extend context length (e.g., from 1K to 16K tokens on a single RTX 3090), with SpaCO offering up to 3x faster training than SeCO under the same setup and negligible performance loss with proper hyperparameter tuning. The work provides actionable, open-source tools for efficient long-context LLM training, highlighting the trade-offs between memory, compute, and gradient accuracy and offering guidance for balancing the three in real-world scenarios.

Abstract

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}.

Figures (7)

• Figure 1: (Left) Computational graph for chunk-wise optimization with $k=4$ chunks. The dense connections among KV caches (green arrows) complicate memory management, leading popular training frameworks deepspeedaccelerate to rely on end-to-end parallel training. (Right) By analyzing the topology of this graph, we propose SeCO, a bootstrapping method leveraging gradient checkpointing along the sequence dimension. SeCO ensures that only the computational graph of a single chunk is stored at any time.
• Figure 2: Qualitative comparison of different methods. (i) SeCO achieves significant memory reduction while maintaining time efficiency comparable to layer-level gradient checkpointing. (ii) Building upon SeCO, SpaCO significantly reduces computational overhead, making the training time converges to inference time as the sequence length expands.
• Figure 3: A visualized illustration of Algorithm \ref{['alg:1']}. Stage 1 corresponds to the first for-loop, generating KV caches for all data chunks through inference-mode, serving as checkpoints. Stage 2 corresponds to the second for-loop, where the computational graph is constructed and localized backpropagation is performed.
• Figure 4: In the Transformer architecture, the gradient flow traverses through at most a number of KV caches equal to the layer depth. This observation was also highlighted in Transformer-XL tsfm-xl.
• Figure 5: SpaCO sparsifies the gradient flow among KV caches. (Left) The original graph. $k=6$. (Right) Only the gradient flow between $t=4$ chunks is retained. By adding a factor $k/t$ to each path, the gradient computed from this sparse graph remains an unbiased estimate.