Efficiently Scaling Dynamic Sparse Attention in Distributed Training

Determine algorithmic and system-level methods to scale dynamic sparse attention efficiently in distributed training of Transformer-based large language models.

Background

Dynamic sparse attention has demonstrated substantial efficiency gains for long-context processing, especially at inference time, and recent work has begun to explore its use during pretraining. However, when moving to distributed training, particularly with context parallelism and ring-attention-style communication, practical obstacles such as worker imbalance and communication overhead complicate achieving efficient scaling.

The paper positions MTraining as a co-designed algorithm–system approach to address these issues, but explicitly notes that, in general, efficiently scaling dynamic sparse attention in distributed training remains unresolved in the literature, motivating the presented techniques and empirical paper.