Scaling Distributed Machine Learning Interatomic Potential Inference with DistMLIP2

Atomistic simulation has been the primary workhorse in computational materials and drug discovery over the past 60 years. However, the quantum mechanical calculations pivotal for the effective simulation of atomistic systems, such as density functional theory (DFT), remain too slow for practical use in large scale systems. In recent years, machine learning interatomic potentials (MLIPs), machine learning models trained on DFT energy and force labels, are capable of reaching DFT-levels of accuracies while performing orders of magnitude faster inference than traditional DFT workflows. Distributing MLIP inference across many GPUs poses as a powerful method for scaling simulation sizes and time scales even further. In this work, we present DistMLIP2, where we introduce four elements enabling a new state of the art in distributed MLIP inference: runtime skin recomputation, partial graph reconstruction, hierarchical partitioning and partition cycling, and graph computation pruning. We implement four widely used MLIPs into DistMLIP2: UMA, MACE, Orbv3, and NequIP, and show DistMLIP2 pareto-dominates the previous state of the art distributed inference methods such as ALCHEMI, LAMMPS-KOKKOS, and DistMLIP when controlling for MLIP size, architecture, kernels, and inference settings. Most notably, DistMLIP2 achieves 2.4x faster inference time than ALCHEMI while simulating 2.16x more atoms compared to LAMMPS. With partition cycling enabled, we show that DistMLIP2 can simulate systems containing 4-5x the atom count of its previous limit.