ParLearning 2018 Invited Talk 2.

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 talk, I will present some of our recent efforts on learning to optimize model parallelism for TensorFlow computational graphs. Key to our method is the use of deep reinforcement learning 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 deep model. Our main result is that on important computer vision, language modeling and neural machine translation tasks, our model finds non-trivial ways to parallelise the model that outperform hand-crafted heuristics and traditional algorithmic methods.