# DistIR: An Intermediate Representation and Simulator for Efficient Neural Network Distribution

arXiv: https://arxiv.org/abs/2111.05426 · HTML: https://mlsanity.com/paper/2111.05426/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 43 references that are in the map.
- First posted on arXiv 9 Nov 2021.

## Builds on (most cited first, 19 of 19)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [One weird trick for parallelizing convolutional neural networks](https://mlsanity.com/paper/1404.5997/) (2014)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [Efficient Algorithms for Device Placement of DNN Graph Operators](https://mlsanity.com/paper/2006.16423/) (2020)
- [Daydream: Accurately Estimating the Efficacy of Optimizations for DNN Training](https://mlsanity.com/paper/2006.03318/) (2020)
- [Stateful Dataflow Multigraphs: A Data-Centric Model for Performance Portability on Heterogeneous Architectures](https://mlsanity.com/paper/1902.10345/) (2019)
- [AMPNet: Asynchronous Model-Parallel Training for Dynamic Neural Networks](https://mlsanity.com/paper/1705.09786/) (2017)
- [Dynamic Control Flow in Large-Scale Machine Learning](https://mlsanity.com/paper/1805.01772/) (2018)
- [Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloads](https://mlsanity.com/paper/2007.04069/) (2020)
- [A Language for Describing Optimization Strategies](https://mlsanity.com/paper/2002.02268/) (2020)

## Built on by (most cited first, 2 of 2)

- [Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces](https://mlsanity.com/paper/2305.14516/) (2023)
- [MLCommons Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces](https://mlsanity.com/paper/2605.11333/) (2026)

## Abstract (arXiv)

The rapidly growing size of deep neural network (DNN) models and datasets has given rise to a variety of distribution strategies such as data, tensor-model, pipeline parallelism, and hybrid combinations thereof. Each of these strategies offers its own trade-offs and exhibits optimal performance across different models and hardware topologies. Selecting the best set of strategies for a given setup is challenging because the search space grows combinatorially, and debugging and testing on clusters is expensive. In this work we propose DistIR, an expressive intermediate representation for distributed DNN computation that is tailored for efficient analyses, such as simulation. This enables automatically identifying the top-performing strategies without having to execute on physical hardware. Unlike prior work, DistIR can naturally express many distribution strategies including pipeline parallelism with arbitrary schedules. Our evaluation on MLP training and GPT-2 inference models demonstrates how DistIR and its simulator enable fast grid searches over complex distribution spaces spanning up to 1000+ configurations, reducing optimization time by an order of magnitude for certain regimes.