PartIR: Composing SPMD Partitioning Strategies for Machine Learning

Sami Alabed, Daniel Belov, Bart Chrzaszcz +14

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