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

TL;DR

TOAST introduces a fast, scalable auto-partitioning system for large machine learning models by combining static analysis with Monte Carlo Tree Search, improving efficiency and solution quality over existing methods.

Contribution

It presents a novel static compiler analysis integrated with Monte Carlo Tree Search to efficiently explore partitioning options for large models.

Findings

Outperforms state-of-the-art industrial partitioners

Discovers previously unknown, superior partitioning solutions

Fully automated process for complex models

Abstract

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.