Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR
Citations Over Time
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.
Related Papers
- → The SPMD Model: Past, Present and Future(2001)122 cited
- → Exploring load balancing in a scientific SPMD parallel application(2002)12 cited
- → Load Balancing in SPMD Applications: Concepts and Experiments(2004)3 cited
- A simulation methodology for the prediction of SPMD programs performance(1993)
- What does Semi-Permeable membrane device (SPMD) sampling mean? an insight on the "time-weight integrative" and "SPMD-availability" hypothesis(2008)