A Data-Centric Optimization Framework for Machine Learning

TL;DR

This paper introduces a flexible, data-centric optimization framework for deep learning that enhances training efficiency across diverse models by minimizing data movement, with interactive tools for customization and demonstrated speedups.

Contribution

It presents a novel, extensible pipeline for optimizing arbitrary neural networks through data movement reduction, applicable to various models and integrated with user-friendly interfaces.

Findings

Achieved competitive performance or speedups on ten different networks.

Discovered new optimization opportunities in EfficientNet.

Provided an interactive, extensible optimization pipeline for deep learning.

Abstract

Rapid progress in deep learning is leading to a diverse set of quickly changing models, with a dramatically growing demand for compute. However, as frameworks specialize performance optimization to patterns in popular networks, they implicitly constrain novel and diverse models that drive progress in research. We empower deep learning researchers by defining a flexible and user-customizable pipeline for optimizing training of arbitrary deep neural networks, based on data movement minimization. The pipeline begins with standard networks in PyTorch or ONNX and transforms computation through progressive lowering. We define four levels of general-purpose transformations, from local intra-operator optimizations to global data movement reduction. These operate on a data-centric graph intermediate representation that expresses computation and data movement at all levels of abstraction, including expanding basic operators such as convolutions to their underlying computations. Central to the design is the interactive and introspectable nature of the pipeline. Every part is extensible through a Python API, and can be tuned interactively using a GUI. We demonstrate competitive performance or speedups on ten different networks, with interactive optimizations discovering new opportunities in EfficientNet.