Scalable Distributed DNN Training using TensorFlow and CUDA-Aware MPI: Characterization, Designs, and Performance Evaluation

The current wave of advances in Deep Learning (DL) have been triggered by the availability of large-scale datasets, efficient CPU and GPU hardware, and development of software frameworks like TensorFlow (TF). However, little exists in literature that addresses TensorFlow's distributed training capabilities. In this paper, we provide an in-depth performance characterization and design analysis for distributed TensorFlow. We present three key insights: 1) Horovod designs achieve better performance compared to the official gRPC-based approaches, 2) performance of Horovod design is heavily influenced by the time spent in gradient aggregation that uses the Allreduce primitive, and 3) performance of existing Horovod-MPI implementation is significantly worse compared to Horovod-NCCL. To address this limitation in Horovod-MPI, we propose a novel and efficient CUDA-Aware MPI Allreduce design that 1) exploits CUDA kernels to perform large reductions on the GPU, 2) uses a com-bination of bandwidth-optimal and latency-optimal algorithms, and 3) maintains a pointer cache to avoid CUDA-driver query overheads in the critical path. The proposed designs deliver 5×, 17×, and 29% better performance compared to NCCL2 for small, medium, and large messages. Our designs enable Horovod-MPI to beat state-of-the-art Horovod-NCCL2 by 3% and achieve 90% scaling efficiency for ResNet-50 training on 64 Pascal GPUs.