PaperlayerView exact revision on arXiv

Source-linked AI summary

DeePMD-kit v3: A Multiple-Backend Framework for Machine Learning Potentials

Jinzhe Zeng, Duo Zhang, Anyang Peng, Xiangyu Zhang, Sensen He, Yan Wang, Xinzijian Liu, Hangrui Bi, Yifan Li, Chun Cai, Chengqian Zhang, Yiming Du, Jia-Xin Zhu, Pinghui Mo, Zhengtao Huang, Qiyu Zeng, Shaochen Shi, Xuejian Qin, Zhaoxi Yu, Chenxing Luo, Ye Ding, Yun-Pei Liu, Ruosong Shi, Zhenyu Wang, Sigbjørn Løland Bore, Junhan Chang, Zhe Deng, Zhaohan Ding, Siyuan Han, Wanrun Jiang, Guolin Ke, Zhaoqing Liu, Denghui Lu, Koki Muraoka, Hananeh Oliaei, Anurag Kumar Singh, Haohui Que, Weihong Xu, Zhangmancang Xu, Yong-Bin Zhuang, Jiayu Dai, Timothy J. Giese, Weile Jia, Ben Xu, Darrin M. York, Linfeng Zhang, Han Wang

arXiv:2502.19161v2physics.chem-ph

TL;DR

DeePMD-kit v3 addresses integration challenges from packages built around different machine-learning frameworks. It introduces a pluggable multiple-backend architecture with a unified interface, enabling backend selection for efficient molecular-dynamics computations.

  • Problem

    Packages built around different machine-learning frameworks create significant integration challenges for developers.

  • Method

    DeePMD-kit v3 introduces a pluggable multiple-backend framework that provides a unified interface while leveraging external machine-learning frameworks.

  • Results

    The framework enables selecting the most efficient backend for molecular-dynamics performance and achieves approximately a 40% reduction in training time compared with the dynamic graph.

  • Takeaways & Limitations

    Backend flexibility supports more efficient molecular-dynamics workflows while allowing users to work through a unified interface.

  • Takeaways & Limitations

    Only the DPA-1 (L=0) model supports model compression, and some operations are not standard in TensorFlow.

Abstract

from arXiv · show

In recent years, machine learning potentials (MLPs) have become indispensable tools in physics, chemistry, and materials science, driving the development of software packages for molecular dynamics (MD) simulations and related applications. These packages, typically built on specific machine learning frameworks such as TensorFlow, PyTorch, or JAX, face integration challenges when advanced applications demand communication across different frameworks. The previous TensorFlow-based implementation of DeePMD-kit exemplified these limitations. In this work, we introduce DeePMD-kit version 3, a significant update featuring a multi-backend framework that supports TensorFlow, PyTorch, JAX, and PaddlePaddle backends, and demonstrate the versatility of this architecture through the integration of other MLPs packages and of Differentiable Molecular Force Field. This architecture allows seamless backend switching with minimal modifications, enabling users and developers to integrate DeePMD-kit with other packages using different machine learning frameworks. This innovation facilitates the development of more complex and interoperable workflows, paving the way for broader applications of MLPs in scientific research.

1 Introduction

MLP packages increasingly require interoperability across machine-learning frameworks, but framework-specific implementations and inconsistent interfaces make integration difficult. DeePMD-kit v3 addresses this with interchangeable TensorFlow, PyTorch, JAX, and PaddlePaddle backends while preserving established interfaces and broad software integration.

  • MLPs support applications across physics, chemistry, molecular biology, and materials science, including molecular dynamics and free-energy simulations.These applications require accurate potential-energy and force calculations.
  • Increasingly complex applications require interoperability among packages built on different machine-learning frameworks, creating challenges for developers.Frameworks may differ in performance optimization, hardware compatibility, and available capabilities.
  • The earlier TensorFlow-based DeePMD-kit complicated integration with packages using other deep-learning frameworks and made consistent training, inference, and simulation interfaces inconvenient and inefficient.
  • DeePMD-kit v3 introduces a multi-backend framework integrating TensorFlow with PyTorch, JAX, and PaddlePaddle backends designed for switching with minimal modifications.
  • The framework reuses established interfaces and maintains integration with molecular-dynamics and related software packages while extending compatibility with packages using the new backends.The ecosystem includes packages such as LAMMPS, i-PI, AMBER, CP2K, OpenMM, GROMACS, ASE, and ABACUS.
  • Multiple backends also allow DeePMD-kit to use framework-specific features and optimizations, improving its flexibility and performance.

