Habitat: A Runtime-Based Computational Performance Predictor for Deep Neural Network Training
Geoffrey X. Yu, Yubo Gao, Pavel Golikov, Gennady Pekhimenko
Abstract
Deep learning researchers and practitioners usually leverage GPUs to help train their deep neural networks (DNNs) faster. However, choosing which GPU to use is challenging both because (i) there are many options, and (ii) users grapple with competing concerns: maximizing compute performance while minimizing costs. In this work, we present a new practical technique to help users make informed and cost-efficient GPU selections: make performance predictions with the help of a GPU that the user already has. Our technique exploits the observation that, because DNN training consists of repetitive compute steps, predicting the execution time of a single iteration is usually enough to characterize the performance of an entire training process. We make predictions by scaling the execution time of each operation in a training iteration from one GPU to another using either (i) wave scaling, a technique based on a GPU's execution model, or (ii) pre-trained multilayer perceptrons. We implement our technique into a Python library called Habitat and find that it makes accurate iteration execution time predictions (with an average error of 11.8%) on ResNet-50, Inception v3, the Transformer, GNMT, and DCGAN across six different GPU architectures. Habitat supports PyTorch, is easy to use, and is open source. 1
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 31a8e1cd-eaa5-4a0f-9068-57d8235278d1Cited by top-tier papers16
- HybridFlow: A Flexible and Efficient RLHF FrameworkGuangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu et al.EuroSys 2025 · 61 citations
- Lyra: Elastic Scheduling for Deep Learning ClustersJiamin Li, Hong Xu, Yibo Zhu, Zherui Liu et al.EuroSys 2023 · 59 citations
- Mobile Foundation Model as FirmwareJinliang Yuan, Chen Yang, Dongqi Cai, Shihe Wang et al.MobiCom 2024 · 40 citations
- Forecasting GPU Performance for Deep Learning Training and InferenceSeonho Lee, Amar Phanishayee, Divya MahajanASPLOS 2025 · 31 citations
- Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-OptimizationZhanda Zhu, Christina Giannoula, Muralidhar Andoorveedu, Qidong Su et al.EuroSys 2025 · 8 citations
Builds on4
- Heterogeneity-Aware Cluster Scheduling Policies for Deep Learning WorkloadsDeepak Narayanan, Keshav Santhanam, Fiodar Kazhamiaka, Amar Phanishayee et al.OSDI 2020 · 286 citations
- Balancing efficiency and fairness in heterogeneous GPU clusters for deep learningShubham Chaudhary, Ramachandran Ramjee, Muthian Sivathanu, Nipun Kwatra et al.EuroSys 2020 · 135 citations
- Daydream: Accurately Estimating the Efficacy of Optimizations for DNN TrainingHongyu Zhu, Amar Phanishayee, Gennady PekhimenkoUSENIX ATC 2020 · 74 citations
- Skyline: Interactive In-Editor Computational Performance Profiling for Deep Neural Network TrainingGeoffrey X. Yu, Tovi Grossman, Gennady PekhimenkoUIST 2020 · 16 citations
Related papers
- Centimani: Enabling Fast AI Accelerator Selection for DNN Training with a Novel Performance PredictorZhen Xie, Murali Emani, Xiaodong Yu, Dingwen Tao et al.USENIX ATC 2024 · 4 citations
- Path Forward Beyond Simulators: Fast and Accurate GPU Execution Time Prediction for DNN WorkloadsYing Li, Yifan Sun, Adwait JogMICRO 2023 · 16 citations
- ADA-GP: Accelerating DNN Training By Adaptive Gradient PredictionVahid Janfaza, Shantanu Mandal, Farabi Mahmud, Abdullah MuzahidMICRO 2023 · 3 citations
- Prediction-Assisted Online Distributed Deep Learning Workload Scheduling in GPU ClustersZiyue Luo, Jia Liu, Myungjin Lee, Ness B. ShroffINFOCOM 2025 · 5 citations
- Orion: Interference-aware, Fine-grained GPU Sharing for ML ApplicationsFoteini Strati, Xianzhe Ma, Ana KlimovicEuroSys 2024 · 96 citations
