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

arXiv:2506.04055 · physics.comp-ph, cs.LG, physics.chem-ph · Submitted 2025-06-04 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations".

Jane: The gist: chemtrain-deploy presents a modelagnostic deployment framework that enables modelagnostic deployment of MLPs in LAMMPS, which supports any JAX-defined semi-local potential,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at a paper called "chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations." This sounds like it’s tackling some big problems with how we use these machine learning potentials in molecular dynamics.

Jane: It really does sound ambitious, Tom. The authors are trying to solve the issue where most existing software is locked into specific model architectures or just can't run efficiently on multiple GPUs for huge systems.

Lu: Exactly. They’re proposing a framework that lets you take a trained machine learning potential, say one defined in JAX, and deploy it seamlessly into LAMMPS. That’s the core idea—model agnostic deployment in LAMMPS, according to the abstract <ref:2506.04055#pg1>.

Meng: So if I'm an engineer looking at this, what does that actually mean for my workflow? Can I just plug a model into my existing simulation setup without rewriting everything?

Tom: That’s the main selling point. They build this framework in three parts: exporting the trained model, importing it into LAMMPS via a plugin, and then efficiently evaluating that model on high-performance hardware <ref:2506.04055#pg1>.

Jane: And they handle the evaluation side using XLA and PJRT to make sure it runs fast on specific hardware without needing you to rewrite LAMMPS or anything else <ref:2506.04055#pg1>.

Lu: They specifically mention supporting any JAX-defined semi-local potential, which is a big deal because it’s not limited to just one type of machine learning model <ref:2506.04055#pg1>.

Meng: The summary talks about benchmarking three specific models: MACE, Allegro, and PaiNN. That tells us they aren't just talking theoretically; they’re testing them on real-world problems like water-vapor coexistence and solid state aluminum <ref:2506.04055#pg3>.

Tom: Right, but what’s the actual problem they are solving? The paper frames it as bridging the gap between the flexibility of JAX training and the robustness of established MD packages like LAMMPS <ref:2506.04055#pg3>.

Jane: They show how this framework allows for large-scale simulations with systems containing millions of atoms across multiple GPUs, which is where many current tools fall short <ref:2506.04055#pg1>.

Lu: The paper details how they compute potential energies in parallel by partitioning the full graph into one subgraph per processor using spatial system decomposition and neighbor list generations from the MD software itself <ref:2506.04055#pg2>.

Title and authors: Meng: That sounds like a lot of work just to get the energies right across all those processors without communication bottlenecks <ref:2506.04055#pg2>.

Tom: They go into the details about how they handle forces too. They show that instead of needing to include the 2T-th order neighbors for every atom, they can compute the forces by summing up partial forces from corresponding copies on other processors <ref:2506.04055#pg2>.

Jane: That’s clever because it avoids those intermediate communication operations when calculating all the particles' forces and potential energies <ref:2506.04055#pg2>.

Lu: They also have these runtime optimizations using XLA, but they address a problem where XLA re-optimizes every time the input shape changes, which takes too long <ref:2506.04055#pg1>.

Meng: So they use buffering to keep the input shape fixed and only recompile when necessary, which sounds like it should keep things snappy during a simulation run <ref:2506.04055#pg1>.

Tom: They even have a specific strategy for collective recompilations across multiple devices by synchronizing errors if one device has an overflow, which shortens those warm-up periods <ref:2506.04055#pg1>.

Jane: It sounds like they’ve put a lot of work into making the evaluation process as efficient as possible on the hardware it’s running on <ref:2506.04055#pg1>.

Lu: They also benchmarked MACE, Allegro, and PaiNN, showing they are all trained to comparable accuracy across different systems like aluminum and water <ref:2506.04055#pg3>.

Meng: The results suggest that the Allegro model is very efficient for systems with just a few atom types like water or aluminum because it scales well <ref:2506.04055#pg3>.

Tom: But for systems with more chemically diverse atoms, like Chignolin, they found that the MACE model provides a better balance between accuracy and scalability <ref:2506.04055#pg3>.

Jane: And they conclude that the Allegro model is highly efficient for systems with a few atom types because of those scaling results <ref:2506.04055#pg3>.

