SparseBalance: 동적 희소 어텐션을 활용한 로드 밸런싱 기반의 장문 컨텍스트 학습

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

희소 어텐션은 장문 컨텍스트 LLM 학습의 계산 병목 현상을 완화하지만, 분산 학습 과정에서 시퀀스 길이와 희소성 민감도 간의 극심한 불균형으로 인해 심각한 불균형 문제가 발생하고 모델 정확도가 저하되는 경우가 많습니다. 기존 알고리즘 및 학습 프레임워크는 일반적으로 단일 문제에 초점을 맞추어 이러한 두 가지 문제를 체계적으로 동시에 최적화하지 못합니다. 따라서 본 논문에서는 모델 정확도와 시스템 효율성을 동시에 최적화하기 위한 새로운 알고리즘-시스템 공동 설계 프레임워크인 SparseBalance를 제안합니다. 첫째, 워크로드에 대한 인식을 바탕으로 동적인 희소성 조절을 수행하여 지연 현상을 제거하고, 추가적인 정확도 향상을 위해 내재된 여유 자원을 활용합니다. 둘째, 희소성을 고려한 배치 전략을 제안하여 거친 수준에서의 균형을 달성하며, 이는 동적 희소성 조절을 보완합니다. 실험 결과는 SparseBalance가 전체적인 속도를 최대 1.33배 향상시키면서 LongBench 벤치마크에서 장문 컨텍스트 처리 능력을 0.46% 향상시켰음을 보여줍니다.

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.