Making Room for AI: Multi-GPU Molecular Dynamics with Deep Potentials in GROMACS

Luca Pennati; Andong Hu; Ivy Peng; Lukas Müllender; Stefano Markidis

cs.DCApr 8, 2026

Making Room for AI: Multi-GPU Molecular Dynamics with Deep Potentials in GROMACS

Authors: Luca Pennati, Andong Hu, Ivy Peng, Lukas Müllender, Stefano Markidis

Organizations: KTH Royal Institute of Technology, Stockholm, Sweden

Abstract

GROMACS is a de-facto standard for classical Molecular Dynamics (MD). The rise of AI-driven interatomic potentials that pursue near-quantum accuracy at MD throughput now poses a significant challenge: embedding neural-network inference into multi-GPU simulations retaining high-performance. In this work, we integrate the MLIP framework DeePMD-kit into GROMACS, enabling domain-decomposed, GPU-accelerated inference across multi-node systems. We extend the GROMACS NNPot interface with a DeePMD backend, and we introduce a domain decomposition layer decoupled from the main simulation. The inference is executed concurrently on all processes, with two MPI collectives used each step to broadcast coordinates and to aggregate and redistribute forces. We train an in-house DPA-1 model (1.6 M parameters) on a dataset of solvated protein fragments. We validate the implementation on a small protein system, then we benchmark the GROMACS-DeePMD integration with a 15,668 atom protein on NVIDIA A100 and AMD MI250x GPUs up to 32 devices. Strong-scaling efficiency reaches 66% at 16 devices and 40% at 32; weak-scaling efficiency is 80% to 16 devices and reaches 48% (MI250x) and 40% (A100) at 32 devices. Profiling with the ROCm System profiler shows that >90% of the wall time is spent in DeePMD inference, while MPI collectives contribute <10%, primarily since they act as a global synchronization point. The principal bottlenecks are the irreducible ghost-atom cost set by the cutoff radius, confirmed by a simple throughput model, and load imbalance across ranks. These results demonstrate that production MD with near ab initio fidelity is feasible at scale in GROMACS.

Figures & tables

Fig. 1: Structure of the 1HCI protein (15,668 atoms) obtained after 500 time steps of a multi-GPU domain decomposed GROMACS-DeePMD simulation, using an in-house DPA-1 model.

Fig. 2: General deep model architecture.

Z_{i}

denotes the atom type,

\mathcal{R}^{i}

atom positions,

\mathcal{D}^{i}

the atom descriptor, and

e_{i}

the atom energy.

Panel of Figure 3 — Fig. 3: DP-SE (a), DPA-1 (b), DPA-2 (c), and DPA-3 (d) descriptor architectures. $Z_{i}$ denotes the atom type, $\mathcal{R}^{i}$ atom positions, and $\mathcal{D}^{i}$ the atom descriptor. The standard DPA-3 configuration (d) uses a graph series of order 2.

Fig. 4: Neighbor list in case of domain decomposition. Atoms ’A’ and ’B’ are local in subdomain 1, while atoms ’C’ and ’D’ belong to subdomain 2.

Fig. 5: GROMACS MD simulation main loop. The conceptual step order is: (1) initialization, (2) domain decomposition and load balancing, (3) position exchange, (4) neighbor-list construction, (5) interaction evaluation, (6) special-force, (7) force reduction and atom update, and (8) output. The actual implementation pipelines several of these stages.

Fig. 6: DeePMD-kit integration in the GROMACS MD engine in the case of DD simulations, example with four MPI ranks.

	System-1	System-2
GROMACS	v2025.2	v2025.2
DeePMD-kit	v3.1.0	v3.1.0
PyTorch	v2.7.1	v2.7.1
Compiler	GCC 13.2.1	GCC 13.3.0
SDK	ROCm 6.3.3	CUDA 12.6.0
MPI	Cray-MPICH 8.1.31	OpenMPI 5.0.3

TABLE I: Software experimental setup.

	Small Protein 1YRF			Large Protein 1HCI
	EM	NVT/NPT	MD	EM	NVT/NPT	MD
$\Delta$ t [fs]	-	2	2	-	2	2
Steps	-	40,000	10,000	-	40,000	200
$r_{c}$ [nm]	1.2	1.2	0.8	1.2	1.2	0.8
DP model	No	No	DPA-1	No	No	DPA-1
DP Group	-	-	Protein	-	-	Protein

TABLE II: Protein simulations parameters.

Fig. 7: Evolution of the force RMSE during training of the DPA-1 model. The error on the training dataset is shown in blue, while the error computed against the validation dataset is reported in orange.

Fig. 8: Comparison between the protein gyration radii about the three Cartesian axes x, y, z obtained with an MD simulation with the DPA-1 model (solid curves) and a classical MD simulation (dashed curves).

Fig. 9: Memory footprint and performance overhead of a GROMACS-DeePMD simulation (red) with respect to a classical MD simulation (blue). Tests are performed on System-1, using the small 1YRF protein, with one MPI process and one GPU.

Fig. 10: Strong scaling test on NVIDIA A100 and AMD MI250x GPUs for one 1HCI protein. The left y axis shows the measured simulation throughput vs the number of GPUs/GCDs (bars), and the throughput computed according to Eq. 8 . The right y axis reports the scaling efficiency. The simulation cannot be run on four NVIDIA A100 GPUs due to insufficient VRAM.

Fig. 11: Weak scaling test on NVIDIA A100 and AMD MI250x GPUs for the 1HCI protein. The left y axis shows the measured simulation throughput vs the number of GPUs/GCDs (bars), the right y axis reports the scaling efficiency.

