Year: 2025
Type: preprint
Abstract: Machine learning interatomic potentials, particularly those based on deep equivariant neural networks, have demonstrated state-of-the-art accuracy and computational efficiency in atomistic modeling tasks like molecular dynamics and high-throughput screening. The size of datasets and demands of downstream workflows are growing rapidly, making robust and scalable software essential. This work presents a major overhaul of the NequIP framework focusing on multi-node parallelism, computational performance, and extensibility. The redesigned framework supports distributed training on large datasets and removes barriers preventing full utilization of the PyTorch 2.0 compiler at train time. We demonstrate this acceleration in a case study by training Allegro models on the SPICE 2 dataset of organic molecular systems. For inference, we introduce the first end-to-end infrastructure that uses the PyTorch Ahead-of-Time Inductor compiler for machine learning interatomic potentials. Additionally, we implement a custom kernel for the Allegro model's most expensive operation, the tensor product. Together, these advancements speed up molecular dynamics calculations on system sizes of practical relevance by up to a factor of 18.
Source: arXiv (Cornell University)
Authors: Tan, Chuin Wei, Descoteaux, Marc L., Kotak, Mit, Nascimento, Gabriel de Miranda, Kavanagh, Seán R., Zichi, Laura, Wang, Menghang, Saluja, Aadit, Hu, Yizhong R., Smidt, Tess, Johansson, Anders, Witt, William C., Kozinsky, Boris, Musaelian, Albert
Language: English
Subfield: Materials Chemistry
Field: Materials Science
Domain: Physical Sciences
Open Access status: green
Funders: National Science Foundation, U.S. Department of Energy, Harvard University, Robert Bosch, National Energy Research Scientific Computing Center, Office of Science, National Nuclear Security Administration, Center for the Environment, Harvard University, Materials Research Science and Engineering Center, Harvard University, Office of Naval Research, Division of Materials Research, Office of Advanced Cyberinfrastructure, Basic Energy Sciences, FAS Division of Science, Harvard University, Sandia National Laboratories
Awards: G1019100964, G1071216263, G11630931, G121915436, G122673668, G1261772666, G1286236842, G1286251239, G1591330662, G1645119126, G1677143136, G1863003419, G203625172, HDR Institute: Institute for Data Driven Dynamical Design, G2345624000, G2366492356, G2411683797, G2417641742, G2618837358, G2692471796, G280414146, G288067973, G2900453181, G2988778103, G30714459, G3083819904, Collaborative Research: DMREF: Design of Superionic Conductors by Tuning Lattice Dynamics, G3274770636, AI Institute: AI Research Institute for Fundamental Interactions, G3945850812, G3967535168, G413247060, G419881185, G42916606, G4501827968, G4556251948, G458500238, G4796240486, G491876556, G498139845, G5076365615, G5261056217, G5328548810, G5565249710, G5892109311, G5964803810, G6110219478, G6338696783, G650090426, G660390541, MACHINE LEARNING FOR ACCELERATED UNDERSTANDING OF DYNAMIC CATALYSIS, G6652958823, Graduate Research Fellowship Program (GRFP), G7651947074, G8006577219, G807938293, G816254390, G8269158468, G855168022, Materials Research Science and Engineering Center, G8876996369, G8906985441, G969889393
OpenAlex ID: W4414644745