DistMLIP: A Distributed Inference Platform for Machine Learning Interatomic Potentials

Large-scale atomistic simulations are essential to bridge computational materialsand chemistry to realistic materials and drug discovery applications. In the past fewyears, rapid developments of machine learning interatomic potentials (MLIPs) haveoffered a solution to scale up quantum mechanical calculations. Parallelizing theseinteratomic potentials across multiple devices poses a challenging, but promisingapproach to further extending simulation scales to real-world applications. In thiswork, we present DistMLIP, an efficient distributed inference platform for MLIPsbased on zero-redundancy, graph-level parallelization. In contrast to conventionalspace-partitioning parallelization, DistMLIP enables efficient MLIP parallelizationthrough graph partitioning, allowing multi-device inference on flexible MLIPmodel architectures like multi-layer graph neural networks. DistMLIP presentsan easy-to-use, flexible, plug-in interface that enables distributed inference ofpre-existing MLIPs. We demonstrate DistMLIP on four widely used and state-of-the-art MLIPs: CHGNet, MACE, TensorNet, and eSEN. We show that existingfoundation potentials can perform near-million-atom calculations at the scale of afew seconds on 8 GPUs with DistMLIP.