Efficient Deep Molecular Dynamic Model Training on Heterogeneous System

Molecular dynamics is a widely adopted simulation method for analyzing the movement of atoms and molecules. Traditional molecular dynamics simulation methods are computationally intensive and difficult to simulate a large number of atoms. In contrast, molecular dynamics based on deep potential models such as DeePMD can leverage deep learning techniques to improve simulation efficiency. Although DeePMD has incorporated mainstream deep learning frameworks, it still suffers from low performance and efficiency during its model training on heterogeneous systems such as CPU and GPU. Particularly, a large number of operators cannot be accelerated by GPU, resulting in low utilization of GPU computational resources. In this paper, we comprehensively analyze the computational bottlenecks and the corresponding root causes of DeePMD. We correspondingly propose several novel optimization strategies. Specifically, for preprocessing, we identify the computation redundancies and the GPU parallelization opportunities for performance optimization. For training, we propose optimization strategies such as operator fusion, redundancy elimination, and concurrent execution of multiple streams and threads in the computation process. Moreover, we apply systematical optimization of computational graphs and operators. The evaluation results show that DeePMD can achieve significant speedups in several cases after applying our proposed optimizations, resulting in a maximum overall speedup of 6.36× with acceptable accuracy.