Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR

TL;DR

This paper introduces an automatic partitioner that uses Monte Carlo Tree Search and compiler analysis to identify optimal SPMD strategies for large neural network training, reducing manual effort.

Contribution

It presents a novel goal-oriented search method that automatically discovers effective composite SPMD partitioning strategies for neural networks.

Findings

The partitioner matches expert-level strategies across various models.

Monte Carlo Tree Search effectively explores partitioning options.

Compiler analysis guides the search towards efficient solutions.

Abstract

Large neural network models are commonly trained through a combination of advanced parallelism strategies in a single program, multiple data (SPMD) paradigm. For example, training large transformer models requires combining data, model, and pipeline partitioning; and optimizer sharding techniques. However, identifying efficient combinations for many model architectures and accelerator systems requires significant manual analysis. In this work, we present an automatic partitioner that identifies these combinations through a goal-oriented search. Our key findings are that a Monte Carlo Tree Search-based partitioner leveraging partition-specific compiler analysis directly into the search and guided goals matches expert-level strategies for various models.