Device Placement Optimization with Reinforcement Learning

By Azalia Mirhoseini, Hieu Pham, Quoc V. Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, Joseph Dean. International Conference on Machine Learning, 2017

492 citations recorded by the source corpus.

The past few years have witnessed a growth in size and computational requirements for training and inference with neural networks. Currently, a common approach to address these requirements is to use a heterogeneous distributed environment with a mixture of hardware devices such as CPUs and GPUs. Importantly, the decision of placing parts of the neural models on devices is often made by human experts based on simple heuristics and intuitions. In this paper, we propose a method which learns to optimize device placement for TensorFlow computational graphs. Key to our method is the use of a sequence-to-sequence model to predict which subsets of operations in a TensorFlow graph should run on which of the available devices. The execution time of the predicted placements is then used as the reward signal to optimize the parameters of the sequence-to-sequence model. Our main result is that on Inception-V3 for ImageNet classification, and on RNN LSTM, for language modeling and neural machine translation, our model finds non-trivial device placements that outperform hand-crafted heuristics and traditional algorithmic methods.

Subfields: Artificial Intelligence, Industrial and Manufacturing Engineering, Mechanical Engineering

View Paper: https://www.semanticscholar.org/paper/bfbd10ebffc9494423770a5bd30ebd0f9cbce66d

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Industrial and Manufacturing Engineering

In Mechanical Engineering

Citations in the corpus

MiCS: Near-linear Scaling for Training Gigantic Model on Public Cloud (2022, background): “ColocRL [34] formulates the distributed training as a placement optimization problem to maximize the throughput.”
Learning Decision Trees with Reinforcement Learning (2017, background): “This framework has obtained promising solutions on a variety of problems, like neural architecture search [18], neural optimizer search [3], combinatorial optimization [2] and device placement [10].”
Efficient and Robust Parallel DNN Training through Model Parallelism on Multi-GPU Platform (2018, methodology): “A few prior works have addressed this issue(Harlap et al., 2018; Mirhoseini et al., 2017).”
TASO: optimizing deep learning computation with automatic generation of graph substitutions (2019, methodology): “ColocRL [26] uses reinforcement learning to automatically discover an efficient device placement for parallelizing DNN training across multiple GPUs.”
Optimizing Federated Learning on Non-IID Data with Reinforcement Learning (2020, background): “The recent success of RL [12, 13] comes from the capability that RL agents learn from the interaction with a dynamic environment.”
Tiresias: A GPU Cluster Manager for Distributed Deep Learning (2019, methodology): “To this end, possible approaches include topology-aware schemes [10] and ﬁne-grained placement of computational graphs in DL jobs [32].”
Software-hardware co-design for fast and scalable training of deep learning recommendation models (2021, methodology): “Building on this direction of auto-parallelization, these recent papers [39, 60] use optimal synthesis and reinforcement learning to find optimized device placement to further improve parallelism without the need for manual intervention.”
Deep Learning in Mobile and Wireless Networking: A Survey (2018, background): “Their solution achieves up to 20% faster computation speed than human experts’ designs of such placements [119].”
SwapAdvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping (2020, background): “However, we have yet to explore certain other promising search methods such as reinforcement learning [31, 32].”
Dive into Deep Learning (2020, background): “, (Mirhoseini et al., 2017) for a best effort in this regard.”
Data Movement Is All You Need: A Case Study on Optimizing Transformers (2020, background): “…2019d;b;a; Chen et al., 2018; Paul G. Allen School of Computer Science & Engineering, University of Washington et al., 2017; Sivathanu et al., 2019; Mirhoseini et al., 2017; Cyphers et al., 2018; Baghdadi et al., 2019; Vasilache et al., 2018; Lethin, 2017; Wei et al., 2018; Truong et al., 2016;…”
HARK Side of Deep Learning - From Grad Student Descent to Automated Machine Learning (2019, methodology): “The methods behind such meta-learning approaches are mainly based on Bayesian optimization [48], evolutionary algorithms [43, 46] or more recently on reinforcement learning [41, 49, 52].”
AutoLoss: Learning Discrete Schedules for Alternate Optimization (2018, methodology): “Also of note is another line of works that apply RL to enable automatic machine learning (AutoML), such as device placement optimization (Mirhoseini et al., 2017), neural architecture search (Baker et al., 2016; Zoph and Le, 2016), etc.”
Explainable AI (XAI): Core Ideas, Techniques, and Solutions (2022, methodology): “In the model parallel approach [42, 43] the same model is split and placed on different worker nodes, while processing the same copy of data set.”