Descriptor DPA-2

4.6. Descriptor DPA-2 PyTorch JAX DP#

Note

Supported backends: PyTorch PyTorch, JAX JAX, DP DP

The DPA-2 model implementation. See https://arxiv.org/abs/2312.15492 for more details.

Training example: examples/water/dpa2/input_torch_medium.json, see README for inputs in different levels.

4.6.1. Requirements of installation PyTorch#

If one wants to run the DPA-2 model on LAMMPS, the customized OP library for the Python interface must be installed when freezing the model.

The customized OP library for the Python interface can be installed by setting environment variable DP_ENABLE_PYTORCH to 1 during installation.

If one runs LAMMPS with MPI, the customized OP library for the C++ interface should be compiled against the same MPI library as the runtime MPI. If one runs LAMMPS with MPI and CUDA devices, it is recommended to compile the customized OP library for the C++ interface with a CUDA-Aware MPI library and CUDA, otherwise the communication between GPU cards falls back to the slower CPU implementation.

4.6.2. Limiations of the JAX backend with LAMMPS JAX#

When using the JAX backend, 2 or more MPI ranks are not supported. One must set map to yes using the atom_modify command.

atom_modify map yes

See the example examples/water/lmp/jax_dpa2.lammps.

4.6.3. Data format#

DPA-2 supports both the standard data format and the mixed type data format.

4.6.4. Type embedding#

Type embedding is within this descriptor with the tebd_dim argument.

4.6.5. Model compression#

Model compression is supported when repinit/tebd_input_mode is strip, but only the repinit part is compressed. An example is given in examples/water/dpa2/input_torch_compressible.json. The performance improvement will be limited if other parts are more expensive.