SparseBalance: Load-Balanced Long Context Training with Dynamic Sparse Attention Hongtao Xu1,2 , Jianchao Tan2 , Yuxuan Hu2 , Pengju Lu1 , Hongyu Wang1 , Pingwei Sun2 , Yerui Sun2 , Yuchen Xie2 , Xunliang Cai2 , Mingzhen Li3,* , and Weile Jia3,*
arXiv:2604.13847v1 [cs.LG] 15 Apr 2026
1
School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences, Beijing, China 2 Meituan, Beijing, China 3 University of Chinese Academy of Sciences, Beijing, China
training, significantly degrading the system efficiency. Existing works attempt to alleviate this issue through batching or packing strategies. However, these methods can hardly solve the problem completely due to the inherently discrete nature in this bin-packing problem, especially for highly skewed datasets. This imbalance manifests across multiple parallelism dimensions in distributed training. In inner-level pipeline parallelism (PP), varying workloads exacerbate pipeline bubbles, and the processing time of extremely long sequences essentially dominates the overall pipeline latency. Worse still, this imbalance in PP ultimately propagates to the outer-level parallelism like data parallelism (DP), severely amplifying the degradation of system efficiency. Meanwhile, sparse training introduces another critical heterogeneity in sparsity sensitivity. Although existing trainable sparse attention algorithms typically adopt a predefined sparsity degree, we observe distinct sparsity sensitivities across different sequences and transformer layers. Furthermore, if simply tranI. I NTRODUCTION sitioned to dynamic sparse training without system awareness, The long-context modeling capability is increasingly crucial this inherent sparsity heterogeneity would lead to the exact same for the evolution of Large Language Models (LLMs), serving workload imbalance problem. Therefore, these two dimensions as the backbone for advanced downstream applications such of heterogeneity—sequence length and sparsity sensitivity—are as code generation, in-depth reasoning, and autonomous agent deeply intertwined and jointly impact the runtime workload, systems. Mainstream LLMs typically perform an additional highlighting the critical need for algorithm-system co-design. training stage on specific long-context datasets to extend the However, current algorithms and training frameworks typically context windows [1], [2], [3], [4]. However, as the context address these challenges in isolation, solely focusing on length increases, the standard attention mechanism becomes the workload balance or developing sophisticated sparse algorithms. primary computational bottleneck. Because standard attention Consequently, they fail to systematically co-optimize these two requires computing affinities across all token pairs, it exhibits issues, resulting in sub-optimal performance in either training a quadratic computational complexity with respect to sequence efficiency or model accuracy. length, thus severely hindering the context expansion. To mitiTo address sequence length heterogeneity and spargate this, sparse attention has emerged as a promising solution sity sensitivity heterogeneity simultaneously, we propose in recent released LLMs [5], [6]. By selectively computing SparseBalance, a novel algorithm–system co-design enabling the critical tokens, sparse attention reduces the computational bidirectional sparsity tuning for LLM training. SparseBalance complexity and breaks the performance bottleneck. jointly optimizes model accuracy and system efficiency while However, long-context sparse training still faces a severe considering their inherent trade-off. load imbalance problem caused by the inherent heterogeneity First, we propose a workload-aware Dynamic Sparsity in sequence length distributions. For instance, the Qwen2.5 Tuning (DST) strategy, which performs bidirectional sparsity technical report [3] discloses their data mix strategy at long adjustment at runtime to rebalance the workload at the layer context training stage, which comprises 40% long sequences level. Specifically, DST identifies bottleneck micro-batches and 60% short sequences. This heterogeneity in sequence and reduces their attention budget to accelerate execution, length leads to severe load imbalance problem in distributed while simultaneously increasing the attention budget of non-
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 coarsegrained 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 longcontext capability by 0.46% on the LongBench benchmark. Index Terms—Sparse Attention, Distributed Training, Load Imbalance, Large Language Model
bottleneck micro-batches to exploit pipeline bubbles for free accuracy improvements. To ensure the efficiency and accuracy during such tuning process, we introduce an anchor-guided thresholding mechanism, which determines the direction and bounds the magnitude for each micro-batch. Second, to fully unleash the optimization potential of DST, we propose a Sparsity-Aware Batching (SAB) strategy, which involves lightweight sparsity estimation and latency-based data packing. SAB provides a well-balanced initial workload distribution, which serves as a foundation for the fine-grained runtime adjustments in DST. Additionally, we implement the latency prediction module, which maps the sequence length and sparsity to practical execution latency through offline profiling, providing the accurate performance guide for both DST and SAB module. Together, SAB and DST form a unified optimization pipeline, spanning from coarse-grained data reorganization to fine-grained runtime tuning, enabling efficient load balancing without sacrificing model accuracy. We evaluate SparseBalance on two real-world datasets and conduct comprehensive evaluations across three downstream benchmarks. Our key contributions are as follows: • We provide a novel perspective on improving sparse attention training through algorithm-system co-design to jointly optimize system efficiency and model accuracy. • We propose workload-aware dynamic sparsity tuning (DST) to dynamically rebalance the training workload at runtime, while preserving model accuracy. • We propose sparsity-aware batching (SAB) dedicated for sparse training scenario, it involves a lightweight sparsity estimator and latency-based batching strategy, providing coarse-grained balance for DST. • We implement SparseBalance and experimental results demonstrate that SparseBalance improves the end-to-end training efficiency by up to 1.33× while still improving the model’s downstream capability.
pipeline bubble problem. This issue is severely exacerbated under highly heterogeneous workloads, as any execution delay in a single stage cascades throughout the entire pipeline. Tensor Parallelism (TP) [12] divides the tensor operations across devices, and each device handles a slice of the tensor operations. With TP, each GPU only has part of the input and parameters, resulting in intensive communication during training. Therefore, TP is typically applied within a single node, while other levels of parallelism are applied across nodes. Sequence Parallelism (SP) partitions the input tensor along the sequence length dimension [13], [14], [15], [16]. There are three types of SP according to the communication schemes: (i) ring-based point-to-point communications [15], (ii) AllToAllbased communications [14], and (iii) AllGather-based communications [16]. Additionally, there is another type of SP, which is proposed by Megatron-LM and splits the dropout and normalization module activation [13]. In this paper, we use this Megatron-style SP as default. B. Self-Attention and Trainable Sparse Attention Modern LLMs are typically built by stacking many Transformer layers, each of which contains a self-attention module as the core mechanism to capture complex dependencies between elements within a sequence. In long-context training, the self-attention module becomes the dominant computational bottleneck because its cost scales quadratically with the sequence length. Given the query (Q), key (K), and value (V ) matrices, where Q, K, V ∈ RN ×d with N denoting the sequence length and d representing the hidden dimension. The self-attention computation is formulated as follows: QK ⊤ √ Attention(Q, K, V ) = softmax V (1) d
In this formulation, the dot product QK ⊤ calculates the raw affinity scores, representing the alignment between sequence elements. These scores are subsequently scaled by √1d . Then, the softmax operation is applied row-wise to yield a normalized II. BACKGROUND distribution of attention weights. Finally, these weights are A. Distributed Training Strategies multiplied by the value matrix V to aggregate the relevant Data Parallelism (DP) [7] partitions the global training contextual information into the final output representation. To alleviate the quadratic computational bottleneck of batch across multiple workers. In each iteration, every worker independently processes its assigned subset of data and standard attention in long-context modeling, various trainable computes the local gradients. At the end of each iteration, DP sparse attention methods have been proposed [17], [5], [18]. enforces a rigid synchronization barrier, requiring all workers to These methods leverage the inherent sparsity of the attention synchronize their gradients before updating the model weights. mechanism by selecting only a subset of highly relevant tokens, This inherent synchronization mechanism forces all workers referred to as critical tokens, to approximate the full attention to wait for the slowest worker (i.e., the straggler) to reach the computation. Existing trainable sparse methods typically adopt a blockbarrier, making DP highly sensitive to workload imbalances. Pipeline Parallelism (PP) [8], [9], [10], [11] splits the sparse paradigm, which partition the key and value (KV) model layers into sequential stages, with each stage assigned sequences into discrete blocks. They compute the correlation to a different PP worker. To maximize hardware utilization, between each query token and these blocks using specific existing methods orchestrate model execution into a pipeline routing metrics to determine the critical tokens. This type of by dividing the input data into multiple micro-batches. These sparse attention computation is formulated as follows: micro-batches then sequentially traverse all PP ranks via point to-point communication to deliver activations and gradients QK[I]⊤ √ SparseAttn(Q, K, V ) = softmax V [I] (2) between stages. However, PP inherently suffers from the d
0.15 0.10 0.05 0.00 0k
Time
Longalign-10k ChatQA2-Long-SFT
1234
PP Device
Percentage
0.20
10k
20k
30k
40k
50k
Sequence Lengths
1 4 21
2
3 3
4
24
11121 2 2 1 2 333 2 3444 4 1 1 2 2 3 3
DP Bubble
3
1122 3 1 1 3 2 1 2 4 343 2 44
PP Worker0
60k
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.
1
112233 4
PP Bubble
3 3 4
PP Worker1
4 4
4 4 PP Worker2
PP Worker3
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.