arXiv:2604.07276cs.DCcs.AI2026-04

让分子动力学模拟支持多卡高效运行的深度势能模型。

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

  • 将深度势能模型嵌入GROMACS,实现跨多节点并行推理。
  • 32卡下弱扩展效率达40%-80%,性能瓶颈为原子通信开销。
  • 适合追求高精度分子模拟的科研人员使用。

GROMACS是经典分子动力学(MD)的行业标准。随着基于AI的原子间势能模型在保持MD速度的同时逼近量子精度,如何在多GPU系统中高效集成神经网络推理成为挑战。本文将机器学习势能框架DeePMD-kit集成至GROMACS,实现基于域分解的多节点GPU加速推理。通过扩展NNPot接口并引入独立于主模拟的域分解层,所有进程并发执行推理,并每步使用两次MPI集体通信广播坐标与聚合重分配力。训练了一个含160万参数的自研DPA-1模型,基于溶剂化蛋白片段数据集。在包含15,668个原子的小蛋白系统上验证后,对多种硬件(NVIDIA A100和AMD MI250x)进行基准测试,最多支持32个设备。强扩展效率在16个设备时达66%,32个设备时为40%;弱扩展效率在16个设备时为80%,32个设备时为48%(MI250x)和40%(A100)。ROCm系统分析显示,超过90%的耗时集中在DeePMD推理,而MPI通信贡献不足10%,主要因全局同步点所致。瓶颈来自由截断半径决定的不可避鬼原子开销,且各计算进程负载不均。结果表明,在大规模上实现接近从头算精度的生产级分子动力学模拟已成为可能。

原文摘要 · Abstract (English)

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.

分子动力学深度势能多卡加速高性能计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。