论文
SparseBalance:具有动态稀疏注意力的负载平衡长上下文训练
SparseBalance: Load-Balanced Long Context Training with Dynamic Sparse Attention
摘要
虽然稀疏注意力缓解了长上下文 LLM 训练的计算瓶颈,但其分布式训练过程在 \textit{1)} 序列长度和 \textit{2)} 稀疏敏感性方面表现出极大的异质性,导致严重的不平衡问题和次优模型精度。现有的算法和训练框架通常关注单一问题,未能系统地共同优化这两个问题。因此,我们提出了SparseBalance,一种新颖的算法系统协同设计框架,它利用稀疏性和序列异构性来共同优化模型精度和系统效率。首先,我们提出了工作负载感知的动态稀疏度调整,它采用双向稀疏度调整来消除掉队者并利用固有的气泡来实现自由精度。其次,我们提出了一种稀疏感知批处理策略来实现粗粒度平衡,这补充了动态稀疏调整。实验结果表明,SparseBalance 实现了高达 1.33$\times$ 的端到端加速,同时在 LongBench 基准上仍将长上下文能力提高了 0.46\%。