arXiv:2503.19050cs.DCcs.AI2025-03中稿 · EuroSys 2025被引 18

Mist自动协同优化内存与并行策略,显著提升大模型分布式训练速度。

Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimization

  • 细粒度重叠调度,动态协调多种优化技术
  • 符号化分析加速调优,实现1.28倍以上提速
  • 兼顾跨微批次不平衡,适合大规模模型训练

为加速大语言模型的分布式训练,现有方法采用数据、张量、流水线并行及激活检查点、冗余消除、卸载等内存优化技术。然而,现有自动系统因缺乏重叠感知、无法遍历庞大搜索空间且忽略微批次间不平衡,仅能调优部分优化。为此,本文提出Mist——一种兼顾内存、重叠和不平衡的自动分布式训练系统,全面协同优化所有内存压缩技术与并行策略。核心思想包括:(1) 细粒度重叠中心调度,实现优化技术的重叠执行;(2) 基于符号的性能分析,通过符号表达式快速预测运行时与内存占用;(3) 不平衡感知的分层调优,将问题分解为跨阶段重叠感知的混合整数规划与阶段内双目标约束优化,并通过帕累托前沿采样连接。评估显示,Mist相比最先进的手动系统Megatron-LM平均提速1.28倍(最高1.73倍),相比最先进的自动系统Aceso平均提速1.27倍(最高2.04倍)。

原文摘要 · Abstract (English)

Various parallelism, such as data, tensor, and pipeline parallelism, along with memory optimizations like activation checkpointing, redundancy elimination, and offloading, have been proposed to accelerate distributed training for Large Language Models. To find the best combination of these techniques, automatic distributed training systems are proposed. However, existing systems only tune a subset of optimizations, due to the lack of overlap awareness, inability to navigate the vast search space, and ignoring the inter-microbatch imbalance, leading to sub-optimal performance. To address these shortcomings, we propose Mist, a memory, overlap, and imbalance-aware automatic distributed training system that comprehensively co-optimizes all memory footprint reduction techniques alongside parallelism. Mist is based on three key ideas: (1) fine-grained overlap-centric scheduling, orchestrating optimizations in an overlapped manner, (2) symbolic-based performance analysis that predicts runtime and memory usage using symbolic expressions for fast tuning, and (3) imbalance-aware hierarchical tuning, decoupling the process into an inter-stage imbalance and overlap aware Mixed Integer Linear Programming problem and an intra-stage Dual-Objective Constrained Optimization problem, and connecting them through Pareto frontier sampling. Our evaluation results show that Mist achieves an average of 1.28$\times$ (up to 1.73$\times$) and 1.27$\times$ (up to 2.04$\times$) speedup compared to state-of-the-art manual system Megatron-LM and state-of-the-art automatic system Aceso, respectively.

大模型训练自动调优内存优化

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