Optimizing DNN Compilation for Distributed Training with Joint OP and Tensor Fusion

Abstract

This paper proposes DisCo, an automatic deep learning compilation module for data-parallel distributed training. Unlike most deep learning compilers that focus on training or inference on a single device, DisCo optimizes a DNN model for distributed training over multiple GPU machines. Existing single-device compilation strategies do not work well in distributed training, due mainly to communication inefficiency that they incur. DisCo generates optimized, joint computation operator and communication tensor fusion strategies to enable highly efficient distributed training. A GNN-based simulator is built to effectively estimate per-iteration training time achieved by operator/tensor fusion candidates. A backtracking search algorithm is driven by the simulator, navigating efficiently in the large strategy space to identify good operator/tensor fusion strategies that minimize distributed training time. We compare DisCo with existing DL fusion schemes and show that it achieves good training speed-up close to the ideal, full computation-communication overlap case.

References in corpus (18)

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• TensorFlow: A system for large-scale machine learning
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• Session-based Recommendation with Graph Neural Networks
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• SimplE Embedding for Link Prediction in Knowledge Graphs
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions
• Glow: Graph Lowering Compiler Techniques for Neural Networks
• MLIR: A Compiler Infrastructure for the End of Moore's Law
• Intel nGraph: An Intermediate Representation, Compiler, and Executor for Deep Learning
• A Latent Variable Recurrent Neural Network for Discourse Relation Language Models
• Ithemal: Accurate, Portable and Fast Basic Block Throughput Estimation using Deep Neural Networks
• A simple yet effective baseline for non-attributed graph classification
• A Learned Performance Model for Tensor Processing Units
• A Deep Learning Based Cost Model for Automatic Code Optimization
• Parallax: Sparsity-aware Data Parallel Training of Deep Neural Networks
• FusionStitching: Deep Fusion and Code Generation for Tensorflow Computations on GPUs