Lu: The authors also point out that while some models like MACE have better throughput for Chignolin, they still show good strong scaling across all tested systems <ref:2506.04055#pg3>.

Title and authors: Tom: So, to sum up the main points of "chemtrain-deploy," it’s a framework that lets you take any trained semi-local MLP and deploy it into LAMMPS for large-scale MD simulations across multiple GPUs <ref:2506.04055#pg1>.

Jane: It moves ML potentials out of the niche and into mainstream, making them accessible to bigger systems and more complex research areas like computational biology <ref:2506.04055#pg1>.

Meng: From an engineering standpoint, this means we can test much larger physical systems than before using these models for force calculations <ref:2506.04055#pg1>.

Lu: The real implication is the ability to benchmark different machine learning potentials consistently across diverse simulation settings, which helps guide future development <ref:2506.04055#pg3>.

Tom: We've covered a lot of ground on how this framework works and what the specific models like Allegro and MACE are doing. It really shows how they’re trying to make these powerful ML tools practical for industrial scale simulations <ref:2506.04055#pg3>.

Jane: It’s a significant step toward using these potentials not just for small, highly controlled tests but for those million-atom systems that are really needed in material science and biology <ref:2506.04055#pg1>.

Lu: The future work they hint at is about pushing the boundaries of what these semi-local models can handle when applied to even more complex, real-world chemical environments <ref:2506.04055#pg3>.

Meng: It’s interesting how they managed to keep the accuracy within chemical accuracy while achieving performance comparable to classical force fields in some cases <ref:2506.04055#pg2>.

Tom: So, we're wrapping up our discussion on "chemtrain-deploy," this parallel and scalable framework for machine learning potentials in million-atom MD simulations <ref:2506.04055#pg1>. It’s a solid tool for anyone working with large molecular systems.

Jane: It really opens up exciting opportunities because it makes these ML potentials available to a much wider range of simulation protocols and hardware setups <ref:2506.04055#pg1>.

Lu: The potential here is massive, allowing us to explore biological systems and complex materials at scales that were previously out of reach for this type of computation <ref:2506.04055#pg3>.

Meng: It’s a lot of work, but seeing the scaling results they showed on MACE and Allegro is pretty impressive when you think about the hardware requirements <ref:2506.04055#pg3>.

Tom: We'll stick around for the next paper we look at, but for now, that's our time on "chemtrain-deploy."

The paper's summary: Tom: So, to quickly recap, this paper introduces chemtrain-deploy as a way to take trained machine learning potentials and deploy them into LAMMPS for really big molecular dynamics simulations using multiple GPUs.

Jane: It basically takes models that were trained in one place, like JAX, and makes them work in the actual MD software on high-performance hardware without you having to rewrite all your code from scratch.

Tom: That’s right. The paper lays out this three-part process for exporting the model, importing it into LAMMPS via a plugin or modification, and then running the evaluation efficiently on powerful hardware.

Jane: And what they really hammer home is how they handle the math behind it—they use spatial system decomposition to partition the simulation graph across different processors, making sure each atom's energy gets calculated only once.

Tom: It’s not just about dividing up the work; it’s about keeping everything synchronized so you can compute those forces correctly without needing constant communication between all the GPUs during every single time step.

Jane: They also get into how they manage those dynamic inputs, using XLA to buffer things to a fixed shape, which means they avoid costly recompilations unless absolutely necessary, making the runtime much smoother.

Tom: And the results on MACE and Allegro are interesting because they show that you can use these models for systems with millions of atoms across multiple GPUs while still keeping the accuracy within chemical accuracy.

Jane: That’s a big deal because it opens up new research areas in computational biology and material sciences where simulating things at that scale is currently just not possible with this kind of computation.

Tom: And they give us some guidance on which model to pick based on what you're simulating, like suggesting Allegro for simpler systems such as water and aluminum, or MACE for more complex chemical mixtures.

Jane: So the real implication here is that machine learning potentials are becoming a much more practical tool for large-scale molecular dynamics than we thought they would be.

Tom: It’s a solid foundation because it lets researchers benchmark different AI models consistently across different simulation setups, which helps guide how we develop these potentials in the future.

