Back to AI Research

AI Research

Long-Context Fine-Tuning with Limited VRAM | AI Research

Key Takeaways

  • Long-Context Fine-Tuning with Limited VRAM Training large language models on long sequences is typically limited by the massive amount of memory required to...
  • Parameter-efficient fine-tuning reduces model and optimizer memory, but dense attention still makes long training sequences expensive.
  • We combine Hierarchical Global Attention (HGA) with segment-wise backpropagation and tiered KV storage.
  • Only the active segment remains differentiable in VRAM; older KV is detached into RAM or NVMe, and HGA loads a bounded set of exact historical tokens for each query block.
  • On Qwen3-8B with 4-bit QLoRA and PG19, dense training on a 16 GB Quadro RTX 5000 fits 2,048 tokens but fails at 4,096, whereas HGA reaches 16,384 tokens with 15.28 GB peak VRAM.
Paper AbstractExpand

Parameter-efficient fine-tuning reduces model and optimizer memory, but dense attention still makes long training sequences expensive. We combine Hierarchical Global Attention (HGA) with segment-wise backpropagation and tiered KV storage. Only the active segment remains differentiable in VRAM; older KV is detached into RAM or NVMe, and HGA loads a bounded set of exact historical tokens for each query block. On Qwen3-8B with 4-bit QLoRA and PG19, dense training on a 16 GB Quadro RTX 5000 fits 2,048 tokens but fails at 4,096, whereas HGA reaches 16,384 tokens with 15.28 GB peak VRAM. Under evaluation the same adapter runs through 131,072 tokens on this card; VRAM is not constant but grows gently with the resident chunk summaries, so RAM and NVMe capacity set the practical limit beyond these lengths. At the shared 2K training length, HGA-trained and dense-trained adapters obtain 2.7405 and 2.7383 nat under the same dense-attention readout, while the stock model obtains 2.9541. At this boundary HGA training is already marginally faster (217.75 vs. 207.02 tokens/s), and the HGA-to-dense throughput ratio improves from 1K to 2K; because HGA keeps the attended historical set per token approximately constant while dense work per token grows, we expect this lead to widen as context grows. Dense attention is used for the main quality and retrieval comparisons so that they measure the learned weights and remain compatible with standard generation frameworks. HGA can also be used for retrieval and generation; an optimized production-grade serving implementation is under development.

Long-Context Fine-Tuning with Limited VRAM
Training large language models on long sequences is typically limited by the massive amount of memory required to store attention states and gradients. Even with parameter-efficient methods like QLoRA, dense attention mechanisms force the GPU to process the entire history at once, leading to "out of memory" (OOM) errors as context lengths grow. This paper introduces a combination of Hierarchical Global Attention (HGA), segment-wise backpropagation, and tiered memory storage to allow fine-tuning on significantly longer sequences without requiring more VRAM.

How the approach works

The core strategy is to decouple the GPU's active working memory from the total sequence length. Instead of keeping the entire history in VRAM, the system uses a tiered storage approach: only the current segment being trained remains in VRAM, while older key-value (KV) states are offloaded to system RAM or NVMe storage.
To manage attention efficiently, HGA uses a two-level routing system. It creates compact summaries of historical "chunks" (64 tokens) and "groups" (8 tokens). When the model processes a new query, it uses these summaries to select only the most relevant historical tokens to load into the GPU. This ensures that the amount of work per token remains relatively constant, even as the total context length increases. By using truncated backpropagation through time (TBPTT), the system ensures that gradients are only calculated for the active segment, preventing memory bloat.

Key performance results

The researchers tested this method using the Qwen3-8B model on a 16 GB Quadro RTX 5000 GPU. While standard dense training failed at 4,096 tokens, the HGA-based approach successfully trained on sequences up to 16,384 tokens while staying within the same 16 GB VRAM limit.
At a shared 2,048-token length, HGA-trained adapters performed nearly identically to those trained with dense attention, proving that the routing process does not degrade the quality of the learned weights. Furthermore, because HGA keeps the attention workload per token stable, it is expected to become faster than dense attention as the context length grows, as it avoids the exponential increase in computational work associated with standard attention mechanisms.

Important considerations

While this method effectively scales training, there are a few limitations to keep in mind. The current implementation is designed for fine-tuning rather than training models from scratch; after approximately 100 million training tokens, the model may begin to exhibit "causal leakage," where it indirectly learns information about future tokens due to the way routing decisions are shared across a chunk.
Additionally, while the research demonstrates successful training up to 16,384 tokens and evaluation up to 131,072 tokens, the practical limit is ultimately determined by the capacity of the host RAM or NVMe storage used to hold the historical data. Finally, while the resulting adapters are compatible with standard dense-attention inference frameworks, an optimized production-grade engine for HGA-based serving is still under development.

Comments (0)

No comments yet

Be the first to share your thoughts!