2 Software Description

DeePMD-kit v3 introduces a multiple-backend framework that presents unified interfaces while allowing backend-specific implementations, model conversion, and modular extension. The redesign also adds reusable model components and supports broader model integration, while addressing consistency and backend-specific constraints.

  • Multiple-backend framework: DeePMD-kit v3 supports multiple machine-learning backends through a unified interface while retaining backend-specific implementations.Users specify the backend during training, while interfaces and training parameters remain backend-agnostic.
  • Multiple-backend framework: Models can be loaded through Python or C++ interfaces, automatically matched to a backend, and converted between backend-specific formats.Model serialization and deserialization in the Python interface enable conversion between backends.
  • Multiple-backend framework: Developers can add new backends modularly without modifying existing Python or C/C++ APIs or external-package interfaces.The Array API provides a route for implementing new backends and models when supported by the target framework.
  • Multiple-backend framework: The framework tests equivalent models across backends to check whether they produce the same results.This consistency requirement is identified as a challenge of the multiple-backend design.
  • Backends: The DP, TensorFlow, PyTorch, JAX, and PaddlePaddle backends provide distinct implementation choices, including JAX JIT compilation and PyTorch dynamic graphs.The DP backend serves as a reference implementation, while JAX uses the Array API and PyTorch uses TorchScript for serialization.
  • Backends: The JAX backend addresses dynamic molecular-dynamics inputs by combining compiled JIT computation with TensorFlow processing of dynamically shaped tensors.Neighbor lists and other dynamic-shape tensors are handled by TensorFlow before being passed to the compiled JAX model.
  • Backends: The DP backend provides a shared foundation for other backends, and its static-graph implementation reports approximately a 40% training-time reduction compared with a dynamic graph.The DP backend uses NumPy by default to minimize dependencies, without gradient computation or GPU acceleration.
  • New design principles in version 3: Version 3 reorganizes model development around atomic models and descriptor blocks, improving maintainability while preserving compatibility with standard potential-energy architectures.Descriptor blocks are modular and extensible, and the refactored design supports graph neural network models and efficient multi-node or multi-GPU simulations.

3 Selected Extensions

DeePMD-kit v3 adds a plugin system and multi-backend support that integrate external MLP packages and long-range force-field methods. These extensions enable shared training components, cross-model comparisons, and hybrid descriptions of interactions.

  • Plugin system: DeePMD-kit v3 introduces a plugin system for integrating external models across its backends.The section demonstrates integrations with other MLP packages and Differentiable Molecular Force Field.
  • Motivation: Earlier MLP packages were often restricted to one neural-network type, complicating package use, maintenance, and controlled benchmarking.The earlier DeePMD-kit also faced integration difficulties because most external packages were not TensorFlow-based.
  • External MLPs: The PyTorch backend enabled DeePMD-GNN to integrate the MACE and NequIP models.This was presented as a consequence of developing the multi-backend framework and PyTorch backend.
  • External MLPs: The plugin extends the trainable model module while reusing DeePMD-kit's loss and learning-rate modules, allowing models to be compared under identical conditions.This design separates model integration from shared optimization components.
  • Long-range interactions: PyTorch DMFF integration combines traditional force fields for long-range interactions with MLPs for short- and mid-range interactions.The DMFF implementation includes Ewald summation and charge equilibration methods.

4 Benchmark

