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

Hongtao Xu; Jianchao Tan; Yuxuan Hu; Pengju Lu; Hongyu Wang; Pingwei Sun; Yerui Sun; Yuchen Xie

cs.LGApr 15, 2026

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

Authors: Hongtao Xu, Jianchao Tan, Yuxuan Hu, Pengju Lu, Hongyu Wang, Pingwei Sun, Yerui Sun, Yuchen Xie,

Organizations: School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences, Beijing, China · Meituan, Beijing, China · University of Chinese Academy of Sciences, Beijing, China

Abstract

While sparse attention mitigates the computational bottleneck of long-context LLM training, its distributed training process exhibits extreme heterogeneity in both 1) sequence length and 2) 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 $\times$ end-to-end speedup while still improving the long-context capability by 0.46% on the LongBench benchmark.

Figures & tables

Fig. 1 : Extreme heterogeneity in sequence length distributions across two real-world long-context datasets. The x-axis represents the sequence length, and the y-axis denotes its corresponding proportion within the entire dataset.

Fig. 2 : Illustration of the straggler effect caused by workload imbalance. Within PP group, a bottleneck micro-batch (e.g., Micro-batch 4) dictates the critical path and exacerbates pipeline bubbles. Furthermore, DP synchronization significantly amplifies this imbalance across the entire system.

Fig. 3 : Forward-pass latency of micro-batches over continuous training iterations. The substantial gap between the maximum (

T_{max}

, red solid line), mean (

T_{mean}

, blue dashed line), and minimum (

T_{min}

, green dashed line) latencies reveals severe workload imbalance.

Fig. 4 : Cumulative score coverage under different attention budgets.

Fig. 5 : Overview of SparseBalance, which consists of three components: workload-aware dynamic sparsity tuning (DST), sparsity-aware batching (SAB) and latency prediction module.

Fig. 6 : Practical latency of sparse attention cannot be reliably predicted by either theoretical FLOPs or sequence length alone, motivating our profiling-based latency predictor.

Fig. 7 : Overall performance and step-by-step speedups on two clusters and two datasets. We evaluate the normalized speedup ratios (measured by average iteration time) of SparseBalance and its two main components for specific model settings (from left to right: Qwen2.5-0.5B, Qwen2.5-3B, and Qwen2.5-3B in LoRA). Results demonstrate significant efficiency improvements in SparseBalance and confirm the effectiveness of each component.

Category	Dataset	Baseline SFT	Min 0.1	Min 0.2	Mean 0.1	Mean 0.2
Single-Doc QA	NarrativeQA	22.42	23.62	22.88	22.89	22.03
	Qasper	32.25	35.77	36.62	32.99	34.88
	MultiFieldQA-en	42.06	42.19	41.54	42.92	42.06
Multi-Doc QA	HotpotQA	44.81	46.80	45.65	45.40	47.25
	2WikiMQA	33.42	33.17	33.98	35.62	33.20
	MuSiQue	24.30	26.71	25.50	26.51	25.78

TABLE I : Evaluation results on the LongBench benchmark. The best results for each dataset are highlighted in bold .

Fig. 8 : Sensitivity of end-to-end speedup with respect to the sparsity threshold

p

and execution anchor choice. The upper end of each bar corresponds to Min , the lower end corresponds to Max , and the horizontal line denotes Mean .

Fig. 9 : Sensitivity of end-to-end speedup to different micro-batch sizes (

mbs=1

4

, corresponding to global batch sizes

gbs=8

32

). SparseBalance consistently achieves clear speedups across all tested settings.

Config.	Imbalance $\downarrow$	Iter. (ms)	Overheads (ms)	Speedup $\uparrow$
Baseline	1.81 $\times$	4202.09	0.00	1.00 $\times$
+DST	1.73 $\times$	3901.57	37.14	1.08 $\times$
+SAB	1.41 $\times$	3470.73	1.81	1.21 $\times$
+LBB	1.58 $\times$	3421.98	0.00	1.23 $\times$
+SAB+DST	1.34 $\times$	3107.83	44.23	1.35 $\times$
+LBB+DST	1.41 $\times$	3272.88	39.01	1.28 $\times$

