chemtrain-deploy: A Parallel and Scalable Framework for Machine Learning Potentials in Million-Atom MD Simulations.

Area of Science:

• Computational Chemistry
• Materials Science
• Machine Learning

Background:

• Machine Learning Potentials (MLPs) are revolutionizing Molecular Dynamics (MD) simulations.
• Existing MLP software often lacks flexibility, integration with standard MD packages, and GPU parallelization.
• There is a need for a versatile framework to deploy diverse MLPs within established simulation tools.

Purpose of the Study:

• To introduce chemtrain-deploy, a novel framework for model-agnostic deployment of MLPs in LAMMPS.
• To enable large-scale, high-performance MD simulations using various JAX-defined semilocal potentials.
• To validate the efficiency and scalability of the framework across different ML architectures and systems.

Main Methods:

• Developed chemtrain-deploy, a framework integrating any JAX-defined semilocal potential with LAMMPS.
• Implemented GPU parallelization for efficient computation on multi-GPU systems.
• Validated performance and scalability using graph neural network architectures (MACE, Allegro, PaiNN) on diverse systems.

Main Results:

• chemtrain-deploy demonstrates state-of-the-art efficiency and scalability for systems up to millions of atoms.
• The framework successfully deployed and validated various MLPs (MACE, Allegro, PaiNN) in LAMMPS.
• Performance was confirmed across diverse simulations, including liquid-vapor interfaces, crystalline materials, and solvated peptides.

Conclusions:

• chemtrain-deploy provides a practical and efficient solution for deploying diverse MLPs in LAMMPS.
• The framework facilitates high-performance, large-scale MD simulations, advancing materials science and computational chemistry.
• Results offer guidance for selecting and designing future MLP architectures for MD simulations.