High-performance training and inference for deep equivariant interatomic potentials
Chuin Wei Tan, Marc Descoteaux, Mit Kotak, Gabriel de Miranda Nascimento, Seán R. Kavanagh, Laura Zichi, Menghang Wang, Aadit Saluja, Yizhong R. Hu, Tess Smidt, Anders Johansson, William C. Witt, Boris Kozinsky, Albert Musaelian
Abstract
The NequIP framework is redesigned for scalable distributed training and PyTorch 2.0 compilation. AOT Inductor inference and optimized Allegro kernels accelerate molecular dynamics by factors of 5–18 on practical system sizes.
Topics & Concepts
InferenceTraining (meteorology)Computer scienceScalabilityArtificial intelligenceAlgorithmTraining setEquivariant mapArtificial neural networkDynamics (music)Theoretical computer scienceStatistical physicsKey (lock)MathematicsApproximate inferenceMachine learningPhysicsStatistical inferenceApplied mathematicsMolecular dynamicsDeep neural networksDeep learningMachine Learning in Materials ScienceQuantum many-body systemsModel Reduction and Neural Networks