PartIR: Composing SPMD Partitioning Strategies for Machine Learning
Authors:
Sami Alabed,
Daniel Belov,
Bart Chrzaszcz,
Juliana Franco,
Dominik Grewe,
Dougal Maclaurin,
James Molloy,
Tom Natan,
Tamara Norman,
Xiaoyue Pan,
Adam Paszke,
Norman A. Rink,
Michael Schaarschmidt,
Timur Sitdikov,
Agnieszka Swietlik,
Dimitrios Vytiniotis,
Joel Wee
Abstract:
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 par…
▽ More
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. We present a simple but powerful API for composing sharding strategies and a simulator to validate them. The process is driven by high-level programmer-issued partitioning tactics, which can be both manual and automatic. Importantly, the tactics are specified separately from the model code, making them easy to change. We evaluate PartIR on several different models to demonstrate its predictability, expressibility, and ability to reach peak performance..
△ Less
Submitted 3 March, 2024; v1 submitted 20 January, 2024;
originally announced January 2024.
The complexity of multilayer $d$-dimensional circuits
Authors:
T. R. Sitdikov,
G. V. Kalachev
Abstract:
In this paper we research a model of multilayer circuits with a single logical layer. We consider $λ$-separable graphs as a support for circuits. We establish the Shannon function lower bound $\max \bigl(\frac{2^n}{n}, \frac{2^n (1 - λ)}{\log k} \bigr)$ for this type of circuits where $k$ is the number of layers. For $d$-dimensional graphs, which are $λ$-separable for $λ= \frac{d - 1}{d}$, this gi…
▽ More
In this paper we research a model of multilayer circuits with a single logical layer. We consider $λ$-separable graphs as a support for circuits. We establish the Shannon function lower bound $\max \bigl(\frac{2^n}{n}, \frac{2^n (1 - λ)}{\log k} \bigr)$ for this type of circuits where $k$ is the number of layers. For $d$-dimensional graphs, which are $λ$-separable for $λ= \frac{d - 1}{d}$, this gives the Shannon function lower bound $\frac{2^n}{\min(n, d \log k)}$. For multidimensional rectangular circuits the proved lower bound asymptotically matches to the upper bound.
△ Less
Submitted 13 March, 2021;
originally announced March 2021.
