Isothermal-Isobaric Molecular Dynamics with a Foundation Interatomic Potential in LAMMPS across NVIDIA and Intel GPUs

Abstract

Foundation machine-learning interatomic potentials (MLIPs) offer one model for many chemistries, but their deep, globally coupled architectures do not fit the spatial decomposition that scales classical and short-ranged MLIPs. We present pair_style uma, a LAMMPS pair style that evaluates the Universal Model for Atoms (UMA) through LibTorch on NVIDIA and Intel GPUs from a single source tree. It reimplements FairChem's edge-sharded graph parallelism in C++ over exported TorchScript modules, with one MPI rank per device, NCCL or oneCCL collectives, and the whole energy and force path in double precision. The strain-gradient virial is carried through the exported modules and the cross-device reductions, so that fix npt runs on both vendors' hardware. We measure NPT capacity and cost on one 4xA100-40GB node, one 8xH200 node, and one 12-tile Intel node, and on up to 32 Intel nodes. Without activation checkpointing, one node holds NaCl cells of about 8,000 atoms (A100, 0.21 ns/day), 22,000 atoms (Intel, 0.10 ns/day), and 47,000 atoms (H200, 0.16 ns/day). Graph parallelism raises the capacity of a single device by 8.0x on H200 and 4.6x on A100. Across nodes, capacity grows to about 779,000 atoms on 32 Intel nodes, with a slower speed at 0.007 ns/day. On H200 and A100, a FLOP count of FairChem's graph-parallel predictor gives lower bounds of 31% and 51% of FP64 peak. From one omat checkpoint and with no refit, nine 100 ps NPT heating ramps (NaCl, equiatomic Pt-Ru, and equiatomic (HfNbTaTiZr)N, on all three node types) ran to completion.