Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs

1 INTRODUCTION . Deep Learning frameworks such as MXNet ( Chen et al. , 2015 ) , PyTorch ( Paszke et al. , 2017 ) , and TensorFlow ( TensorFlow Authors , 2016a ) represent neural network models as computation graphs . Efficiently executing such graphs requires optimizing discrete decisions about how to map the computations in a graph onto hardware so as to minimize a relevant cost metric ( e.g. , running time , peak memory ) . Given that execution efficiency is critical for the success of neural networks , there is growing interest in the use of optimizing static compilers for neural network computation graphs , such as Glow ( Rotem et al. , 2018 ) , MLIR ( MLIR Authors , 2018 ) , TVM ( Chen et al. , 2018a ) , and XLA ( XLA team , 2017 ) . Here we consider the model parallelism setting where a computation graph can be executed using multiple devices in parallel . Nodes of the graph are computational tasks , and directed edges denote dependencies between them . We consider jointly optimizing over placement , i.e. , which nodes are executed on which devices , and schedule , i.e. , the node execution order on each device . These decisions are typically made in either one or two passes in the compiler . We consider two different objectives : 1 ) minimize running time , subject to not exceeding device memory limits , and 2 ) minimize peak memory usage . In the optimization literature , such problems are studied under the class of task scheduling , which is known to be NP-hard in typical settings ( Sinnen , 2007 ; Kwok & Ahmad , 1999 ) . As scheduling and placement are just a few of the many complex decisions made in a compiler , it is essential in a production setting that a solution 1 ) produce solutions of acceptable quality fast , even on large graphs ( e.g. , thousands of nodes ) and decision spaces , and 2 ) handle diverse graphs from various types of applications , neural network architectures , and users . In this work we consider ∗Google AI Resident learning an optimizer that satisfies these requirements . Crucially , we aim to learn an optimizer that generalizes to a broad set of previously unseen computation graphs , without the need for training on such graphs , thus allowing it to be fast at test time . Previous works on learning to optimize model parallelism decisions ( Mirhoseini et al. , 2017 ; 2018 ; Addanki et al. , 2019 ) have not considered generalization to a broad set of graphs nor joint optimization of placement and scheduling . In Mirhoseini et al . ( 2017 ; 2018 ) , learning is done from scratch for each computation graph and for placement decisions only , requiring hours ( e.g. , 12 to 27 hours per graph ) . This is too slow to be broadly useful in a general-purpose production compiler . We propose an approach that takes only seconds to optimize similar graphs . In concurrent work to ours , Addanki et al . ( 2019 ) shows generalization to unseen graphs , but they are generated artificially by architecture search for a single learning task and dataset . In contrast , we collect real user-defined graphs spanning a broad set of tasks , architectures , and datasets . In addition , both Mirhoseini et al . ( 2017 ; 2018 ) and Addanki et al . ( 2019 ) consider only placement decisions and rely on TensorFlow ’ s dynamic scheduler ; they do not address the static compiler setting where it is natural to jointly optimize scheduling and placement . The key idea of our approach ( Figure 1 ) is to learn a neural network that , conditioned on the input graph to be optimized , directs an existing optimization algorithm ’ s search such that it finds a better solution in the same search budget . We choose the Biased Random-Key Genetic Algorithm ( BRKGA ( Gonçalves & Resende , 2011 ) ) as the optimization algorithm after an extensive evaluation of several choices showed that it gives by far the best speed-vs-quality trade-off for our application . BRKGA produces good solutions in just a few seconds even for real-world TensorFlow graphs with thousands of nodes , and we use learning to improve the solution quality significantly at similar speed . We train a graph neural network ( Battaglia et al. , 2018 ) to take a computation graph as input and output node-specific proposal distributions to use in the mutant generation step of BRKGA ’ s inner loop . BRKGA is then run to completion with those input-dependent distribution choices , instead of inputagnostic default choices , to compute execution decisions . The distributions are predicted at each node , resulting in a high-dimensional prediction problem . There is no explicit supervision available , so we use the objective value as a reward signal in a contextual bandit approach with REINFORCE ( Williams , 1992 ) . Our approach , “ Reinforced Genetic Algorithm Learning ” ( REGAL ) , uses the network ’ s ability to generalize to new graphs to significantly improve the solution quality of the genetic algorithm for the same objective evaluation budget . We follow the static compiler approach of constructing a coarse static cost model to evaluate execution decisions and optimizing them with respect to it , as done in ( Addanki et al. , 2018 ; Jia et al. , 2018 ) . This is in contrast to evaluating the cost by executing the computation graph on hardware ( Mirhoseini et al. , 2017 ; 2018 ) . A computationally cheap cost model enables fast optimization . It is also better suited for distributed training of RL policies since a cost model is cheap to replicate in parallel actors , while hardware environments are not . Our cost model corresponds to classical NP-hard scheduling problems , so optimizing it is difficult . In this paper we focus fully on learning to optimize this cost model , leaving integration with a compiler for future work . We structure the neural network ’ s task as predicting proposal distributions to use in the search over execution decisions , rather than the decisions themselves directly . Empirically we have found the direct prediction approach to be too slow at inference time for our application and generalizes poorly . Our approach potentially allows the network to learn a more abstract policy not directly tied to detailed decisions that are specific to particular graphs , which may generalize better to new graphs . It can also make the learning task easier as the search may succeed even with sub-optimal proposal distribution predictions , thus smoothening the reward function and allowing the network to incrementally learn better proposals . The node-specific proposal distribution choices provide a rich set of knobs for the network to flexibly direct the search . Combining learning with a search algorithm has been shown to be successful ( e.g. , ( Silver et al. , 2017 ; 2018 ) ) , and our work can be seen as an instance of the same high-level idea . This paper makes several contributions : • We are the first to demonstrate learning a policy for jointly optimizing placement and scheduling that generalizes to a broad set of real-world TensorFlow graphs . REGAL significantly outperforms all baseline algorithms on two separate tasks of minimizing runtime and peak memory usage ( section 5.3 ) on datasets constructed from 372 unique real-world TensorFlow graphs , the largest dataset of its kind in the literature and at least an order of magnitude larger than the ones in previous works ( Mirhoseini et al. , 2017 ; 2018 ; Chen et al. , 2018b ; Addanki et al. , 2018 ; 2019 ) . • We use a graph neural network to predict mutant sampling distributions of a genetic algorithm , specifically BRKGA , for the input graph to be optimized . This directs BRKGA ’ s search in an input-dependent way , improving solution quality for the same search budget . • We compare extensively to classical optimization algorithms , such as enumerative search , local search , genetic search , and other heuristics , and analyze room-for-improvement in the objective value available to be captured via learning . Both are missing in previous works . 2 RELATED WORK . Learning to optimize computation graphs : AutoTVM ( Chen et al. , 2018b ) applies learning to the very different problem of optimizing low-level implementations of operators in a tensor program , while we focus on optimizing higher-level decisions such as placement and scheduling of ops . Mao et al . ( 2019 ) use graph neural nets and RL to learn a scheduling policy for data processing jobs on clusters . These works are conceptually similar to ours in their use of learning , applied to a different domain . Learning for combinatorial optimization : Our work is an instance of applying learning for combinatorial optimization ( Bengio et al. , 2018 ) . Previous works on learning graph combinatorial optimization algorithms ( e.g. , Li et al . ( 2018 ) ; Khalil et al . ( 2017 ) ) have focused on problems such as Minimum Vertex Cover , Maximum Clique , Maximum Independent Set , etc . The task scheduling problem we consider is significantly different in that the objective value is a more complex function on node-level decisions . Also , we focus on large-scale , real-world TensorFlow graphs , while e.g. , Khalil et al . ( 2017 ) uses small-scale , synthetic graph distributions . Learning a proposal distribution for stochastic search : Bunel et al . ( 2017 ) learns a policy for predicting instance-dependent proposal distributions to be used in the stochastic optimizer STOKE ( Schkufza et al. , 2013 ) for superoptimizing programs . However , it uses handcrafted instance features and shows results on relatively simple , small programs . In contrast , we automatically learn the instance representations and show results on real-world graphs . An earlier work by Paige & Wood ( 2016 ) similarly learns a neural network to predict input-dependent proposal distributions for sequential Monte Carlo search for inference in a graphical model . Optimization without learning : Parallel task scheduling ( Sinnen , 2007 ; Kwok & Ahmad , 1999 ) is a classical problem for scheduling ops in a computational graph to minimize runtime . Learning is not traditionally a part of the approaches proposed in this literature . Mayer et al . ( 2017 ) studies greedy task scheduling approaches for TensorFlow . Jia et al . ( 2018 ) develops a simulation-based optimizer for deep learning computation graphs that uses a larger decision space by combining data , model , and attribute parallelism . Our approach can potentially be extended to such larger decisions spaces to achieve even bigger improvements in execution cost . 3 BACKGROUND . Figure 1 shows an overview of our approach . Given an input graph to optimize , instead of applying BRKGA directly with the default uniform distribution at all nodes , a graph neural network predicts beta distribution choices at each node . BRKGA is run with these choices to optimize placement and scheduling decisions with respect to the objective defined by the performance model . We first explain the performance model and BRKGA in this section , and the learning component in the next . 3.1 PERFORMANCE MODEL . A computation graph has a set of ops to run . Each op produces zero or more tensors and requires zero or more tensors as input . The runtime of each op is known and fixed ( e.g. , given by a simulator as in Jia et al . ( 2018 ) ) . The memory use of each tensor is known ( an assumption that holds in static compilers like XLA ) . We assume a collection of d homogeneous devices that have separate local memory and can run at most one op at a time . An op can run only when its input tensors are present in the local memory . Tensors can be transferred across devices by synchronous ( blocking ) transfers . Tensors are freed from local memory after all local consumers have run . In this setting , we consider the problem of finding an assignment of ops to devices and an overall schedule such that each op is run once with the objectives of ( 1 ) minimizing the peak local memory use across devices ( e.g. , to find a feasible way to run a large computation graph ) , or ( 2 ) minimizing the runtime subject to a constraint on the peak memory used on any device . The performance model does not consider rematerialization of tensors , fragmentation when computing memory use , and asynchronous transfers between devices . Despite these simplifications , the model yields slight variants of problems that are known to be NP-hard ( Eyraud-Dubois et al. , 2015 ) and therefore remains a challenging setting in which to study how to learn an optimizer . See section A.4 for more details of the model .

Summary: In this paper, the authors proposed a framework to generating a task scheduling for a compiler to reduce the execution cost of neural networks. A computation graph is first fed into a GNN to produce a beta distribution, which is then fed into the BRKGA algorithm to yield the encoded solutions. The motivation is interesting, and the proposed method is technically reasonable. The details are also included in the appendix. To improve the quality, the following concerns may be considered: