A Computational-Graph Partitioning Method for Training Memory-Constrained DNNs

Abstract

We propose P DNN, an automatic, generic, and non-intrusive partitioning strategy for large DNN models that do not t into single device memory. P DNN decides a placement of DNN's underlying computational graph operations across multiple devices so that the devices' memory constraints are met and the training time is minimized. P DNN is completely independent of the deep learning aspects of a DNN and requires no modi cation neither at the model nor at the systems level implementation of operation kernels. It partitions DNNs having billions of parameters and hundreds of thousands of operations in seconds to few minutes. Our experiments with TensorFlow on 16 GPUs demonstrate e cient training of 5 very large models while achieving superlinear scaling for both the batch size and training throughput. In comparison to related work (Mesh-TensorFlow and Gradient Checkpointing), P DNN either outperforms or qualitatively improves upon them.