Jane: Next, we're going to look at some of the specific performance numbers they show on those big-atom systems and what those scaling results actually mean for your own research pipeline.

The paper's improvements: Tom: So, we’re talking about how they suggested improving the chemtrain-deploy framework beyond just getting it to run on more GPUs.

Jane: They focused on making sure the runtime optimization is really smart, especially when you’re dealing with systems that keep changing their size.

Tom: Exactly. They addressed the issue where XLA usually re-optimizes everything every single time the shape of your input graph changes, which can take too long during a simulation run.

Jane: They introduced a buffering strategy to hold those dynamically shaped inputs to a fixed size and only trigger a full recompilation when it’s really needed.

Tom: That means you can keep running simulations faster because you aren't waiting around for the hardware to re-optimize the program constantly <ref:2506.04055#pg1>.

Jane: They also have this specific trick for handling simultaneous recompilations across multiple devices by making one device report an error and having that error trigger a collective recompilation on everyone at once.

Tom: That shortens up the warm-up time significantly, which is super important when you're running massive parallel applications <ref:2506.04055#pg1>.

Jane: They’re essentially fine-tuning the software so that it handles the complexities of running millions of atoms in parallel with less overhead than before.

Tom: It moves the framework from just being functional to being genuinely performant for large-scale science <ref:2506.04055#pg1>.

Jane: And they’re also pushing the boundaries on what these semi-local models can actually do, suggesting that this deployment method lets us explore new simulation protocols, like those involving external fields or enhanced sampling techniques.

Tom: That means researchers aren't stuck just running standard molecular dynamics anymore; they can test entirely new ways to study materials or biology <ref:2506.04055#pg1>.

Jane: It’s about making these potentials versatile enough to handle a wider variety of scientific problems than we currently see them being used for <ref:2506.04055#pg1>.

Tom: So, they aren't just giving us a tool; they're showing how this deployment method opens up entirely new avenues for MD research that were previously blocked by software limitations.

Jane: The real implication here is that the barrier to entry for using these advanced machine learning potentials in big simulations is dropping quite a bit <ref:2506.04055#pg1>.

Tom: It’s about making high-performance computation accessible to more diverse scientific fields, which is a big step toward computational chemistry and materials science <ref:2506.04055#pg3>.

Jane: And they want everyone to see how this framework can help accelerate the adoption and development of these machine learning potentials for transformative advances in large-scale molecular simulations <ref:2506.04055#pg3>.

Tom: So, the next thing we’re going to look at is how these specific model choices, like MACE versus Allegro, actually perform when you're dealing with systems that have a high degree of chemical diversity.

Conclusion: Tom: So, to wrap things up, we’re talking about "chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations."

Jane: It’s a tool that takes trained AI models and deploys them into LAMMPS so they can run big molecular dynamics simulations across multiple GPUs efficiently.

Tom: It really solves the problem of making these ML potentials practical for simulating millions of atoms, which is a huge step forward in materials science and biology <ref:2506.04055#pg1>.

Jane: The framework shows how to compute forces and energies in parallel by splitting the graph up, ensuring every atom’s energy is calculated only once <ref:2506.04055#pg2>.

Tom: And they’ve given us guidance on choosing between models like Allegro for simple systems and MACE for chemically diverse ones based on the scaling results <ref:2506.04055#pg3>.

Jane: It’s about making sure that when you’re running these complex simulations, you aren't fighting with slow runtime optimizations or constant re-compilations <ref:2506.04055#pg1>.

Tom: And they also show how the framework can be extended to support non-standard simulation protocols, which opens up a lot of new ways to use these potentials <ref:2506.04055#pg1>.

Jane: It gives researchers more flexibility, moving us beyond just running standard MD simulations <ref:2506.04055#pg1>.

Tom: The implication is that machine learning potentials are becoming a much more accessible and powerful tool for large-scale scientific discovery across different fields <ref:2506.04055#pg3>.

Jane: It’s about enabling simulations of complex materials and biological systems at scales that were just out of reach before <ref:2506.04055#pg1>.

Tom: Lu, what’s your take on the potential here for AI in science?

Lu: I think this framework really shows how we can bridge the gap between cutting-edge machine learning and established physics simulation software, which is a huge cultural shift <ref:2506.04055#pg1>.

