Enabling AI Deep Potentials for Ab Initio-quality Molecular Dynamics Simulations in GROMACS

TL;DR

This paper integrates AI deep potentials into GROMACS for efficient, high-quality molecular dynamics simulations, evaluating two models and optimizing GPU performance for large-scale biomolecular systems.

Contribution

It introduces a method to incorporate AI deep potentials into GROMACS, enabling ab initio-quality simulations with optimized GPU inference for different deep learning models.

Findings

DPA2 outperforms DPA3 in GPU throughput

GPU kernel-launch overhead is a key bottleneck

Domain-decomposed inference improves performance

Abstract

State-of-the-art AI deep potentials provide ab initio-quality results, but at a fraction of the computational cost of first-principles quantum mechanical calculations, such as density functional theory. In this work, we bring AI deep potentials into GROMACS, a production-level Molecular Dynamics (MD) code, by integrating with DeePMD-kit that provides domain-specific deep learning (DL) models of interatomic potential energy and force fields. In particular, we enable AI deep potentials inference across multiple DP model families and DL backends by coupling GROMACS Neural Network Potentials with the C++/CUDA backend in DeePMD-kit. We evaluate two recent large-atom-model architectures, DPA2 that is based on the attention mechanism and DPA3 that is based on GNN, in GROMACS using four ab initio-quality protein-in-water benchmarks (1YRF, 1UBQ, 3LZM, 2PTC) on NVIDIA A100 and GH200 GPUs. Our results show that DPA2 delivers up to 4.23x and 3.18x higher throughput than DPA3 on A100 and GH200 GPUs, respectively. We also provide a characterization study to further contrast DPA2 and DPA3 in throughput, memory usage, and kernel-level execution on GPUs. Our findings identify kernel-launch overhead and domain-decomposed inference as the main optimization priorities for AI deep potentials in production MD simulations.