Source-linked AI summary
Device Placement Optimization with Reinforcement Learning
Azalia Mirhoseini, Hieu Pham, Quoc V. Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, Jeff Dean
TL;DR
Growing neural-network workloads make device placement across heterogeneous CPUs and GPUs difficult, especially for complex graphs and dynamic environments. The paper trains a sequence-to-sequence model with measured execution time as reinforcement-learning feedback to propose placements. Across image classification, language modeling, and machine translation, it finds non-trivial placements that surpass human-designed and algorithmic baselines.
Problem
Device placement for growing neural networks is often specified by human practitioners, while complex graphs and dynamic environments challenge existing solvers.
Method
A sequence-to-sequence model proposes device placements from operation and dependency information, and measured execution time trains it through reinforcement learning.
Results
Across image classification, language modeling, and machine translation, the method finds non-trivial placements that surpass human experts and optimized algorithmic solvers.
Takeaways & Limitations
The approach learns computation–communication tradeoffs from the hardware environment and applies reinforcement learning to large-scale neural-network device placement.
Abstract
from arXiv · showhide
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.
1. Introduction
Neural networks increasingly require heterogeneous distributed hardware, making device placement difficult for complex graphs and dynamic environments. The paper proposes reinforcement learning that learns placements from execution-time feedback and reports faster placements than human and algorithmic approaches.
- Neural networks’ growing computational requirements motivate using heterogeneous environments that combine CPUs and GPUs.
- Device placement becomes challenging for networks with many branches or larger minibatches, while existing solvers are inflexible in dynamic environments.
- The method uses a sequence-to-sequence model to read operations and dependencies, propose placements, execute them, and train from execution time as reward.
- The method finds non-trivial multi-device placements for Inception-V3, recurrent language modeling, and neural machine translation.
- Up to 3.5 times faster single-step placements and up to 20% faster real-time training were achieved relative to the reported comparison baselines.
2. Related Work
Prior work includes neural reinforcement learning for combinatorial optimization, reinforcement learning for system performance, and graph-partitioning algorithms. The paper distinguishes its approach by optimizing measured running time directly for large-scale computational graphs with noisy rewards.
- Neural combinatorial-optimization methods use sequence-to-sequence models and policy gradients, but earlier experiments focused on toy datasets.
- Related reinforcement-learning work optimizes hand-crafted reward objectives, whereas this paper optimizes configurations’ running time directly.
- Graph-partitioning algorithms use iterative refinement or spectral analysis, but applying them to computational graphs requires costly, inaccurate graph cost models.
- Scotch maps tasks across connected processing nodes by balancing computational load and reducing communication through graph-partitioning techniques.
- Scotch represents hardware resources and computational workloads as target-architecture and source graphs, respectively.
3. Method
The method models device placement as a policy that assigns graph operations to devices and learns from measured execution rewards. It combines an attentional sequence-to-sequence architecture with policy-gradient training, co-location heuristics, and distributed asynchronous evaluation.
- Problem formulation: A placement P assigns each operation in a TensorFlow graph to one of D devices, with the objective of minimizing complete-graph execution time.The reward is based on the measured execution time under the sampled placement.
- Policy-gradient training: Using the square root of running time as R(P) makes learning more robust to noisy early measurements and weak late-stage differences.Failed placements receive a manually specified large failing signal, and late updates are restricted to placements that execute after 5,000 steps.
- Policy-gradient training: The policy is an attentional sequence-to-sequence model trained with Adam and REINFORCE policy gradients.The model defines π(P|G; θ), samples placements, and optimizes its parameters from reward signals.
- Architecture: The encoder embeds operation types, output shapes, and direct adjacency information, while the decoder predicts one device per operation using attention and device embeddings.The decoder has as many time steps as the graph has operations.
- Architecture: Because TensorFlow graphs can contain thousands of operations, the method reduces the placement space by manually forcing related operations into co-location groups.Groups use TensorFlow defaults and recursively merge operations when one operation’s output is consumed only by another.
- Distributed training: Asynchronous distributed training uses multiple controllers, a shared parameter server, and workers that execute sampled placements and measure their running times.Each controller interacts with K workers; experiments use up to 20 controllers with 4 or 8 workers each, taking 12 to 27 hours to find the best placement.
4. Experiments
The experiments evaluate learned placements on RNNLM, Neural MT, and Inception-V3 against strong baselines, measuring both single-step and end-to-end training efficiency. The learned method finds placements that exploit device-specific tradeoffs and improve runtime across these models.
- Experiment Setup: The evaluation covers RNNLM, Neural MT, and Inception-V3 using established baselines and training configurations.The models use 1 CPU with either 2 or 4 Nvidia Tesla K80 GPUs; RNNLM and Neural MT use two LSTM layers.
- Single-Step Runtime Efficiency: The learned placements are on par with or better than other placement methods while learning tradeoffs between parallelism and inter-device communication.The model receives running times and the number of available devices rather than hand-designed intermediate cost information.
- Single-Step Runtime Efficiency: More than twice faster than the best published human-designed baseline, the learned placement fits the entire RNNLM graph on one GPU to avoid inter-device communication latency.This placement chooses co-location on one GPU rather than distributing the graph across devices.
- Single-Step Runtime Efficiency: For Neural MT, the method achieves up to 20.6% speedup with 4 GPUs by placing less computationally expensive operations such as embedding lookups on the CPU.The authors suggest this reduces the load on the GPUs.
- Single-Step Runtime Efficiency: With 4 GPUs, Inception-V3 runtime falls from 4.60 seconds to 3.85 seconds, whereas neither baseline improves on assigning all operations to one GPU.With only 2 GPUs, the placer also assigns all operations to one GPU because model-parallel freedom is limited.
- End-to-End Runtime Efficiency: End-to-end training is 27.8% faster for Neural MT and 19.7% faster for Inception-V3 than the specified expert-designed or synchronous-tower baselines.Neural MT training takes 165.73 hours instead of 229.57 hours; Inception-V3's synchronous approach converges faster initially, but the RL curve eventually crosses the asynchronous-tower curve.
- Runtime Analysis: The method balances Neural MT workloads better, while Inception-V3 gains runtime mainly by reducing data copying because its dependencies limit model parallelism.The authors associate the Inception-V3 improvement with keeping model parameters on the same device as operations that use them.
5. Conclusion
The paper presents an adaptive sequence-to-sequence method that learns device placements from execution time and surpasses human-designed and algorithmic alternatives across several tasks.
- The method uses a sequence-to-sequence model to propose device placements from neural-network operations.The model is trained to optimize execution time.
- The approach learns properties of the hardware environment, including the computation–communication tradeoff.
- Across image classification, language modeling, and machine translation, the method surpasses human-designed placements and optimized algorithmic solvers.