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

This paper addresses a practical but increasingly central bottleneck in machine learning potentials (MLPs) for atomistic simulation: most MLP software is tightly coupled to a single deep-learning framework (TensorFlow, PyTorch, JAX, or PaddlePaddle). As scientific workflows become more complex—mixing training pipelines, inference engines, differentiable force-field components, and simulation backends—developers and users face friction when components rely on different frameworks. The research question is therefore not about a new potential model’s physics accuracy, but about software architecture: can DeePMD-kit be redesigned so that the same user-facing interfaces (training, inference, molecular dynamics integration) work across multiple machine-learning backends, enabling interoperability and reducing integration costs?

The significance is twofold. First, it directly impacts reproducibility and usability: users can switch backends to match hardware constraints, performance needs, or ecosystem compatibility without rewriting workflows. Second, it enables “composability” of the broader MLP ecosystem. The authors position DeePMD-kit v3 as a hub that can integrate external models and differentiable force-field modules that may be implemented in different frameworks. In the broader context of materials science and computational chemistry, where end-to-end pipelines often span data generation (e.g., active learning), model training, and large-scale molecular dynamics (MD), backend interoperability is a prerequisite for scalable and maintainable research software.

Methodologically, the paper is primarily a software engineering and benchmarking study. The authors introduce DeePMD-kit v3’s multiple-backend framework, designed to be pluggable while preserving backward compatibility. They explicitly state that no breaking changes were made to the existing Python and C/C++ APIs from DeePMD-kit v2, so existing integrations with MD and workflow packages can continue to work. The framework uses a unified set of interfaces for users and developers; backend-specific implementations are invoked under the hood. A key engineering challenge is ensuring that the “same model” yields consistent results across backends. To address this, the authors implement serialization/deserialization and develop tests that compare outputs when models are converted between backends.

The paper also describes several design principles that enable backend-agnostic development. They introduce metaprogramming to generate backend-specific classes from a reference implementation (the DP backend). They add an “atomic model” abstraction that decomposes learned quantities into sums of atomic contributions, enabling developers to implement only per-atom terms while the framework handles summation and derives forces and virials via derivatives with respect to coordinates and cell vectors. They further refactor computations that previously relied on custom TensorFlow operators (e.g., neighbor list and coordinate matrix gradients) so that other backends can use standard operators and automatic differentiation to support higher-order derivatives such as Hessians.

A major practical extension is support for graph neural network (GNN) style models (e.g., DPA-2). For distributed/multi-GPU MD, the authors implement a customized C++ operator using MPI to exchange atom and edge features between processors or GPU cards, avoiding the high overhead of rebuilding neighbor lists for ghost atoms in an extended cutoff.

For benchmarking, the authors evaluate MD performance (time per MD step in milliseconds) using DeePMD-kit v3 integrated with LAMMPS. They test three model families: DPA-1 without attention layers ( L=0), DPA-1 with two attention layers ( L=2), and DPA-2 (medium). Benchmarks are run in both single and double precision on three GPU types: NVIDIA H100 (80 GB), A800 (40 GB), and 4090 (24 GB). Each calculation is repeated 500 times to obtain an average speed. The benchmark uses a water system with varying atom counts, and results are reported in the main figures and supporting tables. The paper also notes backend limitations at the time of writing: only DPA-1 (L=0) supports model compression; TensorFlow does not support DPA-2; JAX does not support model compression; and PaddlePaddle is still under development.

The key quantitative results are the relative performance trends across backends and model types, plus concrete examples of backend-dependent speed and memory behavior. From the supporting tables:

1) On H100 (80 GB), for DPA-1 (L=0) in FP64, the fastest backend at 12288 atoms is TensorFlow compressed (TFc) at 34.46 ms/step, compared with PyTorch at 41.09 ms/step and JAX at 64.76 ms/step. At 6144 atoms, TFc is 17.50 ms/step versus PyTorch 21.47 and JAX 34.19. 2) On H100 (80 GB), for DPA-1 (L=2) in FP64 at 12288 atoms, the fastest is TensorFlow at 14.09 ms/step (JAX is 21.46 and PyTorch is 30.84). For DPA-1 (L=2) in FP32 at 12288 atoms, TensorFlow is 12.11 ms/step versus PyTorch 18.43 and JAX 51.97. 3) For DPA-2 (medium) on H100 in FP64, the fastest backend at 12288 atoms is PyTorch at 217.21 ms/step, while TensorFlow is 286.22 ms/step and JAX is “OOM” (out of memory) at that size. In FP32 at 12288 atoms, PyTorch is again best at 132.45 ms/step, with TensorFlow 146.37 and JAX 240.44. 4) On A100 (40 GB), for DPA-2 in FP64 at 6144 atoms, TensorFlow is 231.34 ms/step, PyTorch is 315.86 ms/step, and JAX is 239.78 ms/step; the best is TensorFlow. At 12288 atoms in FP64, TensorFlow is 69.48 ms/step (note: this corresponds to DPA-1 (L=0) in the table; for DPA-2 at 12288 atoms, JAX and others are marked OOM in that row), illustrating that memory constraints can dominate. 5) On 4090 (24 GB), for DPA-1 (L=0) in FP32 at 1536 atoms, TensorFlow is 6.24 ms/step, PyTorch is 7.87, and JAX is 12.22. For DPA-1 (L=2) in FP32 at 3072 atoms, TensorFlow is 5.49 ms/step, PyTorch is 8.23, and JAX is 21.69.

Overall, the authors’ main performance claim is qualitative but supported by these numbers: no single backend consistently outperforms the others across all models and precisions. They state that for DPA-1 without attention layers, compressed TensorFlow is fastest; for more computationally demanding models (DPA-1 with attention and DPA-2), JAX generally achieves the highest performance, with explicit exceptions: DPA-2 FP32 at 12288 atoms on H100 where PyTorch outperforms JAX, and DPA-1 (L=2) FP64 on 4090 where TensorFlow outperforms JAX. They also emphasize that GPU memory usage differs by backend, affecting which backend is feasible for a given system size.

Limitations are acknowledged in two ways. First, the paper explicitly lists current feature gaps: TensorFlow backend lacks DPA-2 support, JAX lacks model compression, and PaddlePaddle is under development. Second, the benchmark scope is limited to a specific set of models (DPA-1 variants and DPA-2 medium), a water system, and performance measured as ms/step; the paper does not provide accuracy comparisons (e.g., energy/force errors) across backends, implying that the primary contribution is interoperability and performance engineering rather than new scientific results.

Practically, the implications are clear for different stakeholders. Researchers using DeePMD-kit can choose the backend that best matches their hardware and model type, and can integrate DeePMD-kit with external PyTorch/JAX/Paddle-based components. Developers can add new backends or models more easily due to the atomic model and descriptor block abstractions, and due to the Array API-based approach for backend implementation. The paper also demonstrates extensibility via plugins: DeePMD-GNN integrates external GNN potentials (MACE and NequIP) into the PyTorch backend, and a DMFF plugin integrates long-range interaction methods (Ewald summation and charge equilibration) into DeePMD-kit’s PyTorch backend, enabling hybrid MLP + classical long-range physics workflows.

In sum, DeePMD-kit v3’s core contribution is a multi-backend, interoperable architecture that preserves existing APIs while enabling backend switching, model conversion, and integration with external MLP and differentiable force-field packages. The benchmark results show that backend choice can materially affect MD throughput and feasibility under memory constraints, reinforcing the value of the multi-backend design for real-world simulation workflows.