Fig. 12: Trace of one MD simulation step obtained with the ROCm Systems Profiler on the system equipped with AMD MI250x GPUs during a simulation on 16 ranks. The trace shows the main CPU thread timeline, the HIP GPU queue, and the PyTorch main CPU thread. The block NNPotForceProvider::calculateForces shows the total time spent in the NNPot module. The standard MD operations are barely visible in the main timeline and are highlighted in the d1) box. DeepmdModel::evaluateModel accounts for the time spent in the inference task. All inference kernels are launched with non-blocking APIs, only the final device-to-host copy uses a blocking call. For this reason the trace shows long time spent in the hipMemcpyWithStream process. In reality, as indicated in the HIP queue timeline, the GPU is executing other kernels. The actual time required by the device-to-host copy is shown in the d2) box.

Molecular Dynamics Simulation

Molecular dynamics simulation models atomic motion to study molecular structure, material behavior, reactions, and thermodynamic properties.

See all 55 papers

Jul 31, 2026

Implicit Machine Learning Force Fields Accelerate Molecular Dynamics Simulations

We introduce implicit machine learning force fields (I-MLFFs), which replace explicit stacks of neural network layers with self-consistent fixed-point equations. In molecular simulations, this formulation enables intermediate representations to be reused across successive timesteps, thereby warm-starting force evaluation. The resulting models effectively combine the computational footprint of a shallow, single-layer MLFF with the representational capacity and accuracy of a deep neural network. Our approach unlocks architecture-agnostic efficiency gains that are inaccessible when force prediction and trajectory integration are considered separately. We demonstrate this across three major classes of graph neural networks: invariant, equivariant Cartesian tensor, and SO(3)-equivariant spherical-tensor architectures. Each yields a two- to five-fold reduction in compute and memory footprint. Crucially, these gains are achieved while retaining full atomistic resolution and the original integration timestep, avoiding spatial or temporal coarse graining. Our contribution therefore advances the scaling frontier of quantum-mechanically faithful molecular simulation, enabling longer trajectories and larger atomistic systems within fixed GPU memory and compute budgets, and thereby opening access to new insights across biomolecular and material systems.

Johannes Maeß, Leon Werner, J. Thorben Frank +5

BIFOLD, Berlin, Germany. · Machine Learning Group, TU Berlin, Berlin, Germany. · Department of Artificial Intelligence, Korea University, Seoul, South Korea. +1

Implicit Neural Representations Machine Learning Interatomic Potentials Molecular Dynamics Simulation

Ask PDF

Apr 17, 2026

Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials

Universal Machine Learning Interatomic Potentials (uMLIPs), pre-trained on massively diverse datasets encompassing inorganic materials and organic molecules across the entire periodic table, serve as foundational models for quantum-accurate physical simulations. However, uMLIP training requires second-order derivatives, which lack corresponding parallel training frameworks; moreover, scaling to the billion-parameter regime causes explosive growth in computation and communication overhead, making its training a tremendous challenge. We introduce MatRIS-MoE, a billion-parameter Mixture-of-Experts model built upon invariant architecture, and {Janus}, a pioneering high-dimensional distributed training framework for uMLIPs with hardware-aware optimizations. Deployed across two Exascale supercomputers, our code attains a peak performance of 1.2/1.0 EFLOPS (24%/{35.5%} of theoretical peak) in single precision at over 90% parallel efficiency, compressing the training of billion-parameter uMLIPs from weeks to hours. This work establishes a new high-water mark for AI-for-Science (AI4S) foundation models at Exascale and provides essential infrastructure for rapid scientific discovery.

Yuanchang Zhou, Hongyu Wang, Yiming Du +12

State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · Independent Researcher · Sun Yat-Sen University

Machine Learning Interatomic Potentials Communication-Efficient Distributed Training High-Performance Computing Computational Materials Science

Ask PDF

Aug 7, 2026

Scalable High-Fidelity Macromolecular Docking for GPU-Accelerated Supercomputers

Flexible macromolecular docking offers high-fidelity predictions of biomolecular interactions, but remains prohibitively expensive at scale. Among existing approaches, LightDock leverages Glowworm Swarm Optimization (GSO) for accuracy, yet suffers from limited parallelism, irregular computation, and severe load imbalance, preventing efficient execution on GPU supercomputers. We present SparkleDock, a scalable GSO-based docking framework enabling near-real-time flexible docking. We redesign GSO to expose massive fine-grained parallelism at the glowworm-agent level, and restructure the dominant energy scoring computation into a Tensor Core-compatible formulation, enabling efficient execution of irregular pairwise interactions through structured matrix operations. We further introduce a performance-model-driven scheduling for load balancing and out-of-core scaling across GPUs. SparkleDock achieves 9.7

\times

and 18.9

\times

speedups over LightDock on single A100 and H100 GPU, and delivers over two orders of magnitude acceleration at scale. On 512 GPUs, it reduces docking time from hours to seconds, enabling large-scale, high-fidelity virtual screening previously impractical with flexible docking.

Xiangyu Meng, Peng Chen, Mingzhen Li +7

College of Computer Science and Technology, Shandong Key Laboratory of Intelligent Oil & Gas Industrial Software, China University of Petroleum (East China), Qingdao, China · A*STAR Institute of Advanced Intelligence and Computing, Singapore · RIKEN Center for Computational Science, Kobe, Japan +2

GPU Kernel Optimization GPU Acceleration High-Performance Computing Metaheuristic Optimization Structure-Based Drug Design

Ask PDF

Abstract

Molecular Dynamics Simulation

Explore similar work

Implicit Machine Learning Force Fields Accelerate Molecular Dynamics Simulations

Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials

Scalable High-Fidelity Macromolecular Docking for GPU-Accelerated Supercomputers