Unable to run LAMMPS using DPA-2 finetune model

[guiling-2000]

Hi , guys, I'm using the new version of DeePMD-kit v3.0.0b3 and DPA-2.2.0 to train a deep potential models, DeePMD-kit v3.0.0b3 was installed via off-line packages.In the train phase, it runs without problems. However, when I run lmp with frozen model files model.pth , the following error occurs:

LAMMPS (2 Aug 2023)
OMP_NUM_THREADS environment is not set. Defaulting to 1 thread. (src/comm.cpp:98)
  using 1 OpenMP thread(s) per MPI task
DeePMD-kit: Successfully load libcudart.so.12
2024-08-05 21:14:47.348873: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-08-05 21:14:47.407645: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-08-05 21:14:47.407709: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-08-05 21:14:47.409155: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
DeePMD-kit WARNING: Environmental variable DP_INTRA_OP_PARALLELISM_THREADS is not set. Tune DP_INTRA_OP_PARALLELISM_THREADS for the best performance. See https://deepmd.rtfd.io/parallelism/ for more information.
DeePMD-kit WARNING: Environmental variable DP_INTER_OP_PARALLELISM_THREADS is not set. Tune DP_INTER_OP_PARALLELISM_THREADS for the best performance. See https://deepmd.rtfd.io/parallelism/ for more information.
DeePMD-kit WARNING: Environmental variable OMP_NUM_THREADS is not set. Tune OMP_NUM_THREADS for the best performance. See https://deepmd.rtfd.io/parallelism/ for more information.
Loaded 1 plugins from /home/deepmd-kit/lib/deepmd_lmp
Reading data file ...
  triclinic box = (0 0 0) to (117.42006 101.68876 52.741749) with tilt (-58.710031 0 0)
WARNING: Triclinic box skew is large. LAMMPS will run inefficiently. (src/domain.cpp:220)
  1 by 1 by 1 MPI processor grid
  reading atoms ...
  19200 atoms
  read_data CPU = 0.046 seconds
Summary of lammps deepmd module ...
  >>> Info of deepmd-kit:
  installed to:       /home/deepmd-kit
  source:             
  source branch:      HEAD
  source commit:      cbf2de6
  source commit at:   2024-07-27 05:11:58 +0000
  support model ver.: 1.1 
  build variant:      cuda
  build with tf inc:  /home/deepmd-kit/lib/python3.11/site-packages/tensorflow/include;/home/deepmd-kit/include
  build with tf lib:  /home/deepmd-kit/lib/python3.11/site-packages/tensorflow/libtensorflow_cc.so.2
  build with pt lib:  torch;torch_library;/home/deepmd-kit/lib/python3.11/site-packages/torch/lib/libc10.so;/home/conda/feedstock_root/build_artifacts/deepmd-kit_1722057349216/_build_env/targets/x86_64-linux/lib/stubs/libcuda.so;/home/deepmd-kit/lib/libnvrtc.so;/home/deepmd-kit/lib/libnvToolsExt.so;/home/deepmd-kit/lib/libcudart.so;/home/deepmd-kit/lib/python3.11/site-packages/torch/lib/libc10_cuda.so
  set tf intra_op_parallelism_threads: 0
  set tf inter_op_parallelism_threads: 0
  >>> Info of lammps module:
  use deepmd-kit at:  /home/deepmd-kitload model from: model.pth to gpu 0
DeePMD-kit WARNING: Environmental variable DP_INTRA_OP_PARALLELISM_THREADS is not set. Tune DP_INTRA_OP_PARALLELISM_THREADS for the best performance. See https://deepmd.rtfd.io/parallelism/ for more information.
DeePMD-kit WARNING: Environmental variable DP_INTER_OP_PARALLELISM_THREADS is not set. Tune DP_INTER_OP_PARALLELISM_THREADS for the best performance. See https://deepmd.rtfd.io/parallelism/ for more information.
DeePMD-kit WARNING: Environmental variable OMP_NUM_THREADS is not set. Tune OMP_NUM_THREADS for the best performance. See https://deepmd.rtfd.io/parallelism/ for more information.
  >>> Info of model(s):
  using   1 model(s): model.pth 
  rcut in model:      9
  ntypes in model:    2

CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE

Your simulation uses code contributions which should be cited:
- USER-DEEPMD package:
The log file lists these citations in BibTeX format.

CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE

