Parallelization Strategies for DeepMD-Kit Using OpenMP: Enhancing Efficiency in Machine Learning-Based Molecular Simulations

2025/08/04 by Qi Du, Feng Wang, Chengkun Wu
Materials Science · Biochemistry, Genetics and Molecular Biology · Decision Sciences · #Machine Learning in Materials Science #Advanced Biosensing Techniques and Applications #Scientific Computing and Data Management

paper · doi:10.1109/tc.2025.3595078

Abstract

DeepMD-kit enables deep learning-based molecular dynamics (MD) simulations that require efficient parallelization to leverage modern HPC architectures. In this work, we optimize DeepMD-kit using advanced OpenMP strategies to improve scalability and computational efficiency on an ARMv8 processor-based server. Our optimizations include data parallelism for neural network inference, force calculation acceleration, NUMAaware memory management, and synchronization reductions, leading to up to 4.1× speedup and 82% higher memory band-width efficiency compared to the baseline implementation. Strong scaling analysis demonstrates superlinear speedup at mid-range core counts, with improved workload balancing and vectorized computations. However, challenges remain at ultra-large scales due to increasing synchronization overhead.

Citations

• Bridging the gap between programming languages and hardware weak memory models
• LAMMPS - a flexible simulation tool for particle-based materials modeling at the atomic, meso, and continuum scales

Related

• Coupling streaming AI and HPC ensembles to achieve 100-1000x faster biomolecular simulations
• Implementing a neural network interatomic model with performance portability for emerging exascale architectures
• Scaling Deep Learning Training with MPMD Pipeline Parallelism
• Integrating State of the Art Compute, Communication, and Autotuning Strategies to Multiply the Performance of the Application Programm CPMD for Ab Initio Molecular Dynamics Simulations
• 29-Billion Atoms Molecular Dynamics Simulation With Ab Initio Accuracy on 35 Million Cores of New Sunway Supercomputer