# SparseBalance: Load-Balanced Long Context Training with Dynamic Sparse Attention

arXiv: https://arxiv.org/abs/2604.13847 · HTML: https://mlsanity.com/paper/2604.13847/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 34 references that are in the map.
- First posted on arXiv 15 Apr 2026.

## Builds on (most cited first, 23 of 23)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ChatGLM: A Family of Large Language Models from GLM-130B to GLM-4 All Tools](https://mlsanity.com/paper/2406.12793/) (2024)
- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference](https://mlsanity.com/paper/2406.10774/) (2024)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [SWIFT:A Scalable lightWeight Infrastructure for Fine-Tuning](https://mlsanity.com/paper/2408.05517/) (2024)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models](https://mlsanity.com/paper/2309.12307/) (2023)
- [MoBA: Mixture of Block Attention for Long-Context LLMs](https://mlsanity.com/paper/2502.13189/) (2025)
- [LongAlign: A Recipe for Long Context Alignment of Large Language Models](https://mlsanity.com/paper/2401.18058/) (2024)
- [ChatQA 2: Bridging the Gap to Proprietary LLMs in Long Context and RAG Capabilities](https://mlsanity.com/paper/2407.14482/) (2024)
- [Twilight: Adaptive Attention Sparsity with Hierarchical Top-p Pruning](https://mlsanity.com/paper/2502.02770/) (2025)
- [WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training](https://mlsanity.com/paper/2503.17924/) (2025)
- [Efficient Long Context Fine-tuning with Chunk Flow](https://mlsanity.com/paper/2503.02356/) (2025)
- [Skrull: Towards Efficient Long Context Fine-tuning through Dynamic Data Scheduling](https://mlsanity.com/paper/2505.19609/) (2025)

## Abstract (arXiv)

While sparse attention mitigates the computational bottleneck of long-context LLM training, its distributed training process exhibits extreme heterogeneity in both textit1) sequence length and textit2) sparsity sensitivity, leading to a severe imbalance problem and sub-optimal model accuracy. Existing algorithms and training frameworks typically focus on single issue, failing to systematically co-optimize these two problems. Therefore, we propose SparseBalance, a novel algorithm-system co-design framework, which exploits the sparsity and sequence heterogeneity to optimize model accuracy and system efficiency jointly. First, we propose workload-aware dynamic sparsity tuning, which employs a bidirectional sparsity adjustment to eliminate stragglers and exploit inherent bubbles for free accuracy. Second, we propose a sparsity-aware batching strategy to achieve coarse-grained balance, which complements dynamic sparsity tuning. Experimental results demonstrate that SparseBalance achieves up to a 1.33× end-to-end speedup while still improving the long-context capability by 0.46% on the LongBench benchmark.