v3.0.0a0

[njzjz]

DeePMD-kit v3: A multiple-backend framework for deep potentials

We are excited to announce the first alpha version of DeePMD-kit v3. DeePMD-kit v3 allows you to train and run deep potential models on top of TensorFlow or PyTorch. DeePMD-kit v3 also supports the DPA-2 model, a novel architecture for large atomic models.

Highlights

Multiple-backend framework

DeePMD-kit v3 adds a pluggable multiple-backend framework to provide consistent training and inference experiences between different backends. You can:

• Use the same training data and the input script to train a deep potential model with different backends. Switch backends based on efficiency, functionality, or convenience:

# Training a model using the TensorFlow backend
dp --tf train input.json
dp --tf freeze

# Training a mode using the PyTorch backend
dp --pt train input.json
dp --pt freeze

• Use any model to perform inference via any existing interfaces, including dp test, Python/C++/C interface, and third-party packages (dpdata, ASE, LAMMPS, AMBER, Gromacs, i-PI, CP2K, OpenMM, ABACUS, etc). Take an example on LAMMPS:

# run LAMMPS with a TensorFlow backend model
pair_style deepmd frozen_model.pb
# run LAMMPS with a PyTorch backend model
pair_style deepmd frozen_model.pth
# Calculate model deviation using both models
pair_style deepmd frozen_model.pb frozen_model.pth out_file md.out out_freq 100

• Convert models between backends, using dp convert-backend, if both backends support a model:

dp convert-backend frozen_model.pb frozen_model.pth
dp convert-backend frozen_model.pth frozen_model.pb

• Add a new backend to DeePMD-kit much more quickly if you want to contribute to DeePMD-kit.

PyTorch backend: a backend designed for large atomic models and new research

We added the PyTorch backend in DeePMD-kit v3 to support the development of new models, especially for large atomic models.

DPA-2 model: Towards a universal large atomic model for molecular and material simulation

DPA-2 model is a novel architecture for Large Atomic Model (LAM) and can accurately represent a diverse range of chemical systems and materials, enabling high-quality simulations and predictions with significantly reduced efforts compared to traditional methods. The DPA-2 model is only implemented in the PyTorch backend. An example configuration is in the examples/water/dpa2 directory.

The DPA-2 descriptor includes two primary components: repinit and repformer. The detailed architecture is shown in the following figure.

Training strategies for large atomic models

The PyTorch backend has supported multiple training strategies to develop large atomic models.

Parallel training: Large atomic models have a number of hyper-parameters and complex architecture, so training a model on multiple GPUs is necessary. Benefiting from the PyTorch community ecosystem, the parallel training for the PyTorch backend can be driven by torchrun, a launcher for distributed data parallel.

torchrun --nproc_per_node=4 --no-python dp --pt train input.json

Multi-task training: Large atomic models are trained against data in a wide scope and at different DFT levels, which requires multi-task training. The PyTorch backend supports multi-task training, sharing the descriptor between different An example is given in examples/water_multi_task/pytorch_example/input_torch.json.

Finetune: Fine-tune is useful to train a pre-train large model on a smaller, task-specific dataset. The PyTorch backend has supported --finetune argument in the dp --pt train command line.

Developing new models using Python and dynamic graphs

Researchers may feel pain about the static graph and the custom C++ OPs from the TensorFlow backend, which sacrifices research convenience for computational performance. The PyTorch backend has a well-designed code structure written using the dynamic graph, which is currently 100% written with the Python language, making extending and debugging new deep potential models easier than the static graph.

Supporting traditional deep potential models

People may still want to use the traditional models already supported by the TensorFlow backend in the PyTorch backend and compare the same model among different backends. We almost rewrote all of the traditional models in the PyTorch backend, which are listed below:

• Features supported:
  • Descriptor: se_e2_a, se_e2_r, se_atten, hybrid;
  • Fitting: energy, dipole, polar, fparam/apram support
  • Model: standard, DPRc
  • Python inference interface
  • C++ inference interface for energy only
  • TensorBoard
• Features not supported yet:
  • Descriptor: se_e3, se_atten_v2, se_e2_a_mask
  • Fitting: dos
  • Model: linear_ener, DPLR, pairtab, linear_ener, frozen, pairwise_dprc, ZBL, Spin
  • Model compression
  • Python inference interface for DPLR
  • C++ inference interface for tensors and DPLR
  • Paralleling training using Horovod
• Features not planned:
  • Descriptor: loc_frame, se_e2_a + type embedding, se_a_ebd_v2
  • NVNMD

Warning

As part of an alpha release, the PyTorch backend's API or user input arguments may change before the first stable version.

DP backend and format: reference backend for other backends

