[Paper Review] Device Placement Optimization with Reinforcement Learning

Azalia Mirhoseini, Hieu Pham+8 more|arXiv (Cornell University)|Jun 13, 2017

Industrial Vision Systems and Defect Detection47 references220 citations

TL;DR

The paper learns to optimize TensorFlow device placement for neural nets using a sequence-to-sequence policy optimized by REINFORCE, achieving faster placements than hand-crafted heuristics and Scotch baselines.

ABSTRACT

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.

Motivation & Objective

• Motivate reducing training/inference costs through better device placement in heterogeneous hardware.
• Propose a learned strategy that assigns graph operations to devices to minimize execution time.
• Demonstrate improvements over human-designed placements and traditional graph-partitioning methods across multiple models.

Proposed method

• Model device placement as a discrete optimization over TF graph operations with a policy π(P|G;θ).
• Use a seq2seq model with attention to predict a device for each operation in the graph.
• Train via policy gradients (REINFORCE) using R(P)=sqrt(r(P)) as the reward signal, with a moving average baseline.
• Incorporate co-location groups to reduce sequence length and manage large graphs.
• Implement asynchronous distributed training with multiple controllers and workers to sample and evaluate placements.
• Measure running times by executing placements on actual hardware and use those times as rewards.

Experimental results

Research questions

• RQ1 Can a learned policy outperform hand-crafted and Scotch-based baselines for device placement on TF graphs?
• RQ2 What are the tradeoffs between computation and communication that a learned placement exploits across different models (Inception-V3, NMT, RNNLM)?
• RQ3 How does end-to-end training time and per-step latency change when using RL-based placements versus expert-designed placements?

Key findings

• RL placements find non-trivial configurations that outperform hand-crafted baselines and Scotch on multiple models.
• Single-step run times with RL placements are up to 3.5x faster than baselines.
• End-to-end training with RL-based placements speeds up training by up to about 28% for NMT and ~20% for Inception-V3 compared to expert designs.
• For NMT, RL placement balances computational load across devices better than expert placements, reducing bottlenecks during backpropagation.
• For Inception-V3, RL placement reduces inter-device data copies by colocating parameters with their consumers, enabling faster per-step times under multi-GPU settings.