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

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.

💡 Research Summary

This paper presents a comprehensive integration of state‑of‑the‑art AI deep potential (DP) models into the production‑grade molecular dynamics engine GROMACS, leveraging the C++/CUDA backend of DeePMD‑kit. The authors first identify a fundamental mismatch: most DP inference pipelines are Python‑centric and rely on deep‑learning frameworks, whereas GROMACS is a high‑performance C++/CUDA/MPI code built around classical force fields. To bridge this gap, they extend GROMACS’s Neural Network Potentials (NNPot) module with a new backend that directly calls DeePMD‑kit’s C++/CUDA inference API, enabling support for all DeePMD model families (including DP‑A, DP‑A2, DP‑A3) and multiple DL frameworks (PyTorch, TensorFlow, JAX, etc.).

Two recent large‑atom‑model (LAM) architectures are evaluated: DP‑A2, which uses an attention‑based transformer‑like descriptor, and DP‑A3, which employs a graph neural network (GNN) with line‑graph series. Both models require a broader “ghost‑atom” region than traditional MD neighbor lists, and they are non‑linear, demanding full atom information on each sub‑domain. The authors solve these challenges by (1) expanding the ghost‑atom halo to L × Rcut, (2) performing an extra communication step before inference to rebuild symmetric ghost topologies, and (3) aggregating atoms from all MPI ranks to a single rank for inference, then redistributing forces. Although this single‑rank inference limits scalability, it simplifies domain‑decomposition handling and provides a baseline for future parallel inference extensions.

Performance is benchmarked on four solvated proteins (1YRF, 1UBQ, 3LZM, 2PTC) with atom counts ranging from ~600 to ~4100, using NVIDIA A100 and the newer GH200 GPUs. DP‑A2 consistently outperforms DP‑A3, achieving up to 4.23× higher throughput on A100 and 3.18× on GH200. Memory profiling shows DP‑A2 uses less GPU memory per atom, and kernel‑level profiling reveals that the dominant overhead stems from frequent kernel launches with small batch sizes, as well as sub‑optimal memory access patterns. The current implementation’s single‑rank inference also emerges as a scalability bottleneck for larger systems.

Based on these observations, the authors outline three primary optimization priorities: (i) reducing kernel‑launch overhead by increasing batch sizes or fusing operations into custom CUDA kernels; (ii) enabling true parallel inference across MPI ranks, which would require a more sophisticated domain‑decomposed inference scheme and efficient ghost‑atom exchange; and (iii) improving memory bandwidth utilization through prefetching, streaming, and better data layout.

In summary, the work demonstrates that AI deep potentials can be tightly coupled with a mainstream MD package to deliver near‑ab‑initio accuracy at a fraction of the cost of density‑functional‑theory AIMD. By providing an open‑source implementation and a detailed performance characterization, the paper establishes a practical pathway for high‑fidelity, large‑scale biomolecular simulations and sets the stage for future scaling and optimization efforts.