CDMPP: A Device-Model Agnostic Framework for Latency Prediction of Tensor Programs

Abstract

Deep Neural Networks (DNNs) have shown excellent performance in a wide range of machine learning applications. Knowing the latency of running a DNN model or tensor program on a specific device is useful in various tasks, such as DNN graph- or tensor-level optimization and device selection. Considering the large space of DNN models and devices that impede direct profiling of all combinations, recent efforts focus on building a predictor to model the performance of DNN models on different devices. However, none of the existing attempts have achieved a cost model that can accurately predict the performance of various tensor programs while supporting both training and inference accelerators. We propose CDMPP, an efficient tensor program latency prediction framework for both cross-model and cross-device prediction. We design an informative but efficient representation of tensor programs, called compact ASTs, and a pre-order-based positional encoding method, to capture the internal structure of tensor programs. We develop a domain-adaption-inspired method to learn domain-invariant representations and devise a KMeans-based sampling algorithm, for the predictor to learn from different domains (i.e., different DNN operators and devices). Our extensive experiments on a diverse range of DNN models and devices demonstrate that CDMPP significantly outperforms state-of-the-art baselines with 14.03% and 10.85% prediction error for cross-model and cross-device prediction, respectively, and one order of magnitude higher training efficiency. The implementation and the expanded dataset are available at https://github.com/joapolarbear/cdmpp.

arXiv commentAccepted by EuroSys 2024

References in corpus (9)

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Beyond Data and Model Parallelism for Deep Neural Networks
• Robust Unsupervised Domain Adaptation for Neural Networks via Moment Alignment
• The Non-IID Data Quagmire of Decentralized Machine Learning
• Ansor: Generating High-Performance Tensor Programs for Deep Learning
• A Learned Performance Model for Tensor Processing Units
• Shift-Robust GNNs: Overcoming the Limitations of Localized Graph Training Data
• Margin Calibration for Long-Tailed Visual Recognition