Automap: Towards Ergonomic Automated Parallelism for ML Models

Schaarschmidt, Michael; Grewe, Dominik; Vytiniotis, Dimitrios; Paszke, Adam; Schmid, Georg Stefan; Norman, Tamara; Molloy, James; Godwin, Jonathan; Rink, Norman A.; Nair, Vinod; Belov, Dan

doi:10.48550/arxiv.2112.02958

Cited by 2 publications

(7 citation statements)

References 7 publications

Supporting

Mentioning

Contrasting

Order By: Relevance

“…Under this setting, we report excellent parallel efficiency that remains above 99% as the number of GPUs is increased. While we have only considered data-parallelism in this study, we may be able to obtain further speed-ups by considering a combination of data-and function-parallelism techniques [54] in future studies. Figure 14 also reports the effect of batch-size of training on the resulting L 2 accuracy for the first time-window (t ∈ [0, 0.1]).…”

Section: Navier-stokes Equationmentioning

confidence: 99%

Respecting causality is all you need for training physics-informed neural networks

Wang¹,

Sankaran²,

Perdikaris³

2022

Preprint

View full text Add to dashboard Cite

While the popularity of physics-informed neural networks (PINNs) is steadily rising, to this date PINNs have not been successful in simulating dynamical systems whose solution exhibits multi-scale, chaotic or turbulent behavior. In this work we attribute this shortcoming to the inability of existing PINNs formulations to respect the spatio-temporal causal structure that is inherent to the evolution of physical systems. We argue that this is a fundamental limitation and a key source of error that can ultimately steer PINN models to converge towards erroneous solutions. We address this pathology by proposing a simple re-formulation of PINNs loss functions that can explicitly account for physical causality during model training. We demonstrate that this simple modification alone is enough to introduce significant accuracy improvements, as well as a practical quantitative mechanism for assessing the convergence of a PINNs model. We provide state-of-the-art numerical results across a series of benchmarks for which existing PINNs formulations fail, including the chaotic Lorenz system, the Kuramoto-Sivashinsky equation in the chaotic regime, and the Navier-Stokes equations in the turbulent regime. To the best of our knowledge, this is the first time that PINNs have been successful in simulating such systems, introducing new opportunities for their applicability to problems of industrial complexity.

show abstract

Section: Navier-stokes Equationmentioning

confidence: 99%

Respecting causality is all you need for training physics-informed neural networks

Wang¹,

Sankaran²,

Perdikaris³

2022

Preprint

View full text Add to dashboard Cite

show abstract

“…Search with PartIR actions Automap [27] applies such actions to the module, guided by a worklist consisting of the module arguments (e.g., model parameters, optimizer state and the data tensors passed to a training step function). If the compiler detects that the sharding of an argument can be propagated to another argument, then both are removed from the worklist and no longer considered for further sharding actions (See Figure 2 for an example.…”

Section: Partir Rewriting For Spmd Partitioningmentioning

confidence: 99%

“…In this work, we extend Automap [27] to support multi-axis automatic partitioning, a capability that our PartIR stack already supports through nesting tiling and reduction loops, as well as multi-axis loop propagation. However, multi-axis partitioning increases the search complexity exponentially for every axis.…”

Section: Multi-axis Searchmentioning

confidence: 99%

“…The search must consider the dependency between partitioning decisions and the argument for partition, axis, and dimension. Compiler static analysis together with Automap's [27] worklist mechanism described above helps us prune the partitioning decision branching factor. This section outlines the major additions needed to reduce the complexity of the problem and reach fast expert-level sharding on several representative models.…”

Section: Multi-axis Searchmentioning

confidence: 99%

“…In this work, we extend our automated partitioner, Automap [27], to support partitioning a model across multi-dimensional device topology and discover expert-level composite SPMD partitioning strategies. Our contribution is a design of a goal-oriented Monte Carlo tree search (MCTS) [3] that decomposes the problem into smaller sub-problems.…”

Section: Introductionmentioning

confidence: 99%

See 2 more Smart Citations

Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR

Alabed¹,

Grewe²,

Franco³

et al. 2022

Preprint

View full text Add to dashboard Cite

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.

show abstract

Automap: Towards Ergonomic Automated Parallelism for ML Models

Cited by 2 publications

References 7 publications

Respecting causality is all you need for training physics-informed neural networks

Respecting causality is all you need for training physics-informed neural networks

Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR

Contact Info

Product

Resources

About