DistIR: An Intermediate Representation for Optimizing Distributed Neural Networks

By Keshav Santhanam, Siddharth Krishna, Ryota Tomioka, Andrew William Fitzgibbon, Tim Harris. 2021

21 citations recorded by the source corpus.

The rapidly growing size of deep neural network (DNN) models and datasets has given rise to a variety of distribution strategies such as data, horizontal, and pipeline parallelism. However, selecting the best set of strategies for a given model and hardware configuration is challenging because debugging and testing on clusters is expensive. In this work we propose DistIR, an IR for explicitly representing distributed DNN computation that can capture many popular distribution strategies. We build an analysis framework for DistIR programs, including a simulator and reference executor that can be used to automatically search for an optimal distribution strategy. Our unified global representation also eases development of new distribution strategies, as one can reuse the lowering to per-rank backend programs. Preliminary results using a grid search over a hybrid data/horizontal/pipeline-parallel space suggest DistIR and its simulator can aid automatic DNN distribution.

Subfields: Computer Networks and Communications, Computer Vision and Pattern Recognition, Hardware and Architecture, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Networks and Communications

Ruoming Pang

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Model Architectures

Citations in the corpus

GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, methodology): “DistIR [31] is an MPMD-style intermediate representation for distributed ML computations.”