Title: GDP: GENERALIZED DEVICE PLACEMENT FOR DATAFLOW GRAPHS

Abstract: Runtime and scalability of large neural networks can be significantly affected by the placement of operations in their dataflow graphs on suitable devices. With increasingly complex neural network architectures and heterogeneous device characteristics, finding a reasonable placement is extremely challenging even for domain experts. Most existing automated device placement approaches are impractical due to the significant amount of compute required and their inability to generalize to new, previously held-out graphs. To address both limitations, we propose an efficient end-to-end method based on a scalable sequential attention mechanism over a graph neural network that is transferable to new graphs. On a diverse set of representative deep learning models, including Inception-v3, AmoebaNet, Transformer-XL, and WaveNet, our method on average achieves 16% improvement over human experts and 9.2% improvement over the prior art with 15× faster convergence. To further reduce the computation cost, we pre-train the policy network on a set of dataflow graphs and use a superposition network to fine-tune it on each individual graph, achieving state-of-the-art performance on large hold-out graphs with over 50k nodes, such as an 8-layer GNMT.

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 cannot 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) 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 → D that assigns a placement D ∈ D for all the ops in the given graph G ∈ G, to maximize the reward r G,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:

EQUATION

In the following, we refer to the case when N = 1 as individual training and the case when N > 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 a v for each v ∈ V across all nodes, form one action a = {a v∈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 cannot 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:

EQUATION

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

v , with the aggregated neighborhood vector, h (l) N (v) , and feed this concatenated vector through a fully connected layer f (l+1)

EQUATION

EQUATION

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.

3.2. PLACEMENT NETWORK:

The graph neural network works as a feature aggregation network that learns a trainable feature representation for the computational graph, we still need a policy network that produces actions on a per node basis. Given h v 's, the policy network produces a v 's through conditionally independent predictions, where the prediction for one node v does not depend on the prediction of other nodes.

EQUATION

While f can be represented using multilayer perceptrons (MLPs), where the MLPs is shared across all nodes for prediction the placement output distributions. However, MPLs lack a dependency tracking mechanism across nodes. In practise, the placement of one node can be determined by the placement of another node, where the placed node may consume a large size of data produced by the other node. Intuitively, an attention network can learn this dependency and the relative importance of dependencies across an entire graph. Therefore, we decide to use an attention-based placement network to better track inter-node placement-related dependencies.

Designing a scalable placement network that can generalize to graphs with thousands of nodes is challenging, as the conventional GNMT models proposed for language tasks usually target a shorter sequence length. Hierarchical placement (Mirhoseini et al., 2018) has been proposed to address this issue,however, the proposed grouper network comes with limited flexibility and generality. For example, the grouper network leverages an aggregated feature representation by averaging feature vectors for nodes within the same group. The non-differentiable grouping procedure prevents training the graph-embedding and placement networks end-to-end.

To remove the two-stage hierarchical workflow in HDP for improved scalability, we propose to use a Transformer-based attentive network to generate operation placements in an end-to-end fashion. As the graph embedding already contains spatial (topological) information for each node, we remove the positional embedding in the original transformer to prevent the model from overfitting node identifications. To capture long-term dependencies efficiently among a large set of nodes, we adopt segment-level recurrence introduced in Transformer-XL (Dai et al., 2019; Dai, 2019) , where hidden states computed for the previous set of nodes are cached (with gradient flows disabled) and reused as an extended context during the training of the next segment. Besides achieving extra long context, we empirically find the segment-level recurrent attention much faster than a conventional LSTMbased GNMT model. In our experimental evaluation, we compare both the performance and speed up of our placement network with that of the LSTM-based hierarchical device placement.

3.3. BATCH TRAINING WITH PARAMETER SUPERPOSITION:

Since the parameterization for the architecture of the end-to-end policy is designed to be invariant over input graphs with different topologies, the same placement policy can be shared across a wide set of workloads. We therefore propose a batch training strategy, and further enhance the aforementioned architecture to handle such generalization across graphs.

Naïve batch training is challenging in our context as different dataflow graphs contain different number of operations connected in different topologies. In addition, unlike previous device placement methods, GDP aims to handle graphs from potentially different application domains (e.g. computer vision, language, and speech), where the number of operations can range from a few thousand to one million. These graphs have drastically different network architecture, in terms of computational operations, data shape, and network topology. As an example, recurrent networks have completely different operation types and connections compared to multi-branch convolutional networks that are widely used in computer vision. It would be highly desirable to train a single shared network that maximizes information sharing across these heterogeneous tasks, without hurting the performance on each of them due to their distinct learning dynamics.

Along a similar direction of multi-task learning and few-shot learning (Oreshkin et al., 2018) , we propose a feature conditioning mechanism similar to parameter superposition (Cheung et al., 2019) . The idea is to train one shared policy, but condition its parameters based on the input features to mitigate the potentially undesirable interference among different input graphs. Since dense layers (affine transforms followed by nonlinearity) serve as the fundamental building blocks in all of our network components, we introduce an additional conditioning layer to enable superposition in all dense layers the placement network:

EQUATION

where g (l) stands for a dense layer in our policy network, c stands for the feature conditioning layer, and x (0) denotes the feature representation of the input graph generated by the graph-embedding network. The feature conditioning layer is implemented with minimum overhead by adding an additional transformer layer to our placement network.

4. EXPERIMENT:



4.1. EXPERIMENT SETUP:

In this section, we evaluate our training strategy on widely used machine learning models in computer vision, natural language processing, and speech domains. We compare our approach to human expert placement, TensorFlow METIS placement, and hierarchical device placement (HDP) (Mirhoseini et al., 2018) . Our experiments are run on machines with one Intel Broadwell CPU and up to eight Nvidia P100 GPUs. Note that the prior work (Mirhoseini et al., 2017; 2018; Gao et al., 2018) were evaluated on different GPU devices, preventing direct comparison of results. Therefore, we re-evaluate HDP on our own system environment and report those numbers.

The performance of a placement is evaluated by the resulted training step time (run time) of the neural network. We use the negative square root of the normalized run time as the reward, where the run time is normalized with the best run time from a baseline. We use the average reward of all the previous trials as a bias term. The advantage value is computed by subtracting the reward by the average reward. During the search, we apply a large negative reward (-10) for invalid placements (e.g. a violation of co-location constraint, out of memory, etc.). For operation scheduling, we rely on the Tensorflow default FIFO scheduling.

4.2. PERFORMANCE ON INDIVIDUAL GRAPHS:

We evaluate GDP by training the model separately on six important graphs, including RNN Language Modeling, GNMT (Sutskever et al., 2014) , Transformer-XL, Inception, AmoebaNet, and WaveNet. We name this approach GDP-one. For all the tasks, GDP-one consistently outperforms human expert placement, TensorFlow METIS (Karypis & Kumar, 1998) placement, and HDP. For extremely large graphs, GDP-one is only 6% worse on 8-layer NMT (over 60k nodes), compared to human placement, but is 6.8% better than HDP. Overall, GDP-one achieves on average more than 16% run time reduction across the evaluated 12 graphs, compared to human expert placement. Compared to hierarchical device placement, GDP-one achieves an average 9.2% speed up, and scales better to large graphs such as 8-layer NMT and 4-layer RNNLM. Importantly, with the efficient end-to-end training and sample efficient reinforcement learning algorithm, GDP-one has a 15x speed up in convergence time of the placement network over HDP.

4.3. GENERALIZATION:

GDP enables the training of multiple heterogeneous graphs in a single batch, sharing parameters in the graph-embedding network and the placement network. We name this training strategy GDPbatch. We empirically show that GDP-batch generates better placements for many workloads such as transformer-XL (7.6%), WaveNet (15%), and 8-layer GNMT (8%). Table 2 compares the run time of 11 tasks using GDP-batch, with the same end-to-end architecture as described in section 4.2. GDP-batch yields slightly better run time compared to GDP-one in majority of the tasks, while being only slightly worse on AmoebaNet. Compared to training graphs separately, GDP-batch reduces network parameters and enables transfer learning among different graphs.

We further evaluate the effect of transfer learning by mixing redundant tasks in a batch. We find that mixing different graphs such as RNNLM and GNMT models with different number of layers results in both faster and better learning for RNNLM and GNMT with large number of layers (8layer). As a matter of fact, both Placeto (Addanki et al., 2019) and HDP had problems matching human placement performance for 8-layer GNMT or 8-layer RNNLM. With batch training, GDP is the first device placement work to match human expert performance for both 8-layer GNMT and 8-layer RNNLM. We also for the first time show that GDP-batch not only improves the search time (since we do not retrain the policy for every new graph), it can also improve the performance of the found placements. More detailed results are shown in Appendix Table 5 . Generalization to hold-out graphs: Here we show another set of experiments where we treat GDPbatch as a pre-training strategy and remove the target graph from the batch training dataset. We then fine-tune the pre-trained model on the hold-out graphs for fewer than 50 steps, which takes less than one minute. We name this GDP-generalization+finetune. Figure 2 shows that GDP fine-tuning for hold-out graphs outperforms human expert placement and HDP consistently on all six batch training datasets, and performs only slightly worse than GDP-one. 2-layer RNNLM and 2-stack WaveNet almost match the performance of GDP-one. We also run inference (generate placement) directly on the pre-trained model for the target hold-out graphs, and name this GDP-generalization-zeroshot.

We find that GDP-generalization-zeroshot only marginally hurts performance as compared to GDP-generalization+finetune, while being slightly better than human placement and HDP. This indicates that both graph embedding and the learned policies transfer and generalize to the unseen data. Comparisons with other generalized placement approaches: Placeto (Addanki et al., 2019) , to our knowledge, is the only other method beside GDP that shows true (and non-simulated) generalized device placement results. Direct comparison is not possible since Placeto uses a different hardware platform and different input graphs (Inception-V3, NMT, and NASNet). Placeto's search time is on average 2.65x faster than HDP, while GDP is on average 15x faster than HDP on our larger set of graphs. Apart from search time speed up, Placeto on average reduces placed graph run time by 3% (for its different graphs and hardware) while GDP on average reduces placed graph run time by 9.2%, compared to HDP. One advantage of GDP over Placeto is that it does not rely on any initial feasible placement. Providing a reasonable initial placement is often non-trivial for domain experts, especially for larger graphs such as 8-layer GNMT. As such, we are the first to report superhuman results on 8-layer GNMT (with GDP-batch).

4.4. ABLATION STUDIES:

Attention and Superposition. We did an ablation study on the attention and the superposition layer in the transformer-XL placer network. We find that attention improves placement run time by an average of 18% compared to a placer network with no attention, and superposition improves placement run time by an average of 6.5% where all the graphs are trained in a single batch as described in Section 4.3. Without superposition network, batch training fails for AmoebaNet and Inception when mixing with larger RNNLM or GNMT models (4-layer).

Pre-training graph embeddings. We also evaluate a fine-tuning strategy by pre-training the graph embedding and placement network and fine-tuning the network on the down stream tasks. The difference here compared to Section 4.3 is that we also include the target graphs in the pre-training dataset. When GDP-batch is used as a pre-training strategy, the graph embedding and placement network assimilate meaningful graph representations and placement policies from a wide set of graphs, thus can be used as a strong baseline network for fine-tuning on downstream tasks. We compare the generated placement run time and the placement search time, normalized to GDP-one. We find that fine-tuning further reduces the the placed graph run time by an average of 5% and placement search time by an average of 86%, compared to GDP-one.

5. CONCLUSION:

In this paper, we present a generalized device placement strategy that uses a graph neural network and super-positioning to generalize to arbitrary and held out graphs. Through experimental evaluation over a wide set of representative graphs from different domains including computer vision, speech, and NLP, we demonstrated over 15 times faster convergence while achieving a 16% and 9.2% reductions in step time over human expert placement and HDP, respectively.

6. APPENDIX 6.1 PROXIMAL POLICY OPTIMIZATION:

In device placement, the objective is to minimize the training step time of a given computational graph or a batch of dataflow graphs for a target system configuration (e.g. a 8-GPU cluster or a TPU pod), by placing operations onto different devices to enable model-level parallelism. This process corresponds to maximizing the expected performance in the MDP. For better sample efficiency, we adopted a Proximal Policy Optimization (PPO) (Schulman et al., 2017) algorithm. The objective is to maximize a surrogate objective:

L π = E a [0:n] ∼π [ q (a n |s n ) q(a n |s n ) A π (s n , a n )] L π = max π 1 N N -1 n=0,an∼π [min( q (a n |s n ) q(a n |s n ) (R -R), clip( q (a n |s n ) q(a n |s n ) , 1 -, 1 + )(R -R))]

Within a loop, GDP PPO continuously samples placements from the distribution and evaluates their training times in real systems. For a rollout of K, we perform a minimatch of m stochastic gradient ascent steps with respective to the objective of proximal policy optimization, which makes incremental policy improvements. The rollout steps K and minibatch size m are hyper parameters for PPO. We find a set of optimized hyper parameters and keep them fixed for all the experiments presented. As the rewards are generated on-the-fly based on real system measurements, we no longer need to re-evaluate the placement solutions in a separate phase.

6.2. HYPERPARAMETERS:

In this section, we list out all the selected hyperparameters in our experiments for reproducibility in Table 3 and Table 4 . 

6.3. INPUT GRAPHS:

We used a variety of widely used workloads from computer vision, speech, and NLP. In this section, we give a detailed explanation on the selected models and hyperparameters. 

:



Figures and Tables:

Figure 1: Overview of GDP: An end-to-end placement network that combines graph embedding and sequential attention. N : Number of Nodes, h: Hidden Size, d: Number of Devices.

Figure 2: Finetuning on hold-out graphs.Comparisons with other generalized placement approaches: Placeto(Addanki et al., 2019), to our knowledge, is the only other method beside GDP that shows true (and non-simulated) generalized device placement results. Direct comparison is not possible since Placeto uses a different hardware platform and different input graphs (Inception-V3, NMT, and NASNet). Placeto's search time is on average 2.65x faster than HDP, while GDP is on average 15x faster than HDP on our larger set of graphs. Apart from search time speed up, Placeto on average reduces placed graph run time by 3% (for its different graphs and hardware) while GDP on average reduces placed graph run time by 9.2%, compared to HDP. One advantage of GDP over Placeto is that it does not rely on any initial feasible placement. Providing a reasonable initial placement is often non-trivial for domain experts, especially for larger graphs such as 8-layer GNMT. As such, we are the first to report superhuman results on 8-layer GNMT (with GDP-batch).

Figure 4: Normalized run time (step time for the generated placement) and normalized training time (search time) for fine-tuning. Time is normalized to GDP-one.

Run time comparison between GDP-one, human expert, Tensorflow METIS, and hierarchical device placement (HDP) on six graphs (RNNLM, GNMT, Transformer-XL, Inception, Amoe-baNet, and WaveNet). Graph runtime speed up is compared with Human Placement (HP) and Hierarchical Device Placement (HDP). Search speed up is the policy network training time speed up compared to HDP (reported values are averages of six runs).

Table content: <table><tr><td>Model (#devices)</td><td>GDP-one (s)</td><td>HP (s)</td><td>METIS (s)</td><td>HDP (s)</td><td>Run time speed up over HP / HDP</td><td>Search speed up</td></tr><tr><td>2-layer RNNLM (2)</td><td>0.234</td><td>0.257</td><td>0.355</td><td>0.243</td><td>9.8% / 4%</td><td>2.95x</td></tr><tr><td>4-layer RNNLM (4)</td><td>0.409</td><td>0.48</td><td>OOM</td><td>0.490</td><td>17.4% / 19.8%</td><td>1.76x</td></tr><tr><td>2-layer GNMT (2)</td><td>0.301</td><td>0.384</td><td>OOM</td><td>0.376</td><td>27.6% / 24.9%</td><td>30x</td></tr><tr><td>4-layer GNMT (4)</td><td>0.409</td><td>0.469</td><td>OOM</td><td>0.520</td><td>14.7% / 27.1%</td><td>58.8x</td></tr><tr><td>8-layer GNMT (8)</td><td>0.649</td><td>0.610</td><td>OOM</td><td>0.693</td><td>-6% / 6.8%</td><td>7.35x</td></tr><tr><td>2-layer Transformer-XL (2)</td><td>0.386</td><td>0.473</td><td>OOM</td><td>0.435</td><td>22.5% / 12.7%</td><td>40x</td></tr><tr><td>4-layer Transformer-XL (4)</td><td>0.580</td><td>0.641</td><td>OOM</td><td>0.621</td><td>11.4% / 7.1%</td><td>26.7x</td></tr><tr><td>8-layer Transformer-XL (8)</td><td>0.748</td><td>0.813</td><td>OOM</td><td>0.789</td><td>8.9% / 5.5%</td><td>16.7x</td></tr><tr><td>Inception (2)</td><td>0.405</td><td>0.418</td><td>0.423</td><td>0.417</td><td>3.2% / 3%</td><td>13.5x</td></tr><tr><td>AmoebaNet (4)</td><td>0.394</td><td>0.44</td><td>0.426</td><td>0.418</td><td>26.1% / 6.1%</td><td>58.8x</td></tr><tr><td>2-stack 18-layer WaveNet (2)</td><td>0.317</td><td>0.376</td><td>OOM</td><td>0.354</td><td>18.6% / 11.7%</td><td>6.67x</td></tr><tr><td>4-stack 36-layer WaveNet (4)</td><td>0.659</td><td>0.988</td><td>OOM</td><td>0.721</td><td>50% / 9.4%</td><td>20x</td></tr><tr><td>GEOMEAN</td><td>-</td><td>-</td><td>-</td><td>-</td><td>16% / 9.2%</td><td>15x</td></tr></table>

Run time comparison on GDP-batch vs. GDP-one.

Table content: <table><tr><td>Model</td><td>Speed up</td><td>Model</td><td>Speed up</td></tr><tr><td>2-layer RNNLM</td><td>0</td><td>Inception</td><td>0</td></tr><tr><td>4-layer RNNLM</td><td>5%</td><td>AmoebaNet</td><td>-5%</td></tr><tr><td>2-layer GNMT</td><td>0</td><td>4-stack 36-layer WaveNet</td><td>3.3 %</td></tr><tr><td>4-layer GNMT</td><td>0</td><td>2-stack 18-layer WaveNet</td><td>15%</td></tr><tr><td>2-layer Transformer-XL</td><td>7.6%</td><td>8-layer Transformer-XL</td><td>1.5%</td></tr><tr><td>4-layer Transformer-XL</td><td>3%</td><td/><td/></tr></table>

Hyperparameters for Policy Network. gs layers: GraphSAGE layers, gs knn: Graph-SAGE maximum neighbors, trf d model: Dimension of the TransformerXL model, trf n head: Number of attention heads, trf layers: Number of TransformerXL layers, trf d heads: Dimension of each attention head, trf d inner: Dimension of inner hidden size in positionwise feedforward.

Table content: <table><tr><td>Parameters</td><td>Value</td><td>Parameters</td><td>Value</td></tr><tr><td>gs layers</td><td>4</td><td>gs dim</td><td>128</td></tr><tr><td>gs knn</td><td>5</td><td>trf layers</td><td>2</td></tr><tr><td>trf d model</td><td>128</td><td>trf n head</td><td>5</td></tr><tr><td>trf d head</td><td>25</td><td>trf d inner</td><td>256</td></tr></table>

Hyperparameters for PPO.

Table content: <table><tr><td>Parameters</td><td>Value</td><td>Parameters</td><td>Value</td></tr><tr><td>learing rate</td><td>0.5</td><td>num of rollouts</td><td>400</td></tr><tr><td>minibatches</td><td>40</td><td>epochs</td><td>5</td></tr><tr><td>epsilon</td><td>0.2</td><td>entropy</td><td>0.05</td></tr></table>

Run time comparison on GDP batch training vs. the best of related methods (human expert, METIS, HDP, and GDP no batch training).

Table content: <table><tr><td>Batch Setting</td><td>Model</td><td>speed up (s)</td></tr><tr><td/><td>Inception</td><td>0</td></tr><tr><td/><td>AmoebaNet</td><td>-4.5%</td></tr><tr><td>Batch 2</td><td>2-layer RNNLM 2-layer GNMT</td><td>0 0</td></tr><tr><td/><td>2-layer Transformer-XL</td><td>6.5%</td></tr><tr><td/><td>2-stack 18-layer Wavenet</td><td>4%</td></tr><tr><td/><td>2-layer RNNLM</td><td>0</td></tr><tr><td/><td>4-layer RNNLM</td><td>0</td></tr><tr><td>Batch 3</td><td>8-layer RNNLM 2-layer GNMT</td><td>4.5% 0</td></tr><tr><td/><td>4-layer GNMT</td><td>0</td></tr><tr><td/><td>8-layer GNMT</td><td>8%</td></tr><tr><td>Batch 4</td><td>3x8-layer GNMT</td><td>5.1%</td></tr><tr><td>Batch 5</td><td>3x8-layer RNNLM</td><td>4.5%</td></tr></table>

Peer review (iclr2020):
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. I vote for weak reject since some issues that I would like to see addressed by the author(s) are not. These include the fact that since the goal is to minimize the makespan, scheduling within each machine the operations should be addressed in a better way. Also, while the objective J(	heta) is reasonable, the distribution for the makespans could be very skewed (e.g., heavy tails over the dataflow graphs). Doesn*t this affect the results? Finally, the novelty from a deep learning perspective is limited. - How do the authors address the issue of scheduling the operations within each machine? - How is mathcal{D} formally defined (i.e., the range of the mapping function)? Do you take into account the different number of machines, their memory footprints that can be significantly different, the different processing units they may have (GPUs, CPUs, TPUs)? Is the number of machines used for the partition automatically learned by the policy? That part was not very clear. - Since the authors compare with METIS, it is worth also comparing with Scotch https://www.labri.fr/perso/pelegrin/scotch/ that is also publicly available. - Can the authors comment on the scalability of their method as a function of n (number of nodes), and k (number of devices)? Update: Thanks to the author(s) for the detailed feedback. I have upgraded my score accordingly.