Meng: From an engineering standpoint, it means we can finally test much larger physical systems without getting bogged down by the computational limitations of the current tools <ref:2506.04055#pg1>.

Lalam: This work shows that AI isn't just about generating text or images; it’s becoming a foundational layer for simulating complex, real-world physical systems at unprecedented scales, which will change how we approach modeling reality <ref:2506.04055#pg1>.

Tom: Jane, what’s the final word on this paper?

Jane: It’s a solid piece of work because it provides a concrete way to deploy ML potentials reliably for massive MD simulations across hardware setups <ref:2506.04055#pg1>.

Lu: The future work they hint at is about pushing the boundaries of what these semi-local models can handle when applied to even more complex, real-world chemical environments <ref:2506.04055#pg3>.

Meng: It’s good to see this kind of progress, because it sets a new standard for how we deploy AI models in scientific computing <ref:2506.04055#pg1>.

Lalam: I see this as a step toward a more intuitive way for scientists to model complex molecular interactions, which will improve how we communicate and collaborate across different disciplines <ref:2506.04055#pg1>.

Technical University of Munich · Atomistic Modeling Center (AMC), Munich Data Science Institute (MDSI) · Valence Labs

physics.comp-ph, cs.LG, physics.chem-ph

Submitted: 2025-06-04

Updated: 2025-06-04

Comments: Source code available at: https://github.com/tummfm/chemtrain

DOI: 10.1021/acs.jctc.5c00996

Code: https://github.com/atomistic-ml/ani-al

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 87/100

The gist: The gist: chemtrain-deploy presents a modelagnostic deployment framework that enables modelagnostic deployment of MLPs in LAMMPS, which supports any JAX-defined semi-local potential, allowing users

Key concepts

Modelagnostic Deployment Framework
This framework is designed to allow any trained machine learning potential model to be used within established molecular dynamics software like LAMMPS without needing to rewrite the core simulation code. It acts as a bridge, handling the necessary steps from saving a trained model file to running it efficiently on specialized hardware.
Distributed Potential Computation
This technique parallelizes the calculation of potential energies across multiple processors (GPUs/CPUs). The system partitions the large molecular system into smaller domains. Each processor calculates the energy for its local atoms and handles necessary boundary interactions with neighboring domains, ensuring all atoms' energies are computed exactly once.
Runtime Optimizations and Buffering
To speed up model evaluation on XLA hardware, chemtrain-deploy buffers inputs to a fixed shape instead of re-optimizing the program constantly. It manages buffer resizing and intelligently coordinates recompilations across multiple devices to minimize startup time during parallel computations.

Terminology

Summary

The gist: chemtrain-deploy presents a modelagnostic deployment framework that enables modelagnostic deployment of MLPs in LAMMPS, which 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>

Framework Overview

chemtrain-deploy is a framework that complements the chemtrain framework [51] to apply trained models in large-scale MD simulations using established MD software> Therefore, chemtrain-deploy comprises three parts: exporting a trained model, importing it into established MD software through a plugin or modification, and efficiently evaluating the model on high-performance hardware> First, chemtrain-deploy extends the framework chemtrain (Fig. 1a) to export trained models to a self-contained format (Fig. 1b)> The export saves the model architecture and parameters through the MLIR framework [53] and properties that define the input and output of the model, such as length and energy units, the maximum length of graph edges, and the format of the input graph> Secondly, chemtrain-deploy consists of a plugin to load and use the exported model for large-scale molecular dynamics simulations in established MD software (Fig. 1c)> The MD software provides basic and advanced algorithms to perform MD simulations> Moreover, the MD software provides algorithms to decompose the system into multiple domains for parallelization and create a neighbor graph representation of the system (Fig. 1d)> Therefore, the plugin interfaces chemtrain-deploy with framework-specific algorithms for parallelization and to run advanced MD simulations with the exported model> Finally, chemtrain-deploy provides a library for the plugin to evaluate the exported model (Fig. 1e)> This library uses XLA [54] and PJRT [55] to translate the model into an efficient backend-specific computation at runtime> Therefore, the library extends the LAMMPS’ capabilities to run efficiently on specific hardware and evaluate new force-field architectures without rewriting or recompiling LAMMPS>

