TOAST: Fast and scalable auto-partitioning based on principled static analysis
Automatic partitioning of large models across distributed accelerators faces three key challenges: an exponentially growing search space, high risk of out-of-memory (OOM) failures, and suboptimal or infeasible solutions due to heuristic pruning in existing tools. This paper proposes a novel hybrid approach integrating principled static compilation analysis with Monte Carlo Tree Search (MCTS). First, it models tensor dimension dependencies to precisely identify homogeneous sharding requirements and conflict constraints, thereby constructing a compact and feasible decision space. Second, it employs MCTS to efficiently explore this space while enforcing memory safety and execution efficiency. Evaluated across diverse hardware platforms and model architectures, our fully automated method discovers partitioning schemes that outperform industrial-grade baselines—including TensorFlow/XLA and DeepSpeed—achieving higher scalability, throughput, and zero GPU memory overflow.