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

TL;DR

DistIR introduces an expressive intermediate representation and simulator for distributed neural network computation, enabling rapid exploration of complex distribution strategies and significantly reducing optimization time without physical hardware execution.

Contribution

The paper presents DistIR, a novel intermediate representation that can naturally express diverse distribution strategies and facilitate efficient simulation and optimization of distributed DNN training.

Findings

Enables fast grid search over 1000+ configurations

Reduces optimization time by an order of magnitude

Supports various distribution strategies including pipeline parallelism

Abstract

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.