We introduce a hierarchical model for efficient placement of computational graphs onto hardware devices, especially in heterogeneous environments with a mixture of CPUs, GPUs, and other computational devices. Our method learns to assign graph operations to groups and to allocate those groups to available devices. The grouping and device allocations are learned jointly. The proposed method is trained with policy gradient and requires no human intervention. Experiments with widely-used
 computer vision and natural language models show that our algorithm can find optimized, non-trivial placements for TensorFlow computational graphs with over 80,000 operations. In addition, our approach outperforms placements by human
 experts as well as a previous state-of-the-art placement method based on deep reinforcement learning. Our method achieves runtime reductions of up to 60.6% per training step when applied to models such as Neural Machine Translation. Deep neural networks have been successfully applied to many practical problems, such as image classification BID20 BID19 BID32 BID29 , speech recognition BID9 , and machine translation BID28 BID1 BID36 . These successes have lead to a surge in demand for the computational resources needed to train and infer with neural networks. A common approach to addressing this demand is to use a distributed environment with a combination of CPUs and GPUs. In this environment, it is typical for a machine learning practitioner to explicitly place the operations of their neural network onto particular computing devices for model parallelism and data parallelism. For example, one might distribute the computation of the first layer in a translation network onto the first GPU and the computation of the second layer onto the second GPU BID28 BID36 . Although these decisions can be made by a human practitioner, such an approach does not scale well or produce optimal results, especially in the case of more complicated networks BID31 a) . Given the growing diversity of hardware devices (e.g., Google TPUs, Intel Nervana, etc.) and recent trends toward automated neural architecture search BID38 BID27 BID2 , where new models are generated, trained and evaluated in an entirely end-to-end fashion, it seems natural to move toward more automated solutions for efficiently distributing computation.Device placement can be framed as the problem of learning to partition a graph across available devices. Given that graph partitioning is a well-studied subject in computer science BID8 BID16 BID26 , traditional graph partitioning methods represent a natural baseline for automated device placement. We ran experiments using Scotch BID26 , a well-established open source library for graph partitioning, which includes optimizations such as k-way Fiduccia-Mattheyses BID8 , Multilevel methods BID3 BID11 BID15 , the Band Method BID5 , the Diffusion Method BID23 , and Dual Recursive Bipartitioning Mapping BID25 . The objective was to balance the computational load across a set of connected processing nodes, while colocating neighboring nodes to minimize communication cost. Despite its promise, this approach yielded disappointing results, likely due to the non-stationarity of the cost function. We target a distributed environment where we use a shared cluster of CPUs and GPUs, and our CPUs may also serve other jobs at the same time. Thus, while cost-based models such as BID21 provide a strong baseline for memory optimizations, since memory usage is deterministic, they cannot be directly applied to environments with dynamic costs.Using deep networks and reinforcement learning for combinatorial optimization has already been proposed BID33 BID4 BID22 . Recent work BID22 ) uses a recurrent neural network (RNN) policy network to predict the placement of operations in a computational graph, optimizing for speed of computation using policy gradient methods. While this approach outperforms traditional graph partitioning heuristics and human expert placements, it is prohibitively expensive for the RNN policy to learn when the number of operations is large. This method is therefore limited to small graphs (with fewer than 1000 nodes) and requires human experts to manually partition the graph into collocation groups as a pre-processing step in order to scale to larger graphs. We refer to the method in BID22 as ColocRL.In this paper, we propose a more flexible approach which learns to optimize device placement for training neural networks that have tens of thousands of operations with no need for manual grouping. Our method consists of a two-level hierarchical network, in which the first model groups the operations of the graph (the Grouper) and the second model places those groups onto devices (the Placer). The Grouper is a feed forward network which reads in information about each operation and its context within the graph, in order to predict the group to which that operation should be assigned. The Placer is a sequence-to-sequence model BID28 that reads in the embedding of the group and predicts the device placement for that group. The entire two-level network is trained jointly using reinforcement learning to optimize for speed of computation and for feasibility (e.g., having sufficient memory available on each device for the computation assigned). Unlike the previous work, our method is end-to-end and does not require human experts to manually group operations as a pre-processing step, making it a fully automated solution to optimizing device placement.Our main result is that our model effectively handles very large graphs and finds non-trivial placements on multiple devices for models such as Inception-V3 BID31 , ResNet BID10 , Language Modeling BID14 , and Neural Machine Translation BID36 . The placements found by our model outperform TensorFlow's default placements BID0 , the Scotch algorithm's placements, and human expert placements, as well as those of ColocRL BID22 ). Our results demonstrate that the proposed approach learns the properties of the environment, including the complex tradeoff between computation and communication in hardware. For example, on a Neural Machine Translation model, our method achieves a 60.6% reduction in training time per iteration. In this paper, we present a hierarchical method for efficiently placing the operations of a computational graph onto devices. Our approach consists of a hierarchical model that first assigns the operations to groups and then places those groups onto devices. We use a policy gradient method to optimize the parameters of the planner. The proposed method enables us to scale to computational graphs containing over 80,000 operations. Unlike previous work, our method is end-to-end and requires no manual effort. On a range of tasks including image classification, language modeling, and machine translation, our method surpasses placements designed by human experts as well as those of previous state-of-the-art deep RL methods. Our approach finds highly granular parallelism within the graph, enabling us to outperform prior methods by up to 60.6%.
We introduce a hierarchical model for efficient, end-to-end placement of computational graphs onto hardware devices. Proposes to jointly learn groups of operators to colocate and to place learned groups on devices to distribute operations for deep learning via reinforcement learning. The authors purpose a fully connect network to replace the co-location step in an auto-placement method proposed to accelerate a TensorFlow model's runtime. Proposes a device placement algorithm to place operations of tensorflow on devices.