Distributed Potential Computation

chemtrain-deploy computes the potential energies in parallel on multiple independent processors (GPUs/CPUs) by partitioning the full graph into one subgraph per processor using spatial system decomposition and neighbor list generations provided in MD software, e.g., LAMMPS [29]> As outlined in Figure 1d, LAMMPS distributes the workload by dividing the system into non-overlapping domains, such that each atom is local to exactly one domain> LAMMPS then assigns the domains and the contained local atoms to the available processors> For each processor, LAMMPS additionally copies all atoms from other domains within a distance of Rcomm of the processor’s domain boundary and constructs a neighbor list graph> Since the maximum distance between an atom i and any of its T-th neighbors can be T R, choosing Rcomm ≥ T R ensures that the neighbor graph of each domain contains all T-th neighbor atoms of all local atoms> Thus, summing up the predicted energies of local atoms results in the total potential energy of the system, as each atom energy Ui is computed exactly once on a subgraph with a complete environment of i>

Running MD simulations requires chemtrain-deploy to compute the forces acting on the atoms> From the sum rule, the total force on an atom fi = −∂Ui(ri, Ri) ∂ri − X j∈N R ≤T ∂Uj (rj, Rj) ∂ri (2)> This decomposes into a sum of partial forces fij = −∂Uj (rj,Rj) ∂ri> The total forces of all local atoms can be computed directly on each processor by computing all nonzero partial forces> However, this approach generally requires extending the domain subgraph to include the 2T-th order neighbors, which are necessary to correctly compute the potential energy of all T-th order neighbors of the local atoms> Alternatively, chemtrain-deploy computes the partial forces of all local atoms with respect to all atoms of the domain subgraph> Then, the total forces on each local atom can be obtained by summing up all partial forces from corresponding copies on other processors> Thus, by constructing graphs containing all T-th order neighbors of local atoms, all particles’ forces and potential energies can be computed with initial and final but without intermediate communication operations>

Runtime Optimizations and Buffering

chemtrain-deploy optimizes and evaluates the model through XLA> However, XLA re-optimizes the program every time the shape of an input changes, which typically requires more time than the actual computation> Thus, chemtrain-deploy buffers all dynamically shaped inputs to a fixed shape and evaluates the model as outlined in Figure 1> First, chemtrain-deploy computes the required buffer shapes> These shapes can vary, for example, if the number of atoms and neighbors in the domain changes> If the buffer capacities are exceeded, chemtrain-deploy enlarges the buffers and recompiles the model using the TensorFlow [56] call module loader> If no buffer overflowed, chemtrain-deploy transforms and copies data to the device and performs the computation> The models use internal buffers to enable optimizations such as graph pruning in the computation> Thus, after the computation, chemtrain-deploy checks whether internal model buffers overflowed> If internal model buffers overflow, chemtrain-deploy repeats the computation with resized model buffers> If no buffer overflowed, chemtrain-deploy copies back the computed forces and returns statistics of the computation> Since chemtrain-deploy only enlarges buffers, the frequency of recompilations decreases for systems at equilibrium> However, recompilations can happen frequently in the initial stages of computations on multiple devices if each device recompiles independently> Thus, chemtrain-deploy enforces collective recompilations of multiple devices per time step by explicitly controlling recompilations> Therefore, the chemtrain-deploy plugin first tries to evaluate the model with recompilation disabled> If a recompilation is necessary on one device, the device raises an exception that will be called by the plugin> The plugin then synchronizes the error to all devices, which will enlarge overflown and nearly filled buffers and recompile the program> Thereby, chemtrain-deploy boosts simultaneous recompilations on multiple devices to shorten warm-up periods in large-scale parallel applications>

Model Architectures and Training