The benchmark evaluates DeePMD-kit v3 across models, precisions, backends, GPU cards, and water-system sizes. Performance depends on the model and hardware: JAX often leads for demanding models, while no backend is consistently fastest.

  • Benchmark setup: The benchmark measures MD speed across TensorFlow, PyTorch, and JAX using three DPA models, two precisions, and three NVIDIA GPUs.Water systems with varying atom counts were simulated, with each calculation repeated 500 times for average speed.
  • Limitations: At the reported time, only DPA-1 (L=0) supports model compression, TensorFlow lacks DPA-2 support, JAX lacks model compression, and PaddlePaddle remains under development.These package limitations constrain the available model-backend combinations.
  • Performance results: For DPA-1 without attention layers, the compressed TensorFlow model is fastest.The benchmark compares DPA-1 (L=0), DPA-1 (L=2), and DPA-2 models.
  • Performance results: For more computationally demanding DPA-1 (L=2) and DPA-2 models, JAX generally achieves the highest performance.Exceptions occur for DPA-2 FP32 with 12,288 atoms on H100, where PyTorch is faster, and DPA-1 (L=2) FP64 on 4090, where TensorFlow is faster.
  • Resource considerations: Different backends use different amounts of GPU memory, influencing backend choice according to available computational resources.The comparison spans H100, A800, and 4090 GPU cards.
  • Conclusion: No single backend consistently outperforms the others across all models, making backend selection useful for identifying the most efficient option for specific needs.The benchmark therefore supports choosing backends according to model and resource conditions.

5 Conclusions

DeePMD-kit v3 addresses predecessor limitations with a multi-backend framework supporting four machine-learning backends. It enables backend switching and integration with external software while supporting more flexible workflows for machine-learning-potential applications.

  • DeePMD-kit v3 addresses critical limitations of its predecessor by supporting TensorFlow, PyTorch, JAX, and PaddlePaddle.
  • Backend switching requires minimal changes, providing a flexible and interoperable platform for researchers and developers.
  • The framework facilitates integration with diverse software ecosystems and supports more complex workflows for molecular dynamics and related applications.
  • The best molecular-dynamics performance can be achieved by selecting the most efficient backend for a specific model and computational resource.
  • The authors anticipate broader accessibility and applicability of machine-learning potentials in physics, chemistry, and materials science.

Supporting Information Available

The supporting information reports molecular-dynamics performance comparisons across model variants, numerical precisions, machine-learning backends, and GPU resources, and includes an example of language-model-assisted backend code conversion.

  • Performance comparisons cover DPA-1 and DPA-2 models across TensorFlow, PyTorch, and JAX backends.
  • The comparisons evaluate both double-precision FP64 and single-precision FP32 molecular-dynamics performance.
  • The reported benchmarks use H100, A800, and 4090 GPU cards as computational resources.
  • The supporting information includes an example of using large language models for code conversion between backends.

TOC Graphic

The supplied TOC-graphic material consists primarily of the paper’s author, affiliation, and preprint metadata, spanning institutions in China, North America, Europe, and Japan.

  • The paper lists a large, multidisciplinary author team including Jinzhe Zeng, Han Wang, and other contributors.
  • The document is identified as arXiv:2502.19161v2, dated 27 February 2025.
  • The affiliations include universities, research institutes, laboratories, and companies across multiple countries.

1 Tables and Figures

The supporting tables and figures compare molecular-dynamics performance for DPA model variants across precisions, backends, water-system sizes, and GPU hardware.

  • Tables S1–S3 report molecular-dynamics performance for DPA-1 and DPA-2 models in TensorFlow, PyTorch, and JAX.
  • The tables cover 80 GB NVIDIA H100, 40 GB NVIDIA A100, and 24 GB NVIDIA 4090 GPUs.
  • The evaluations use FP64 and FP32 precision and water systems with different numbers of atoms.
  • Compressed TensorFlow and PyTorch models are identified as TFc and PTc, while OOM denotes GPU memory exceeding the listed capacity.
  • Figure S1 compares ms/step performance for DPA-1 and DPA-2 models on a single 40 GB NVIDIA A800 GPU.

Conversion between Backends

The paper demonstrates backend conversion by using a large language model to translate DeepPot-SE code into a PyTorch implementation. The converted implementation preserves array operations, tensor shapes, masking, aggregation, reshaping, normalization, and returned features.

  • A DeepSeek R1 example converts DeepPot-SE code into a PyTorch implementation.
  • The implementation selects an array namespace from coordinate, atom-type, and neighbor-list inputs to support backend-compatible operations.
  • The code reshapes extended coordinates and exclusion masks by merging frame and local-atom axes before processing neighbor features.
  • The converted routine returns aggregated features, sliced components, placeholders, and weights using the same output structure as the source implementation.
  • The conversion preserves input dtype and initializes outputs on the same device as the input.
  • Per-type processing applies masks, computes radial gates and aggregated contributions, forms final feature products, and reshapes the resulting features.
Loading 2502.19161v2…