# chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations

arXiv: https://arxiv.org/abs/2506.04055 · HTML: https://mlsanity.com/paper/2506.04055/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 78 references that are in the map.
- First posted on arXiv 4 Jun 2025.

## Builds on (most cited first, 30 of 30)

- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules](https://mlsanity.com/paper/2011.14115/) (2020)
- [Deep Potential Molecular Dynamics: a scalable model with the accuracy of quantum mechanics](https://mlsanity.com/paper/1707.09571/) (2017)
- [Machine Learning Force Fields](https://mlsanity.com/paper/2010.07067/) (2020)
- [Machine learning for molecular simulation](https://mlsanity.com/paper/1911.02792/) (2019)
- [The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials](https://mlsanity.com/paper/2205.06643/) (2022)
- [Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations](https://mlsanity.com/paper/2210.07237/) (2022)
- [Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction](https://mlsanity.com/paper/2502.12147/) (2025)
- [SPICE, A Dataset of Drug-like Molecules and Peptides for Training Machine Learning Potentials](https://mlsanity.com/paper/2209.10702/) (2022)
- [DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics](https://mlsanity.com/paper/1712.03641/) (2017)
- [Machine Learning of coarse-grained Molecular Dynamics Force Fields](https://mlsanity.com/paper/1812.01736/) (2018)
- [TorchMD: A deep learning framework for molecular simulations](https://mlsanity.com/paper/2012.12106/) (2020)
- [Ewald-based Long-Range Message Passing for Molecular Graphs](https://mlsanity.com/paper/2303.04791/) (2023)
- [Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size](https://mlsanity.com/paper/2304.10061/) (2023)
- [Learning neural network potentials from experimental data via Differentiable Trajectory Reweighting](https://mlsanity.com/paper/2106.01138/) (2021)
- [Latent Ewald summation for machine learning of long-range interactions](https://mlsanity.com/paper/2408.15165/) (2024)
- [86 PFLOPS Deep Potential Molecular Dynamics simulation of 100 million atoms with ab initio accuracy](https://mlsanity.com/paper/2004.11658/) (2020)
- [chemtrain: Learning Deep Potential Models via Automatic Differentiation and Statistical Physics](https://mlsanity.com/paper/2408.15852/) (2024)
- [Euclidean Fast Attention -- Machine Learning Global Atomic Representations at Linear Cost](https://mlsanity.com/paper/2412.08541/) (2024)
- [Automated discovery of a robust interatomic potential for aluminum](https://mlsanity.com/paper/2003.04934/) (2020)
- [Uncertainty Estimation for Molecules: Desiderata and Methods](https://mlsanity.com/paper/2306.14916/) (2023)
- [Predicting solvation free energies with an implicit solvent machine learning potential](https://mlsanity.com/paper/2406.00183/) (2024)
- [Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators](https://mlsanity.com/paper/2402.13984/) (2024)
- [Learning Non-Local Molecular Interactions via Equivariant Local Representations and Charge Equilibration](https://mlsanity.com/paper/2501.19179/) (2025)
- [DistMLIP: A Distributed Inference Platform for Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2506.02023/) (2025)

## Built on by (most cited first, 6 of 6)

- [Enhanced Sampling for Efficient Learning of Coarse-Grained Machine Learning Potentials](https://mlsanity.com/paper/2510.11148/) (2025)
- [Refining Machine Learning Potentials through Thermodynamic Theory of Phase Transitions](https://mlsanity.com/paper/2512.03974/) (2025)
- [Facet: highly efficient E(3)-equivariant networks for interatomic potentials](https://mlsanity.com/paper/2509.08418/) (2025)
- [Coarse-Grained Boltzmann Generators](https://mlsanity.com/paper/2602.10637/) (2026)
- [Making Room for AI: Multi-GPU Molecular Dynamics with Deep Potentials in GROMACS](https://mlsanity.com/paper/2604.07276/) (2026)
- [Transferable Implicit Solvent Machine Learning Potential for Drugs and Proteins Approaching Ab Initio Accuracy](https://mlsanity.com/paper/2607.10887/) (2026)

## Abstract (arXiv)

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.