DP is a reference backend for development that uses pure NumPy to implement models without using any heavy deep-learning frameworks. It cannot be used for training but only for Python inference. As a reference backend, it is not aimed at the best performance but only the correct results. The DP backend uses HDF5 to store model serialization data, which is backend-independent.
The DP backend and the serialization data are used in the unit test to ensure different backends have consistent results and can be converted between each other.
In the current version, the DP backend has a similar supporting status to the PyTorch backend, while DPA-1 and DPA-2 are not supported yet.

Authors

The above highlights were mainly contributed by

• Hangrui Bi (@20171130), in Merge deepmd-pytorch into main repo #3180
• Chun Cai (@caic99), in Merge deepmd-pytorch into main repo #3180
• Junhan Chang (@TablewareBox), in Merge deepmd-pytorch into main repo #3180
• Yiming Du (@nahso), in Merge deepmd-pytorch into main repo #3180
• Guolin Ke (@guolinke), in Merge deepmd-pytorch into main repo #3180
• Xinzijian Liu (@zjgemi), in Merge deepmd-pytorch into main repo #3180
• Anyang Peng (@anyangml), in Chore: Update doc for PyTorch Energy Fitting #3362, Feat: add pair table model to pytorch #3192, Feat: numpy pairtab model #3212, Feat: add ZBL weighted DP model #3210, fix: cuda tests of linear and pair atomic model #3248, Chore: refactor InvarFitting #3266, Feat: Refactor dipole fitting pytorch #3281, Feat: add polarizability fitting net #3296, Feat: add DipoleModel and PolarModel #3309, Fix: float32 bug in UTs #3314, Feat: add dipole consistency test #3321, Feat: add polar consistency test #3327, Feat: add se_r descriptor #3338, Fix: se_r prod_env_mat #3351, Doc: Update the developer guide for the v3 #3376, Doc: Update PT tensor fitting #3385
• Xuejian Qin (@qin2xue3jian4), in Merge deepmd-pytorch into main repo #3180
• Han Wang (@wanghan-iapcm), in breaking: pt: change the virial output dim to 9 #3188, breaking: pt: unify the output of descriptors. #3190, feat: breaking: backend indepdent definition for dp model #3208, atomic model is not required to provide the fitting net #3184, breaking: pt: add dp model format and refactor pt impl for the fitting net. #3199, remove duplicated fitting output check. fix codeql #3202, fix bug of output def: the reduced virial is not defined. #3219, fix model doc str: add shape hint to doc string. #3225, test: better structure for the dpmodel uts #3232, fix: pt: energy model forward lower is not tested and has bugs.  #3235, test: pt: mv all model related tests to the pt/test/model folder #3234, fix skipped tests #3241, support separate r_differentiable and c_differentiable #3240, add hessian support in output def. #3246, fix bug of not passing params in model #3260, fix bug of output dimension in pt InvarFitting  #3274, dp: pt: rearrange folder structure of atomic_model and model. #3268, feat: pt: support user specified rcond for fitting stat #3279, feat: dp and pt: implement exclude types in descriptor se_a #3280, feat: dp and pt: implement fitting exclude types #3282, test the case of permuted index in different frames. #3295, fix: replace all distinguished types by mixed types. #3289, refact: pt: mv all plugin support to base descriptor.  #3340, fix: cast input and output types at model's interface #3352, feat: support exclude atypes in atomic model #3357, add mask to atomic model output when an atomic type exclusion presents. #3389, ut: add null test #3391, Recovered all the skipped test for hybrid descriptor #3400
• Jinzhe Zeng (@njzjz), in docs: add TF icons to platform-specific features #3171, Move model deviation and ase calculator to deepmd_utils #3173, add dpdata driver #3174, docs: rewrite README; deprecate manually written TOC #3179, docs: document PyTorch backend #3193, allow disabling TensorFlow backend during Python installation #3200, pt: add tensorboard and profiler support #3204, pt: set nthreads from env #3205, docs: install pytorch in RTD #3333, docs: dpmodel, model conversion #3360, docs: apply type_one_side=True to se_a and se_r #3364, Hybrid descriptor #3365, setup PyTorch C++ interface build environement #3169, add universal Python inference interface DeepPot #3164, add more abstractmethods to universal DeepPot #3175, cc: reimplement read_file_to_string without calling TensorFlow #3176, Merge TF and PT CLI #3187, Fix PT DeepPot and replace ASE calculator #3186, PT: keep the same checkpoint behavior as TF #3191, add size and replace arguments to deepmd.utils.random.choice #3195, drop tqdm #3194, reorganize tests directory #3196, fix: install CU11 PyTorch in the CU11 docker image #3198, throw errors when PyTorch CXX11 ABI is different from TensorFlow #3201, fix GPU test OOM problem #3207, pt: rename atomic_virial to atom_virial in the model output #3226, pt: apply global logger to pt #3222, pt: apply global user set precision to pt #3220, fix DP_ENABLE_TENSORFLOW support #3229, pt: rename atomic_virial to atom_virial in the model output #3226, drop deepmd.tf.cluster.slurm #3239, add category property to OutputVariableDef #3228, issue template: change TF version to backend version #3244, refactor print summary #3243, refactor DeepEval #3213, backend-indepedent dp test #3249, pt: infer model type from ModelOutputDef #3250, tf: support checkpoint path (instead of directory) in dp freeze #3254, add get_type_map method to model; export model methods  #3247, pt: support loading frozen models in DeepEval #3253, add neighbor stat support with NumPy and PyTorch implementation #3271, support TF se_e2_a serialization; add a common test fixture to compare TF, PT, and DP models #3263, pt: add exported methods to BaseAtomicModel #3258, pt: fix torchscript converage #3276, pt: refactor data stat #3285, consistent energy fitting #3286, dp&pt: let DPAtomicModel fetch attributes from Fitting #3292, pluggable backend #3294, pt: process frames in parallel for env mat stat #3293, pt: avoid set_default_dtype in tests #3303, fix neighbor stat mixed_types input #3304, allow both absulute and relative tolerance when testing consistency #3308, pt: explicitly set device #3307, consistent energy model #3306, merge common subcommands in cli #3316, pt: fix se_e2_a precision cast #3315, pt: export model_output_type instead of model_output_def #3318, feat: convert model files between backends #3323, store type in descriptor serialization data #3325, feat: add NumPy DeepPot #3332, store type in fitting serialization data #3331, pt: remove env.DEVICE in all forward functions #3330, feat(pt/dpmodel): support type_one_side in se_e2_a #3339, add BaseModel; store type in serialization #3335, apply PluginVariant and make_plugin_registry to classes #3346, add @version to serialization data #3349, pt: support --init-frz-model #3350, merge compute_output_stat #3310, feat(pt): support fparam/aparam in DeepEval #3356, pt: fix se_a type_one_side performance degradation #3361, pt: apply argcheck to pt #3342, feat: update sel by statistics #3348, feat(pt): support fparam/aparam in C++ DeepPot #3358, fix se_r consistency #3366, sync descriptor alias #3374, feat: atom_ener in energy fitting #3370, docs: DPRc for PT, DPModel #3373, pt: supprot --output in dp train #3377, fix: prevent deepmd.tf be imported globally #3382, pt: print data summary #3383, pt: expand systems before training #3384, pt: add fparam/aparam data requirements #3386, breaking: change DeepTensor output dim from nsel_atoms to natoms #3390, pt: ban torch.testing.assert_allclose #3395, allow loading either nsel or natoms atomic tensor data #3394, do not return g2, h2, sw in hybrid descriptors #3396, format training logging #3397
• Chengqian Zhang (@Chengqian-Zhang), in Merge deepmd-pytorch into main repo #3180
• Duo Zhang (@iProzd), in Merge deepmd-pytorch into main repo #3180, Fix GPU UTs #3203, breaking: pt: remove data stat from model init #3245, breaking: pt: remove data preprocess from data stat #3261, Implement hessian autodiff calculation #3262, Fix single-task training&data stat #3355, pt: Fix multitask neighbor stat #3367, pt: refact training code #3359, Add trainable settings for pt #3371, Doc: Update PT Multi-task #3387, pt: fix params with no docstrs #3388, pt: Add support for dipole and polar training #3380, Merge change_energy_bias and fix finetune #3378
• Xiangyu Zhang (@CaRoLZhangxy), in cc: add torch backend support for Multiple backend #3162, pt:Removed the data system implemented in PyTorch #3287, pt: add necessary jit.export #3337, pt: add real atoms select in c++ interface #3375, doc: add PyTorch parallel training content #3379

