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: This work proposes to use a combination of graph neural networks (GNNs) and proximal policy optimization (PPO) to train policies for generalized device placement in dataflow graphs. Essentially, (1) a GNN is used to learn representations of a dataflow graph (in an inductive manner), (2) a transformer is used to output a device placement action for each node in the graph, and (3) the entire system is trained end-to-end via PPO. Extensive experimental results show very impressive results compared to strong baselines.