Building a Performance Model for Deep Learning Recommendation Model Training on GPUs

We devise a performance model for GPU training of Deep Learning Recommendation Models (DLRM), which has low GPU utilization (i.e., the percentage of per-batch training time when kernels are running on the device) compared to other well-optimized vision (CV) and natural language processing (NLP) models. We show that both the device active time (the sum of kernel runtimes) and idle time are important components of the overall device time, and can be tackled separately by (1) flexibly adopting heuristic- and ML-based kernel performance models for kernels that dominate the device active time, and (2) categorizing operator overheads into five types to quantitatively determine their contribution to the overall device time. Combining these two parts, we propose a critical-path-based algorithm to predict the per-batch training time of DLRM by traversing its execution graph. We achieve less than 10% geometric mean absolute error (GMAE) in all kernel performance modeling, and 5.23% and 7.96% geomean errors, respectively, for GPU active time and overall end-to-end per-batch training time prediction on the highly-customized and multi-factor dominated DLRM architectures. We also demonstrate our performance model’s ability to generalize to other compute-bound DL models targeted by most previous methods and better assist general model-system co-design than previous work.