Breaking changes

• Python 3.7 support is dropped. by @njzjz in breaking: drop Python 3.7 support #3185
• We require all model files to have the correct filename extension for all interfaces so a corresponding backend can load them. TensorFlow model files must end with .pb extension.
• Python class DeepTensor (including DeepDiople and DeepPolar) now returns atomic tensor in the dimension of natoms instead of nsel_atoms. by @njzjz in breaking: change DeepTensor output dim from nsel_atoms to natoms #3390
• For developers: the Python module structure is fully refactored. The old deepmd module was moved to deepmd.tf without other API changes, and deepmd_utils was moved to deepmd without other API changes. by @njzjz in breaking: move deepmd to deepmd.tf #3177, breaking: move deepmd_utils to deepmd #3178

Other changes

Enhancement

• Neighbor stat for the TensorFlow backend is 80x accelerated. by @njzjz in tf: refactor neighbor stat #3275
• i-PI: remove normalize_coord by @njzjz in ipi: remove normalize_coord #3257
• LAMMPS: fix_dplr.cpp delete redundant setup and set atom->image when pre_force by @shiruosong in fix_dplr.cpp delete redundant setup #3344, fix_dplr.cpp set atom->image when pre_force #3345
• Bump scikit-build-core to 0.8 by @njzjz in bump scikit-build-core to 0.8 #3369
• Bump LAMMPS to stable_2Aug2023_update3 by @njzjz in bump LAMMPS to stable_2Aug2023_update3 #3399
• Add fparam/aparam support for fine-tune by @njzjz in tf: add fparam/aparam support for finetune #3313
• TF: remove freeze warning for optional nodes by @njzjz in tf: remove freeze warning for optional nodes #3381

