GDP: Generalized Device Placement for Dataflow Graphs

1 INTRODUCTION . Neural networks have demonstrated remarkable scalability–improved performance can usually be achieved by training a larger model on a larger dataset ( Hestness et al. , 2017 ; Shazeer et al. , 2017 ; Jozefowicz et al. , 2016 ; Mahajan et al. , 2018 ; Radford et al. ) . Training such large models efficiently while meeting device constraints , like memory limitations , necessitate partitioning of the underlying dataflow graphs for the models across multiple devices . However , devising a good partitioning and placement of the dataflow graphs requires deep understanding of the model architecture , optimizations performed by domain-specific compilers , as well as the device characteristics , and is therefore extremely hard even for experts . ML practitioners often rely on their understanding of model architecture to determine a reasonable partitioning and placement for graphs . However , relying solely on the model architecture while ignoring the effect of the partitioning on subsequent compiler optimizations like op-fusion can lead to sub-optimal placements and consequently under-utilization of available devices . The goal of automated device placement is to find the optimal assignment of operations to devices such that the end-to-end execution time for a single step is minimized and all device constraints like memory limitations are satisfied . Since this objective function is non-differentiable , prior approaches ( Mirhoseini et al. , 2017 ; 2018 ; Gao et al. , 2018 ) have explored solutions based on reinforcement learning ( RL ) . However , these RL policies are usually not transferable and require training a new policy from scratch for each individual graph . This makes such approaches impractical due to the significant amount of compute required for the policy search itself , at times offsetting gains made by the reduced step time . In this paper , we propose an end-to-end deep RL method for device placement where the learned policy is generalizable to new graphs . Specifically , the policy network consists of a graph-embedding network that encodes operation features and dependencies into a trainable graph representation , followed by a scalable sequence-to-sequence placement network based on an improved Transformer ( Vaswani et al. , 2017 ; Dai et al. , 2019 ) . The placement network transforms the graph representations into a placement decision with soft attention , removing hard constraints such as hierarchical grouping of operations ( Mirhoseini et al. , 2018 ) or co-location heuristics ( to reduce the placement complexity ) ( Mirhoseini et al. , 2017 ) . Both of our graph-embedding network and placement network can be jointly trained in an end-to-end fashion using a supervised reward , without the need to manipulate the loss functions at multiple levels . We empirically show that the network learns flexible placement policies at a per-node granularity and can scale to problems over 50,000 nodes . To generalize to arbitrary and held-out graphs , our policy is trained jointly over a set of dataflow graphs ( instead of one at a time ) and then fine-tuned on each graph individually . By transferring the learned graph embeddings and placement policies , we are able to achieve faster convergence and thus use less resources to obtain high-quality placements . We also use super-positioning , i.e. , a feature conditioning mechanism based on the input graph embeddings , to effectively orchestrate the optimization dynamics of graphs with drastically different sizes in the same batch . Our contributions can be summarized as follows : 1 . An end-to-end device placement network that can generalize to arbitrary and held-out graphs . This is enabled by jointly learning a transferable graph neural network along with the placement network . 2 . A scalable placement network with an efficient recurrent attention mechanism , which eliminates the need for an explicit grouping stage before placement . The proposed end-to-end network provides 15× faster convergence as compared to the hierarchical LSTM model used in earlier works ( Mirhoseini et al. , 2017 ; 2018 ) . 3 . A new batch pre-training and fine-tuning strategy based on network superposition , which leads to improved transferability , better placements especially for larger graphs , and 10× reduction in policy search time as compared to training individual graphs from scratch . 4 . Superior performance over a wide set of workloads , including InceptionV3 ( Szegedy et al. , 2015 ) , AmoebaNet ( Real et al. , 2018 ) , RNNs , GNMT ( Wu et al. , 2016 ) , Transformer-XL ( Dai et al. , 2019 ) , WaveNet ( van den Oord et al. , 2016 ) , and more . 2 RELATED WORK . Device Placement Reinforcement learning has been used for device placement of a given dataflow graph ( Mirhoseini et al. , 2017 ) and demonstrated run time reduction over human crafted placement and conventional heuristics . For improved scalability , a hierarchical device placement strategy ( HDP ) ( Mirhoseini et al. , 2018 ) has been proposed that clusters operations into groups before placing the operation groups onto devices . Spotlight ( Gao et al. , 2018 ) applies proximal policy optimization and cross-entropy minimization to lower training overhead . Both HDP and Spotlight rely on LSTM controllers that are difficult to train and struggle to capture very long-term dependencies over large graphs . In addition , both methods are restricted to process only a single graph at a time , and can not generalize to arbitrary and held-out graphs . Placeto ( Addanki et al. , 2019 ) represents the first attempt to generalize device placement using a graph embedding network . But like HDP , Placeto also relies on hierarchical grouping and only generates placement for one node at each time step . Our approach ( GDP ) leverages a recurrent attention mechanism and generates the whole graph placement at once . This significantly reduces the training time for the controller . We also demonstrate the generalization ability of GDP over a wider set of important workloads . Parallelization Strategy Mesh-TensorFlow is a language that provides a general class of distributed tensor computations . While data-parallelism can be viewed as splitting tensors and operations along the “ batch ” dimension , in Mesh-TensorFlow the user can specify any tensor-dimensions to be split across any dimensions of a multi-dimensional mesh of processors . FlexFlow ( Jia et al. , 2018 ) introduces SOAP , a more comprehensive search space of parallelization strategies for DNNs which allows parallelization of a DNN in the Sample , Operator , Attribute , and Parameter dimensions . It uses guided randomized search of the SOAP space to find a parallelization strategy for a specific parallel machine . GPipe ( Huang et al. , 2018 ) proposed pipeline parallelism , by partitioning a model across different accelerators and automatically splitting a mini-batch of training examples into smaller micro-batches . By pipelining the execution across micro-batches , accelerators can operate in parallel . Our GDP focuses on a general deep RL method for automating device placement on arbitrary graphs , and is therefore orthogonal to existing parallelization strategies . Compiler Optimization REGAL ( Paliwal et al. , 2019 ) uses deep RL to optimize the execution cost of computation graphs in a static compiler . The method leverages the policy ’ s ability to transfer to new graphs to improve the quality of the genetic algorithm for the same objective budget . However , REGAL only targets peak memory minimization while GDP focuses on graph run time and scalability while also meeting the peak memory constraints of the devices . Specifically , we generalize graph partitioning and placement into a single end-to-end problem , with and without simulation , which can handle graphs with over 50,000 nodes . 3 END-TO-END PLACEMENT POLICY . Given a dataflow graph G ( V , E ) where V represents atomic computational operations ( ops ) and E represents the data dependency , our goal is to learn a policy π : G 7→ D that assigns a placement D ∈ D for all the ops in the given graph G ∈ G , to maximize the reward rG , D defined based on the run time . D is the allocated devices that can be a mixture of CPUs and GPUs . In this work , we represent policy πθ as a neural network parameterized by θ . Unlike prior works that focus on a single graph only , the RL objective in GDP is defined to simultaneously reduce the expected runtime of the placements over a set of N dataflow graphs : J ( θ ) = EG∼G , D∼πθ ( G ) [ rG , D ] ≈ 1 N ∑ G ED∼πθ ( G ) [ rG , D ] ( 1 ) In the following , we refer to the case whenN = 1 as individual training and the case whenN > 1 as batch training . We optimize the objective above using Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ) for improved sample efficiency . Figure 1 shows an overview of the proposed end-to-end device placement network . Our proposed policy network πθ consists a graph embedding network that learns the graphical representation of any dataflow graph , and a placement network that learns a placement strategy over the given graph embeddings . The two components are jointly trained in an end-to-end fashion . The policy p ( a|G ) is applied to make a set of decisions at each node . These decisions , denoted as av for each v ∈ V across all nodes , form one action a = { av∈V } . One decision corresponds to playing one arm of a multi-bandit problem , and specifying the entire a corresponds to playing several arms together in a single shot . Note the architecture is designed to be invariant over the underlying graph topology , enabling us to apply the same learned policy to a wide set of input graphs with different structures . 3.1 GRAPH EMBEDDING NETWORK . We leverage graph neural networks ( GNNs ) ( Hamilton et al. , 2017 ; Xu et al. , 2019 ; You et al. , 2018 ) to capture the topological information encoded in the dataflow graph . Most graph embedding frameworks are inherently transductive and can only generate embeddings for a given fixed graph . These transductive methods do not efficiently extrapolate to handle unseen nodes ( e.g. , in evolving graphs ) , and can not learn to generalize to unseen graphs . GraphSAGE ( Hamilton et al. , 2017 ) is an inductive framework that leverages node attribute information to efficiently generate representations on previously unseen data . While our proposed framework is generic , we adopt the feature aggregation scheme proposed in GraphSAGE to model the dependencies between the operations and build a general , end-to-end device placement method for a wide set of dataflow graphs . In GDP , nodes and edges in the dataflow graph are represented as the concatenation of their meta features ( e.g. , operation type , output shape , adjacent node ids ) and are further encoded by the graph embedding network into a trainable representation . The graph embedding process consists of multiple iterations , and the computation procedure for the l-th iteration can be outlined as follows : First , each node v ∈ V aggregates the feature representations of its neighbors , { h ( l ) u , ∀u ∈ N ( v ) } , into a single vector h ( l ) N ( v ) . This aggregation outcome is a function of all previously generated representations , including the initial representations defined based on the input node features . In this work , we use the following aggregation function with max pooling : h ( l ) N ( v ) = max ( σ ( W ( l ) h ( l ) u + b ( l ) ) , ∀u ∈ N ( v ) ) ( 2 ) where ( W ( l ) , b ( l ) ) define an affine transform and σ stands for the sigmoid activation function . We then concatenate the node ’ s current representation , h ( l ) v , with the aggregated neighborhood vector , h ( l ) N ( v ) , and feed this concatenated vector through a fully connected layer f ( l+1 ) h ( l+1 ) v = f ( l+1 ) ( concat ( h ( l ) v , h ( l ) N ( v ) ) ) ( 3 ) Different from GraphSAGE , parameters in our graph embedding network are trained jointly with a placement network via stochastic gradient descent with PPO , in a supervised fashion , as described in Section 3 . That is , we replace the unsupervised loss with our task-specific objective .

Summary: In this paper the authors propose an end-to-end policy for graph placement and partitioning of computational graphs produced "under-the-hood" by platforms like Tensorflow. As the sizes of the neural networks increase, using distributed deep learning is becoming more and more necessary. Primitives like the one suggested by the authors are very important in many ways, including improving the ability of the NN to process more data, reduce energy consumption etc. The authors compared to prior work propose a method that can take as input more than one data flow graphs, and learns a policy for graph partitioning/placement of the operations in a set of machines that minimizes the makespan. This problem in principle is NP-hard as it entails both graph partitioning and graph scheduling as its components. The authors propose a heuristic that composes of two existing methods: graph neural networks are used to produce an embedding of the computation/data flow graph, and then a seq-2-seq placement network. The method is able to generalize to unseen instances.