Generated 0 of 1 mixed pair_coeff terms from geometric mixing rule
Neighbor list info ...
  update: every = 1 steps, delay = 0 steps, check = yes
  max neighbors/atom: 2000, page size: 100000
  master list distance cutoff = 11
  ghost atom cutoff = 11
  binsize = 5.5, bins = 33 19 10
  1 neighbor lists, perpetual/occasional/extra = 1 0 0
  (1) pair deepmd, perpetual
      attributes: full, newton on
      pair build: full/bin/atomonly
      stencil: full/bin/3d
      bin: standard
Setting up Verlet run ...
  Unit style    : metal
  Current step  : 0
  Time step     : 0.001
ERROR on proc 0: DeePMD-kit C API Error: DeePMD-kit Error: DeePMD-kit PyTorch backend error: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/deepmd/pt/model/model/ener_model.py", line 56, in forward_lower
    comm_dict: Optional[Dict[str, Tensor]]=None) -> Dict[str, Tensor]:
    _5 = (self).need_sorted_nlist_for_lower()
    model_ret = (self).forward_common_lower(extended_coord, extended_atype, nlist, mapping, fparam, aparam, do_atomic_virial, comm_dict, _5, )
                 ~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _6 = (self).get_fitting_net()
    model_predict = annotate(Dict[str, Tensor], {})
  File "code/__torch__/deepmd/pt/model/model/ener_model.py", line 213, in forward_common_lower
    cc_ext, _36, fp, ap, input_prec, = _35
    atomic_model = self.atomic_model
    atomic_ret = (atomic_model).forward_common_atomic(cc_ext, extended_atype, nlist0, mapping, fp, ap, comm_dict, )
                  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _37 = (self).atomic_output_def()
    training = self.training
  File "code/__torch__/deepmd/pt/model/atomic_model/energy_atomic_model.py", line 50, in forward_common_atomic
    ext_atom_mask = (self).make_atom_mask(extended_atype, )
    _3 = torch.where(ext_atom_mask, extended_atype, 0)
    ret_dict = (self).forward_atomic(extended_coord, _3, nlist, mapping, fparam, aparam, comm_dict, )
                ~~~~~~~~~~~~~~~~~~~~ <--- HERE
    ret_dict0 = (self).apply_out_stat(ret_dict, atype, )
    _4 = torch.slice(torch.slice(ext_atom_mask), 1, None, nloc)
  File "code/__torch__/deepmd/pt/model/atomic_model/energy_atomic_model.py", line 93, in forward_atomic
      pass
    descriptor = self.descriptor
    _16 = (descriptor).forward(extended_coord, extended_atype, nlist, mapping, comm_dict, )
           ~~~~~~~~~~~~~~~~~~~ <--- HERE
    descriptor0, rot_mat, g2, h2, sw, = _16
    fitting_net = self.fitting_net
  File "code/__torch__/deepmd/pt/model/descriptor/dpa2.py", line 84, in forward
    repformers3 = self.repformers
    _17 = nlist_dict[_1(_16, (repformers3).get_nsel(), )]
    _18 = (repformers1).forward(_17, extended_coord, extended_atype, g11, mapping0, comm_dict0, )
           ~~~~~~~~~~~~~~~~~~~~ <--- HERE
    g12, g2, h2, rot_mat, sw, = _18
    concat_output_tebd = self.concat_output_tebd
  File "code/__torch__/deepmd/pt/model/descriptor/repformers.py", line 228, in forward
      ret = ops.deepmd.border_op(_25, _26, _27, _28, _29, g10, _31, _32, _33)
      g1_ext, comm_dict6, mapping6 = torch.unsqueeze(ret[0], 0), comm_dict7, mapping2
    _34 = (_00).forward(g1_ext, g23, h2, nlist0, nlist_mask, sw1, )
           ~~~~~~~~~~~~ <--- HERE
    g11, g24, h21, = _34
    if torch.__is__(comm_dict6, None):
  File "code/__torch__/deepmd/pt/model/descriptor/repformer_layer.py", line 172, in forward
      smooth0 = self.smooth
      epsilon0 = self.epsilon
      _19 = (self).symmetrization_op(gg16, h2, nlist_mask, sw, axis_neuron0, smooth0, epsilon0, )
             ~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
      _20 = torch.append(g1_mlp, _19)
      gg15 : Optional[Tensor] = gg16
  File "code/__torch__/deepmd/pt/model/descriptor/repformer_layer.py", line 330, in symmetrization_op
    _52 = __torch__.deepmd.pt.model.descriptor.repformer_layer._cal_hg
    _53 = __torch__.deepmd.pt.model.descriptor.repformer_layer._cal_grrg
    h2g2 = _52(g2, h2, nlist_mask, sw, smooth, epsilon, )
           ~~~ <--- HERE
    return _53(h2g2, axis_neuron, )
