arXiv:2510.18830cs.CLcs.DC2025-10被引 1

MTraining让大模型高效训练超长文本,32张A100显卡实现512K上下文扩展。

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

  • 动态稀疏注意力+分层环形通信,解决分布式训练不平衡问题。
  • 在32张A100上将上下文从32K拓展至512K,训练吞吐提升6倍。
  • 适合需要超长文本理解的模型训练,如推理、摘要与问答场景。

长上下文窗口已成为大语言模型的标准特性,能显著提升复杂推理能力并拓宽应用范围。动态稀疏注意力是降低长上下文计算开销的有前景方法。然而,在分布式环境下高效训练具有动态稀疏注意力的超长上下文模型仍面临重大挑战,主要源于工作节点与训练步级的不平衡。本文提出MTraining,一种新型分布式方法,利用动态稀疏注意力实现超长上下文大模型的高效训练。具体包含三个核心组件:动态稀疏训练模式、均衡稀疏环形注意力和分层稀疏环形注意力。这些组件协同缓解长上下文训练中动态稀疏注意力机制带来的计算不平衡与通信开销。我们在32张A100 GPU集群上成功将Qwen2.5-3B的上下文窗口从32K扩展至512K。在RULER、PG-19、InfiniteBench和Needle In A Haystack等下游任务上的评估显示,MTraining在保持模型精度的同时,训练吞吐最高提升6倍。代码已开源:https://github.com/microsoft/MInference/tree/main/MTraining。

原文摘要 · Abstract (English)

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.

大模型训练稀疏注意力超长上下文分布式

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。