Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MTraining: Distributed Dynamic Sparse Attention for Efficient Ultra-Long Context Training

Wenxuan Li, Chengruidong Zhang, Huiqiang Jiang, Yucheng Li, Yuqing Yang, Lili Qiu

TL;DR

The paper tackles the challenge of training LLMs with ultra-long contexts where quadratic attention costs become prohibitive. It introduces MTraining, a co-design of a training-oriented dynamic sparse pattern plus a hierarchy of imbalance-aware context-parallelism mechanisms (Balanced Sparse Ring Attention and Hierarchical Balanced Sparse Ring Attention) to achieve near-linear scaling in distributed settings. Empirical results on Qwen2.5-3B extended to 512K tokens demonstrate up to 6× end-to-end throughput gains across benchmarks such as RULER, Needle In A Haystack, InfiniteBench, and PG-19 while preserving or improving accuracy. The work shows that addressing worker- and step-level imbalance is essential for scalable dynamic sparse attention, and provides a public codebase to enable broader adoption and extension to larger models.

Abstract

The adoption of long context windows has become a standard feature in Large Language Models (LLMs), as extended contexts significantly enhance their capacity for complex reasoning and broaden their applicability across diverse scenarios. Dynamic sparse attention is a promising approach for reducing the computational cost of long-context. However, efficiently training LLMs with dynamic sparse attention on ultra-long contexts-especially in distributed settings-remains a significant challenge, due in large part to worker- and step-level imbalance. This paper introduces MTraining, a novel distributed methodology leveraging dynamic sparse attention to enable efficient training for LLMs with ultra-long contexts. Specifically, MTraining integrates three key components: a dynamic sparse training pattern, balanced sparse ring attention, and hierarchical sparse ring attention. These components are designed to synergistically address the computational imbalance and communication overheads inherent in dynamic sparse attention mechanisms during the training of models with extensive context lengths. We demonstrate the efficacy of MTraining by training Qwen2.5-3B, successfully expanding its context window from 32K to 512K tokens on a cluster of 32 A100 GPUs. Our evaluations on a comprehensive suite of downstream tasks, including RULER, PG-19, InfiniteBench, and Needle In A Haystack, reveal that MTraining achieves up to a 6x higher training throughput while preserving model accuracy. Our code is available at https://github.com/microsoft/MInference/tree/main/MTraining.

MTraining: Distributed Dynamic Sparse Attention for Efficient Ultra-Long Context Training

TL;DR

The paper tackles the challenge of training LLMs with ultra-long contexts where quadratic attention costs become prohibitive. It introduces MTraining, a co-design of a training-oriented dynamic sparse pattern plus a hierarchy of imbalance-aware context-parallelism mechanisms (Balanced Sparse Ring Attention and Hierarchical Balanced Sparse Ring Attention) to achieve near-linear scaling in distributed settings. Empirical results on Qwen2.5-3B extended to 512K tokens demonstrate up to 6× end-to-end throughput gains across benchmarks such as RULER, Needle In A Haystack, InfiniteBench, and PG-19 while preserving or improving accuracy. The work shows that addressing worker- and step-level imbalance is essential for scalable dynamic sparse attention, and provides a public codebase to enable broader adoption and extension to larger models.

Abstract

The adoption of long context windows has become a standard feature in Large Language Models (LLMs), as extended contexts significantly enhance their capacity for complex reasoning and broaden their applicability across diverse scenarios. Dynamic sparse attention is a promising approach for reducing the computational cost of long-context. However, efficiently training LLMs with dynamic sparse attention on ultra-long contexts-especially in distributed settings-remains a significant challenge, due in large part to worker- and step-level imbalance. This paper introduces MTraining, a novel distributed methodology leveraging dynamic sparse attention to enable efficient training for LLMs with ultra-long contexts. Specifically, MTraining integrates three key components: a dynamic sparse training pattern, balanced sparse ring attention, and hierarchical sparse ring attention. These components are designed to synergistically address the computational imbalance and communication overheads inherent in dynamic sparse attention mechanisms during the training of models with extensive context lengths. We demonstrate the efficacy of MTraining by training Qwen2.5-3B, successfully expanding its context window from 32K to 512K tokens on a cluster of 32 A100 GPUs. Our evaluations on a comprehensive suite of downstream tasks, including RULER, PG-19, InfiniteBench, and Needle In A Haystack, reveal that MTraining achieves up to a 6x higher training throughput while preserving model accuracy. Our code is available at https://github.com/microsoft/MInference/tree/main/MTraining.
Paper Structure (44 sections, 1 theorem, 16 equations, 15 figures, 5 tables, 2 algorithms)

This paper contains 44 sections, 1 theorem, 16 equations, 15 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Ring Attention
  4. Motivation
  5. Long-context Training is Dynamic Sparse
  6. Attention Training Sparsity Exhibits Patterns
  7. Distributed Dynamic Sparse Attention is Imbalanced
  8. MTraining
  9. Dynamic Sparse Training Pattern
  10. Balanced Sparse Ring Attention
  11. Hierarchical Balanced Sparse Ring Attention
  12. Experiments
  13. Implementation Details
  14. Long-context Training
  15. Long-context Extension Training
  16. ...and 29 more sections

Key Result

Theorem 3.1

The expectation of the attention weights after applying RoPE depends solely on the relative position $n-m$, i.e., $E[z_{n,m}] = \sum_{i=0}^{d-1} {\phi_{n-m}^{(i)}~ A_i} + \sum_{i=0}^{d-1} {\psi_{n-m}^{(i)}~ B_i}.$

Figures (15)

  • Figure 1: Workload distribution over 4 CP workers (GPUs) in Striped and Zigzag Ring Attention.
  • Figure 2: (a) Latency breakdown of the training stage. (b) The attention recall of top-k(k=1024) from 128K context in different sample and training step. (c-d) Visualization of attention weights (c) and their gradients (d) during training. Results are based on Qwen2.5-3B yang2024qwen2-5 trained with a 4×8 A100 cluster.
  • Figure 3: Illustration of worker- and step-level load imbalance problems introduced by dynamic sparse attention in distributed LLM training.
  • Figure 4: Overview of MTraining in distributed scenarios.
  • Figure 5: Step-level Computation Schedule of Striped Ring Attention (a) and Hierarchical Striped Ring Attention (b) with 4 CP workers.
  • ...and 10 more figures

Theorems & Definitions (1)

  • Theorem 3.1