LeMo: Enabling LEss Token Involvement for MOre Context Fine-tuning

TL;DR

LeMo addresses the activation-memory bottleneck in long-context fine-tuning by introducing Contextual Token Sparsity, a token-level sparsity mechanism where token informativeness varies across inputs and layers. It deploys three core techniques—Information-driven Token Elimination, Context-aware Pattern Prediction, and High-performance Kernel Optimization—to identify, predict, and exploit informative tokens while avoiding Shadowy Activation. Empirical results show up to $1.93\times$ memory savings and up to $1.36\times$ speedups over state-of-the-art methods, with the ability to handle sequence lengths of $s=32K$–$64K$ on a single GPU and maintain competitive accuracy. The approach is compatible with common LLM architectures, extensible via Two-dimensional Sparsity and Sparsity-sensitive Offload, and demonstrates strong scalability across multi-GPU platforms, offering a practical path to more efficient long-context fine-tuning.

Abstract

The escalating demand for long-context applications has intensified the necessity of extending the LLM context windows. Despite recent fine-tuning approaches successfully expanding context lengths, their high memory footprints, especially for activations, present a critical practical limitation. Current parameter-efficient fine-tuning methods prioritize reducing parameter update overhead over addressing activation memory constraints. Similarly, existing sparsity mechanisms improve computational efficiency but overlook activation memory optimization due to the phenomenon of Shadowy Activation. In this paper, we propose LeMo, the first LLM fine-tuning system that explores and exploits a new token-level sparsity mechanism inherent in long-context scenarios, termed Contextual Token Sparsity. LeMo minimizes redundant token involvement by assessing the informativeness of token embeddings while preserving model accuracy. Specifically, LeMo introduces three key techniques: (1) Token Elimination, dynamically identifying and excluding redundant tokens across varying inputs and layers. (2) Pattern Prediction, utilizing well-trained predictors to approximate token sparsity patterns with minimal overhead. (3) Kernel Optimization, employing permutation-free and segment-based strategies to boost system performance. We implement LeMo as an end-to-end fine-tuning system compatible with various LLM architectures and other optimization techniques. Comprehensive evaluations demonstrate that LeMo reduces memory consumption by up to 1.93x and achieves up to 1.36x speedups, outperforming state-of-the-art fine-tuning systems.

Figures (20)

• Figure 1: Illustration of shadowy activation. (a) LoRA performs full attention computation without incorporating sparsity. (b) LongLoRA adopts two shifted local attention patterns (colored in red and green) to approximate full attention, targeting sparsity at the hidden-dimension level. (c) LeMo employs token-level sparsity, minimizing token involvement and achieving activation savings compared to other methods.
• Figure 2: Memory breakdown during LLM mixed-precision fine-tuning. Compared to consistent model states, activations scale with both input batch size and sequence length, becoming the primary bottleneck in long-context scenarios.
• Figure 3: An illustrated example of LoRA for showcasing the activation memory usage of PEFT. Beyond the shared parts, LoRA requires storing more activations than vanilla.
• Figure 4: Visualization of attention scores across different models, datasets, layers, and sequences (darker is higher).
• Figure 5: LeMo overview. At each layer, token embeddings are first partitioned into blocks and fed into the pattern predictors (❷). Using the predicted informativeness scores from these predictors, the token elimination algorithm (❶) effectively identifies and retains only the most informative tokens for processing. Optimized kernels (❸) then efficiently perform token selection, computation, residual addition, and padding.