Search papers, labs, and topics across Lattice.
This paper introduces SparseBalance, a co-design framework that tackles load imbalance in distributed training of sparse attention models by addressing sequence length and sparsity sensitivity heterogeneity. It employs workload-aware dynamic sparsity tuning, using bidirectional sparsity adjustment to eliminate stragglers and exploit bubbles, coupled with a sparsity-aware batching strategy for coarse-grained balance. Experiments show SparseBalance achieves up to 1.33x end-to-end speedup and improves long-context capability by 0.46% on LongBench.
Training long-context sparse attention models doesn't have to be a slow, imbalanced mess: SparseBalance achieves 1.33x speedup while *improving* accuracy.
While sparse attention mitigates the computational bottleneck of long-context LLM training, its distributed training process exhibits extreme heterogeneity in both \textit{1)} sequence length and \textit{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.
