chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations
Fuente:
arXiv
Saved in:
| Main Authors: | , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866909636987191296 |
|---|---|
| author | Fuchs, Paul Chen, Weilong Thaler, Stephan Zavadlav, Julija |
| author_facet | Fuchs, Paul Chen, Weilong Thaler, Stephan Zavadlav, Julija |
| contents | Machine learning potentials (MLPs) have advanced rapidly and show great promise to transform molecular dynamics (MD) simulations. However, most existing software tools are tied to specific MLP architectures, lack integration with standard MD packages, or are not parallelizable across GPUs. To address these challenges, we present chemtrain-deploy, a framework that enables model-agnostic deployment of MLPs in LAMMPS. chemtrain-deploy supports any JAX-defined semi-local potential, allowing users to exploit the functionality of LAMMPS and perform large-scale MLP-based MD simulations on multiple GPUs. It achieves state-of-the-art efficiency and scales to systems containing millions of atoms. We validate its performance and scalability using graph neural network architectures, including MACE, Allegro, and PaiNN, applied to a variety of systems, such as liquid-vapor interfaces, crystalline materials, and solvated peptides. Our results highlight the practical utility of chemtrain-deploy for real-world, high-performance simulations and provide guidance for MLP architecture selection and future design. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_04055 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations Fuchs, Paul Chen, Weilong Thaler, Stephan Zavadlav, Julija Computational Physics Machine Learning Chemical Physics Machine learning potentials (MLPs) have advanced rapidly and show great promise to transform molecular dynamics (MD) simulations. However, most existing software tools are tied to specific MLP architectures, lack integration with standard MD packages, or are not parallelizable across GPUs. To address these challenges, we present chemtrain-deploy, a framework that enables model-agnostic deployment of MLPs in LAMMPS. chemtrain-deploy supports any JAX-defined semi-local potential, allowing users to exploit the functionality of LAMMPS and perform large-scale MLP-based MD simulations on multiple GPUs. It achieves state-of-the-art efficiency and scales to systems containing millions of atoms. We validate its performance and scalability using graph neural network architectures, including MACE, Allegro, and PaiNN, applied to a variety of systems, such as liquid-vapor interfaces, crystalline materials, and solvated peptides. Our results highlight the practical utility of chemtrain-deploy for real-world, high-performance simulations and provide guidance for MLP architecture selection and future design. |
| title | chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations |
| topic | Computational Physics Machine Learning Chemical Physics |
| url | https://arxiv.org/abs/2506.04055 |