CI/CD

• Build macos-arm64 wheel on M1 runners by @njzjz in build macos-arm64 wheel on M1 runners #3206
• Other improvements and fixes to GitHub Actions by @njzjz in pin docker actions to major versions #3238, improve gh actions #3283, speed up cuda test #3284, fix gh actions issues #3288, gh actions: fix branches ignore pattern & fix activity types #3290, fix AlmaLinux GPG key error #3326
• Enable docstring code format by @njzjz in enable docstring code format #3267

Bugfix

• Fix TF 2.16 compatibility by @njzjz in bump python to 3.12 in the test environment #3343
• Detect version in advance before building deepmd-kit-cu11 by @njzjz in detect version in advance before building deepmd-kit-cu11 #3172
• C API: change the required shape of electric field to nloc * 3 by @njzjz in c: change the required shape of electric field to nloc * 3 #3237

New Contributors

• @anyangml made their first contribution in Feat: add pair table model to pytorch #3192
• @shiruosong made their first contribution in fix_dplr.cpp delete redundant setup #3344

Full Changelog: v2.2.8...v3.0.0a0

---

This discussion was created from the release v3.0.0a0.

[ShangChien]

Great works 🎉
Question about section of DP backend: How to use pure Numpy for Python inference? Any example?

[njzjz]

First you need to convert an existing model into the .dp format:

dp convert-backend aaa.pb bbb.dp
# or
dp convert-backend aaa.pth bbb.dp

Then bbb.dp can be used in dp test, dp model-devi, or any Python interface (deepmd.infer.DeepPot, dpdata.System.predict, etc).

[ShangChien]

To make a long story short, does this mean that we can get the energy and forces of the structure in the browser by pyodide-related projects 👀 ?
Previous related discussion #3100

[njzjz]

It cannot predict forces as numpy doesn't support automatic gradients, unless someone would like to write the backward functions manually.

[ShangChien]

Awesome DP!

[Epikurz]

Is there any examples for finetuning on DPA-2 pretrained model?

[ShangChien]

Several tutorial on Bohrium website, you can search with DPA tag to find them.
for example: Notebook for DPA-2
hope this helps 😸

[Epikurz]

Does It mean that all the functions of dp-pytorch version have been inplemented in deepmd-kit 3?

[njzjz]

This repo was entirely imported in #3180.

[omidshy]

Nice work!

In the DPA-2 descriptor, the outputs of the env block are the rotational invariant gij and rotational equivariant hij representations. Does this make the DPA-2 an equivariant model?

[wanghan-iapcm]

Nice work!

In the DPA-2 descriptor, the outputs of the env block are the rotational invariant gij and rotational equivariant hij representations. Does this make the DPA-2 an equivariant model?

When discussing the representation of a rotation group, specifically SO(3), it is necessary to introduce the concept of representation space. In the majority of "equivariant" models, this space is typically represented by spherical harmonics. To preserve as much information about the atomic environment as possible, one must increase the completeness (denoted as 'l' in the context of spherical harmonics) of the basis. However, this comes at the cost of increased computational expense.

In the development of DP models, the approach does not involve spanning the atomic environment within a representation space. Instead, it focuses on utilizing the original (equivariant) information without any approximation. This may be the reason why DPA-2 exhibits higher generalizability in our benchmarks, as evidenced by supplementary S2 of https://arxiv.org/abs/2312.15492.

[poyaii]

Hello. Thank you for developing the new version of DeepMd. When will the stable version of this release be available? Is it possible to install it via Conda? Thanks a lot.

[njzjz]

Conda packages are available at the conda-forge/label/deepmd-kit_dev channel: https://anaconda.org/conda-forge/deepmd-kit/files?version=3.0.0a0

There is no ETA for the stable version, but at least we need to resolve all issues in https://github.com/orgs/deepmodeling/projects/19