用数学规划自动设计深度学习分布式并行策略,提升训练效率。
Automatic Operator-level Parallelism Planning for Distributed Deep Learning -- A Mixed-Integer Programming Approach
- 将并行规划建模为混合整数规划问题,实现自动化调度。
- 在相同内存下减少一半计算空洞,性能媲美专家设计策略。
- 适合复杂模型部署,可兼顾硬件利用率与内存约束。
随着参数量达数十亿的大模型发展,分布式训练与推理已成为必要。尽管数据并行、模型并行、序列并行和流水线并行等策略已在主流硬件上成功应用,但优化分布式部署仍需大量专家经验与人工调整。现有框架多采用简单链式结构,难以应对复杂非线性架构。混合专家(Mixture-of-Experts)与多模态模型具有复杂的多输入多输出(MIMO)和分支丰富的拓扑结构,亟需细粒度的算子级并行能力。本文提出将并行规划建模为混合整数规划(Mixed-Integer Programming)问题,设计双层求解框架,在最优性与计算效率间取得平衡,自动生成能捕捉现代神经网络异构结构及底层硬件约束的有效分布式方案。实验表明,在与DeepSeek DualPipe等专家设计策略对比中,本框架在相同内存约束下将计算空洞减少50%,性能相当或更优。其灵活性还可扩展至吞吐率优化、硬件利用率最大化、内存容量约束等多元目标,兼具科研探索价值与工业落地潜力。
原文摘要 · Abstract (English)
As the artificial intelligence community advances into the era of large models with billions of parameters, distributed training and inference have become essential. While various parallelism strategies-data, model, sequence, and pipeline-have been successfully implemented for popular neural networks on main-stream hardware, optimizing the distributed deployment schedule requires extensive expertise and manual effort. Further more, while existing frameworks with most simple chain-like structures, they struggle with complex non-linear architectures. Mixture-of-experts and multi-modal models feature intricate MIMO and branch-rich topologies that require fine-grained operator-level parallelization beyond the capabilities of existing frameworks. We propose formulating parallelism planning as a scheduling optimization problem using mixed-integer programming. We propose a bi-level solution framework balancing optimality with computational efficiency, automatically generating effective distributed plans that capture both the heterogeneous structure of modern neural networks and the underlying hardware constraints. In experiments comparing against expert-designed strategies like DeepSeek's DualPipe, our framework achieves comparable or superior performance, reducing computational bubbles by half under the same memory constraints. The framework's versatility extends beyond throughput optimization to incorporate hardware utilization maximization, memory capacity constraints, and other considerations or potential strategies. Such capabilities position our solution as both a valuable research tool for exploring optimal parallelization strategies and a practical industrial solution for large-scale AI deployment.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。