PartIR: Composing SPMD Partitioning Strategies for Machine Learning

Understand

Training of modern large neural networks (NN) requires a combination of parallelization strategies encompassing data, model, or optimizer sharding.

• When strategies increase in complexity, it becomes necessary for partitioning tools to be 1) expressive, allowing the composition of simpler strategies, and 2) predictable to estimate performance analytically.
• We present PartIR, our design for a NN partitioning system.
• PartIR is focused on an incremental approach to rewriting and is hardware-and-runtime agnostic.