Orchestrating optimization passes of machine learning compiler for reducing memory footprints of computation graphs

Abstract

With the emergence of the needs of edge computing, there arises a demand for training and inferring deep learning (DL) models on memory-constrained devices. However, many DL models, namely computation graphs, have complex structure and plenty of parameters, incurring heavy memory consumption at runtime. Hence it is challenging but necessary to reduce their memory footprints at runtime.

This paper proposes OPass, a novel approach to perform hierarchical memory-constrained operator scheduling of machine learning models, and orchestrate optimization passes of Apache’s TVM (a machine learning compilation framework) for lowering memory footprints of computation graphs, finally allowing the graphs to run on memory-constrained devices. Firstly, given a computation graph $G$, OPass optimizes the graph heuristically and iteratively: OPass learns the effects of passes on the graph; it then optimizes $G$ iteratively — each iteration picks up a pass by the reduction of the memory footprint of $G$ and as well the implicit effects of the pass for further optimizations, letting the pass be applied. The second core component of OPass is its memory computation technique, named OPassMem, which hierarchically schedules $G$’s operators. It constructs a hierarchical computation graph and employs an iterative scheduling algorithm to progressively reduce memory footprints.

We evaluate OPass on ReBench (a suite of computation graphs) and two real-world models (Transformer and ResNet). The results show the strength of OPass: it reduces up to 90.83% of graph’s memory footprints, outperforming TVM’s default by 2.34$\times$. Specifically, pass orchestration and graph scheduling reduce memory footprints by up to 54.34% and 81%, respectively.