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

TL;DR

The paper addresses bringing ab initio-quality MD to production workflows by integrating AI Deep Potentials (via DeePMD-kit) into GROMACS. It introduces a tight C++ coupling of DeePMD-kit with GROMACS NNPot that supports multiple DP families (DPA2 and DPA3) and backends, enabling AIMD-like accuracy with significantly reduced cost. Through four solvated protein benchmarks on A100 and GH200 GPUs, the study shows DPA2 provides substantially higher throughput than DPA3 (up to roughly 4–3x), while also revealing distinct kernel- and memory-bound characteristics that guide optimization priorities. The work establishes end-to-end performance baselines for AI MD in production codes and highlights directions such as GPU-graph optimizations and distributed inference to scale ab initio-quality simulations for biomolecular systems, with throughput scaling approximately as $O(N)$ rather than the $O(N^3)$ of DFT.

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.

Figures (5)

• Figure 1: An overview of the GROMACS MD simulation loop.
• Figure 2: An overview of the GROMACS MD architecture with the DeePMD-kit integration.
• Figure 3: Snapshots of 1YRF protein after 10 ps MD simulations. The initial structure (leftmost) is compared against trajectories from DPA2-solvated_protein, DPA2-SPICE2 , DPA3-solvated_protein, and the empirical GROMACS force field (rightmost). Note that deviations in local atomic positions between DP models and the empirical field are expected, as DP models incorporate ab initio-level quantum effects.
• Figure 4: MD simulation throughput for classical approach, ab initio with DPA2, and ab initio with DPA3 for four protein simulations on A100 and GH200. DPA3 results marked with $\times$ are unavailable due to OOM.
• Figure 5: GPU memory usage of DPA2 and DPA3 models on A100 and GH200. The red dashed line represents the GPU capacity. Results marked with $\times$ are unavailable due to OOM.