支持万原子级分子模拟的通用机器学习势能框架
chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations
- 通用框架，兼容任意JAX定义的局部势能模型
- 在多张GPU上实现万原子级模拟，性能业界领先
- 适配主流分子动力学软件LAMMPS，适合高性能计算用户
机器学习势能（MLPs）发展迅速，有望彻底改变分子动力学（MD）模拟。但现有工具大多局限于特定架构，缺乏与标准MD软件集成，或无法跨GPU并行。为此，我们提出chemtrain-deploy，一个可在LAMMPS中部署任意JAX定义的半局域势能模型的通用框架。该框架支持多GPU并行，可实现百万原子级大规模模拟，具备业界领先效率。我们通过图神经网络模型（如MACE、Allegro、PaiNN）在液-气界面、晶体材料和溶剂化肽等系统上验证了其性能与可扩展性。结果表明，chemtrain-deploy在真实高精度模拟中具有实用价值，并为模型架构选择与未来设计提供指导。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。