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: This paper proposes an ML-based method to optimize TensorFlow Graph execution. Specifically, it combines graph neural networks (GNNs) and BRKGA (a genetic algorithm) to search over the joint space of TF node-device placement and scheduling. The core claims on the advantages of this method are that (1) it co-searches placement and scheduling space, (2) the trained model can generalize to different graphs and inference cost is very  small. The experimental results show that REGEL can outperform a few baseline methods on this problem.