class Atten2Map(Module):
  File "code/__torch__/deepmd/pt/model/descriptor/repformer_layer.py", line 509, in _cal_hg
  _88 = __torch__.deepmd.pt.model.descriptor.repformer_layer._apply_switch
  nb, nloc, nnei, _89, = torch.size(g2)
  g20 = _87(g2, nlist_mask, )
        ~~~ <--- HERE
  if torch.__not__(smooth):
    _90 = torch.sum(torch.type_as(nlist_mask, g20), [-1])
  File "code/__torch__/deepmd/pt/model/descriptor/repformer_layer.py", line 537, in _apply_nlist_mask
    nlist_mask: Tensor) -> Tensor:
  _98 = torch.bitwise_not(torch.unsqueeze(nlist_mask, -1))
  return torch.masked_fill(gg, _98, 0.)
         ~~~~~~~~~~~~~~~~~ <--- HERE
def _apply_switch(gg: Tensor,
    sw: Tensor) -> Tensor:

Traceback of TorchScript, original code (most recent call last):
  File "/home/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/model/model/ener_model.py", line 109, in forward_lower
        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
    ):
        model_ret = self.forward_common_lower(
                    ~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            extended_coord,
            extended_atype,
  File "/home/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/model/model/make_model.py", line 261, in forward_common_lower
            )
            del extended_coord, fparam, aparam
            atomic_ret = self.atomic_model.forward_common_atomic(
                         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                cc_ext,
                extended_atype,
  File "/home/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/model/atomic_model/base_atomic_model.py", line 241, in forward_common_atomic
    
        ext_atom_mask = self.make_atom_mask(extended_atype)
        ret_dict = self.forward_atomic(
                   ~~~~~~~~~~~~~~~~~~~ <--- HERE
            extended_coord,
            torch.where(ext_atom_mask, extended_atype, 0),
  File "/home/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/model/atomic_model/dp_atomic_model.py", line 189, in forward_atomic
        if self.do_grad_r() or self.do_grad_c():
            extended_coord.requires_grad_(True)
        descriptor, rot_mat, g2, h2, sw = self.descriptor(
                                          ~~~~~~~~~~~~~~~ <--- HERE
            extended_coord,
            extended_atype,
  File "/home/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/model/descriptor/dpa2.py", line 652, in forward
            g1 = g1_ext
        # repformer
        g1, g2, h2, rot_mat, sw = self.repformers(
                                  ~~~~~~~~~~~~~~~ <--- HERE
            nlist_dict[
                get_multiple_nlist_key(
  File "/home/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/model/descriptor/repformers.py", line 492, in forward
                )
                g1_ext = ret[0].unsqueeze(0)
            g1, g2, h2 = ll.forward(
                         ~~~~~~~~~~ <--- HERE
                g1_ext,
                g2,
  File "/home/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/model/descriptor/repformer_layer.py", line 1126, in forward
            assert gg1 is not None
            g1_mlp.append(
                self.symmetrization_op(
                ~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                    gg1,
                    h2,
  File "/home/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/model/descriptor/repformer_layer.py", line 991, in symmetrization_op
        nb, nloc, nnei, _ = g2.shape
        # nb x nloc x 3 x ng2
        h2g2 = self._cal_hg(g2, h2, nlist_mask, sw, smooth=smooth, epsilon=epsilon)
               ~~~~~~~~~~~~ <--- HERE
        # nb x nloc x (axisxng2)
        g1_13 = self._cal_grrg(h2g2, axis_neuron)
  File "/home/deepmd-kit/lib/python3.11/site-packages/deepmd/pt/model/descriptor/repformer_layer.py", line 90--------------------------------------------------------------------------
MPI_ABORT was invoked on rank 0 in communicator MPI_COMM_WORLD
with errorcode 1.

NOTE: invoking MPI_ABORT causes Open MPI to kill all MPI processes.
You may or may not see output from other processes, depending on
exactly when Open MPI kills them.
--------------------------------------------------------------------------

I am using bohrium to study the dpa-2 course but there is no tutorial on running lmp. Has anyone encountered such error before? Thanks in advance.