chemtrain-deploy supports three widely used GNN models: MACE, Allegro, and PaiNN, all trained to comparable accuracy> We test them on diverse systems including water–vapor coexistence, solid state (fcc) aluminum, and solvated Chignolin> Our results demonstrate the practicality of chemtrain-deploy for real-world, high-performance molecular simulation and provide guidance on GNN selection and future development> In all cases, the models reached comparable MAE or RMSE values, remained within chemical accuracy, and matched reported literature benchmarks (Table 1)> The Allegro model performed best for aluminium and water> The MACE models achieved a similar throughput for all systems, outperforming the speed of Allegro for the Chignolin system> Based on these scaling results, we conclude that the Allegro model is highly efficient for systems with a few atom types, such as water and aluminum> However, for chemically diverse systems, MACE provides a better tradeoff between accuracy, scalability, and robustness>

Scaling to Million-Atom Systems

The Allegro and MACE models showed good strong scaling for all systems> The PaiNN model performed similarly to the Allegro model in the aluminium system> Based on these scaling results, we conclude that the Allegro model is highly efficient for systems with a few atom types, such as water and aluminum> However, for chemically diverse systems, MACE provides a better tradeoff between accuracy, scalability, and robustness>

Conclusion

chemtrain-deploy overcomes key limitations of existing software which are often restricted to specific model architectures, provide limited training support, or face scalability challenges> Our results demonstrate excellent scaling of modern GNN potentials through chemtrain-deploy for different systems, particularly in simulations involving millions of atoms across multiple GPUs> This capability enables the application of semi-local MLPs to new fields in computational biology and material sciences> The broad applicability of chemtrain-deploy opens several exciting opportunities for MD community> Its ability to efficiently handle million-atom systems enables simulations of complex materials and biological systems> By making use of the various functions of LAMMPS, chemtrain-deploy supports non-standard simulation protocols such as those under external fields or using enhanced sampling techniques> Importantly, its model-agnostic design allows easy benchmarking and comparison of different machine learning potentials within consistent simulation settings> Together, we expect that chemtrain-deploy will accelerate the adoption and development of advanced machine learning potentials, enabling transformative advances in large-scale molecular simulations and ultimately driving progress across computational chemistry and materials science>

Bares

chemtrain-deploy is a framework that complements the chemtrain framework to apply trained models in large-scale MD simulations using established MD software> Therefore, chemtrain-deploy comprises three parts: exporting a trained model, importing it into established MD software through a plugin or modification, and efficiently evaluating the model on high-performance hardware> First, chemtrain-deploy extends the framework chemtrain (Fig. 1a) to export trained models to a self-contained format (Fig.

Improvements for AI systems

  1. Improved deployment of MLPs in LAMMPS: The framework enables modelagnostic deployment of MLPs in LAMMPS, allowing users to exploit the functionality of LAMMPS and perform large-scale MLP-based MD simulations on multiple GPUs. This means existing, potentially architecture-specific, GNN models like MACE or Allegro can be seamlessly integrated into standard MD packages without requiring extensive rewriting.

  2. Enhanced scalability for million-atom simulations: By utilizing spatial system decomposition and neighbor list generations provided in MD software, chemtrain-deploy allows simulating systems containing millions of atoms across multiple GPUs. This addresses the limitation where existing JAX/MD tools are restricted to systems less than half a million atoms, enabling research into complex materials and biological systems at much larger scales.

  3. Optimized distributed potential computation: The method computes energies and forces in parallel by partitioning the graph, ensuring each atom energy Ui is computed exactly once on a subgraph with a complete environment of i. This approach handles the complexity of computing forces, which involves decomposing the total force into a sum of partial forces fij = −∂Uj (rj,Rj) ∂ri (2), without requiring intermediate communication operations for all local atoms.

  4. Reduced execution overhead through XLA optimizations: The framework uses XLA to perform hardware-independent optimizations such as common subexpression elimination and operation fusion to reduce computational cost and memory requirements. Furthermore, it manages dynamic inputs by buffering them to a fixed shape, minimizing costly recompilations unless necessary, thus boosting efficiency during runtime.

  5. Architectural guidance for GNN selection: The results provide critical insights into model choice: strictly local models generally exhibit superior scalability due to their limited communication overhead, while acknowledging that semi-local message-passing GNNs tend to provide improved accuracy but can exhibit reduced scalability and increased memory demands. This guides practitioners in selecting architectures based on system requirements (e.g., Allegro for water/aluminum vs. MACE for chemically diverse systems).

Abstract

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.

Sources

Related papers