Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR. (arXiv:2210.06352v1 [cs.NE])

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 …

arxiv discovery partitioning strategies