arXiv:2508.15010cs.LGcs.DC2025-08

用静态分析加速大模型分布式划分,自动找最优解

TOAST: Fast and scalable auto-partitioning based on principled static analysis

  • 通过静态分析识别必须同分块的张量维度和需解决的划分冲突
  • 在多硬件平台下比现有工业方法快数倍,且避免显存溢出
  • 适合大规模模型部署者,尤其关注自动化与性能优化的团队

将大型机器学习模型跨分布式加速器系统进行划分是一项复杂任务,涉及一系列相互依赖的决策,且受内部分块歧义影响。现有自动划分工具常因搜索空间指数级增长而出现显存溢出或运行过慢,因此人为限制搜索范围,但常导致违反设备内存约束或性能不佳的解。本文提出一种结合新型静态编译分析与蒙特卡洛树搜索的系统。该分析通过识别(i)需相同分块的张量维度,以及(ii)需解决的划分冲突,构建高效决策空间。系统在多种硬件平台和模型架构上显著优于当前最先进的工业方法,发现此前未知的更优解,整个过程对复杂大型模型实现完全自动化。

原文摘要 · Abstract (English)

Partitioning large machine learning models across distributed accelerator systems is a complex process, requiring a series of interdependent decisions that are further complicated by internal sharding ambiguities. Consequently, existing auto-partitioners often suffer from out-of-memory errors or are prohibitively slow when exploring the exponentially large space of possible partitionings. To mitigate this, they artificially restrict the search space, but this approach frequently yields infeasible solutions that violate device memory constraints or lead to sub-optimal performance. We propose a system that combines a novel static compiler analysis with a Monte Carlo Tree Search. Our analysis constructs an efficient decision space by identifying (i) tensor dimensions requiring identical sharding, and (ii) partitioning "conflicts" that require resolution. Our system significantly outperforms state-of-the-art industrial methods across diverse hardware platforms and model architectures, discovering previously unknown, superior solutions, and the process is fully automated even for complex and large models.

模型划分静态分析分布式训练自动化

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