TABLE II : LBB denotes a length-based batching strategy. Imbalance measures the average micro-batch-level imbalance degree (lower is better). Overheads report the additional cost introduced by SparseBalance.

Fig. 10 : Training loss comparison between the MoBA baseline and SparseBalance with different dynamic tuning configurations.

Model	$\leq$ 8K	16K	32K	64K
Baseline SFT	100.00	100.00	100.00	97.60
Min 0.1	100.00	100.00	99.93	96.33
Min 0.2	100.00	100.00	99.93	95.20
Mean 0.1	100.00	100.00	100.00	97.53
Mean 0.2	100.00	100.00	99.93	96.47

TABLE III : Average accuracy (%) of Needle-in-a-Haystack (NIAH) tasks evaluated on Qwen-2.5-3B model. Results are averaged across the niah_single_1 , niah_single_2 , and niah_single_3 evaluations for each model.

Task	Base-SFT	Min 0.1	Min 0.2	Mean 0.1	Mean 0.2
ARC-Chal.	0.4727	0.4735	0.4684	0.4744	0.4659
ARC-Easy	0.7332	0.7348	0.7277	0.7315	0.7260
BoolQ	0.7936	0.7994	0.7951	0.8003	0.8031
HellaSwag	0.7356	0.7354	0.7355	0.7349	0.7366
Lambada	0.6598	0.6583	0.6594	0.6604	0.6592
Piqa	0.7884	0.7867	0.7894	0.7911	0.7900

TABLE IV : General benchmark results on Qwen2.5-3B. Bold indicates the best performance.

Fig. 11 : Trade-off between training speedup and downstream model performance under different configurations in SparseBalance.

Sparse Attention

Sparse attention reduces the token interactions computed by Transformers to lower the cost of long-context modeling and inference.

See all 243 papers

AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism
Apr 29, 2026Ahan Gupta, Zhihao Wang, Neel Dani +3Efficient Language Model Training LLM Training
SSAIL Lab, University of Illinois Urbana-Champaign Champaign, IL, USA · Anyscale San Francisco, CA, USA · Snowflake San Mateo, CA, USA
Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps
May 16, 2026Yanke Zhou, Yiduo Li, Hanlin Tang +6Self-Attention Long-Context Language Model Inference
Nanjing University · Alibaba Group
Concertina: Data-Centric Adaptive Pipeline Parallelism for Efficient Heterogeneous Long-Context LLM Training
Sep 25, 2025Shiju Wang, Yujie Wang, Fangcheng Fu +6Long-Context Language Modeling Efficient Language Model Training
Tsinghua University · Peking University · Shanghai Jiao Tong University +1

Apr 29, 2026J/K move · Enter open · S save

AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism

Ahan Gupta, Zhihao Wang, Neel Dani, Masahiro Tanaka+2

SSAIL Lab, University of Illinois Urbana-Champaign Champaign, IL, USA · Anyscale San Francisco, CA, USA · Snowflake San Mateo, CA, USA

Efficient Language Model Training LLM Training

Open paper AskPDF arXiv

Large-language-models (LLMs) demonstrate enormous utility in long-context tasks which require processing prompts that consist of tens to hundreds of thousands of tokens. However, existing LLM training libraries do not provide easy to use abstractions to optimize for long-context training, instead focusing on optimizations for models with large parameter counts through ZeRO-3/FSDP, Tensor and Pipeline parallelism. This forces users to rewrite LLM training libraries to incorporate compositions of various complex long-context optimizations, such as sequence-parallelism, to training pipelines; a process that requires in-depth expertise, reducing developer productivity. To tackle these challenges, we introduce AutoSP: the first automated solution to automatically optimize LLM training for longer-contexts. AutoSP compiles models and applies a targeted set of optimizations: automated sequence parallelism, and long-context aware activation-checkpointing, to drastically enhance LLM trainability at negligible cost to throughput. Our evaluation demonstrates AutoSP's capability on both NVIDIA and AMD hardware, increasing training contexts by upto 2.7

\times

and 2.5

\times

respectively over competitive hand-written baseline at negligible cost to runtime performance.

Abstract

Sparse Attention

Explore similar work

AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism

Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps

Concertina: Data-Centric Adaptive Pipeline Parallelism for Efficient Heterogeneous Long-Context LLM Training