support DPA2/DPA3 inference

Author: HydrogenSulfate

deepmd/pd/entrypoints/main.py

@@ -388,14 +388,25 @@ def freeze(
         model.forward_lower = paddle.jit.to_static(
             model.forward_lower,
             input_spec=[
-                InputSpec([1, -1, 3], dtype="float64", name="coord"),  # extended_coord
-                InputSpec([1, -1], dtype="int32", name="atype"),  # extended_atype
+                InputSpec(
+                    [1, -1, 3], dtype="float64", name="extended_coord"
+                ),  # extended_coord
+                InputSpec(
+                    [1, -1], dtype="int32", name="extended_atype"
+                ),  # extended_atype
                 InputSpec([1, -1, -1], dtype="int32", name="nlist"),  # nlist
                 InputSpec([1, -1], dtype="int64", name="mapping"),  # mapping
                 None,  # fparam
                 None,  # aparam
                 True,  # do_atomic_virial
-                None,  # comm_dict
+                [
+                    InputSpec([-1], "int32", name="send_list"),
+                    InputSpec([-1], "int32", name="send_proc"),
+                    InputSpec([-1], "int32", name="recv_proc"),
+                    InputSpec([-1], "int32", name="send_num"),
+                    InputSpec([-1], "int32", name="recv_num"),
+                    InputSpec([-1], "int64", name="communicator"),
+                ],  # comm_dict
             ],
             full_graph=True,
         )

deepmd/pd/model/descriptor/dpa2.py

@@ -798,8 +798,8 @@ def forward(
             assert self.tebd_transform is not None
             g1 = g1 + self.tebd_transform(g1_inp)
         # mapping g1
-        if comm_dict is None:
-            assert mapping is not None
+        if comm_dict is None or len(comm_dict) == 0:
+            assert mapping.numel() > 0
             mapping_ext = (
                 mapping.reshape([nframes, nall])
                 .unsqueeze(-1)

deepmd/pd/model/descriptor/repflows.py

@@ -11,6 +11,10 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
+from deepmd.pd.cxx_op import (
+    ENABLE_CUSTOMIZED_OP,
+    paddle_ops_deepmd,
+)
 from deepmd.pd.model.descriptor.descriptor import (
     DescriptorBlock,
 )
@@ -35,6 +39,9 @@
 from deepmd.pd.utils.exclude_mask import (
     PairExcludeMask,
 )
+from deepmd.pd.utils.spin import (
+    concat_switch_virtual,
+)
 from deepmd.pd.utils.utils import (
     ActivationFn,
 )
@@ -49,6 +56,29 @@
     RepFlowLayer,
 )
 
+if not ENABLE_CUSTOMIZED_OP:
+
+    def border_op(
+        argument0,
+        argument1,
+        argument2,
+        argument3,
+        argument4,
+        argument5,
+        argument6,
+        argument7,
+        argument8,
+    ) -> paddle.Tensor:
+        raise NotImplementedError(
+            "border_op is not available since customized Paddle OP library is not built when freezing the model. "
+            "See documentation for DPA3 for details."
+        )
+
+    # Note: this hack cannot actually save a model that can be run using LAMMPS.
+    paddle_ops_deepmd_border_op = border_op
+else:
+    paddle_ops_deepmd_border_op = paddle_ops_deepmd.border_op
+
 
 @DescriptorBlock.register("se_repflow")
 class DescrptBlockRepflows(DescriptorBlock):
@@ -418,13 +448,14 @@ def forward(
     ):
         parallel_mode = comm_dict is not None
         if not parallel_mode:
-            assert mapping is not None
+            if paddle.in_dynamic_mode():
+                assert mapping is not None and mapping.numel() > 0
         nframes, nloc, nnei = nlist.shape
         nall = extended_coord.reshape([nframes, -1]).shape[1] // 3
         atype = extended_atype[:, :nloc]
         # nb x nloc x nnei
         exclude_mask = self.emask(nlist, extended_atype)
-        nlist = paddle.where(exclude_mask != 0, nlist, -1)
+        nlist = paddle.where(exclude_mask != 0, nlist, paddle.full_like(nlist, -1))
         # nb x nloc x nnei x 4, nb x nloc x nnei x 3, nb x nloc x nnei x 1
         dmatrix, diff, sw = prod_env_mat(
             extended_coord,
@@ -447,7 +478,7 @@ def forward(
             :, :, : self.a_sel
         ]
         a_nlist = nlist[:, :, : self.a_sel]
-        a_nlist = paddle.where(a_dist_mask, a_nlist, -1)
+        a_nlist = paddle.where(a_dist_mask, a_nlist, paddle.full_like(a_nlist, -1))
         _, a_diff, a_sw = prod_env_mat(
             extended_coord,
             a_nlist,
@@ -497,7 +528,8 @@ def forward(
         angle_input = cosine_ij.unsqueeze(-1) / (np.pi**0.5)
 
         if not parallel_mode and self.use_loc_mapping:
-            assert mapping is not None
+            if paddle.in_dynamic_mode():
+                assert mapping is not None and mapping.numel() > 0
             # convert nlist from nall to nloc index
             nlist = paddle.take_along_axis(
                 mapping,
@@ -542,7 +574,8 @@ def forward(
 
         # nb x nall x n_dim
         if not parallel_mode:
-            assert mapping is not None
+            if paddle.in_dynamic_mode():
+                assert mapping is not None and mapping.numel() > 0
             mapping = (
                 mapping.reshape([nframes, nall])
                 .unsqueeze(-1)
@@ -552,14 +585,81 @@ def forward(
             # node_ebd:     nb x nloc x n_dim
             # node_ebd_ext: nb x nall x n_dim [OR] nb x nloc x n_dim when not parallel_mode
             if not parallel_mode:
-                assert mapping is not None
+                if paddle.in_dynamic_mode():
+                    assert mapping is not None and mapping.numel() > 0
                 node_ebd_ext = (
                     paddle.take_along_axis(node_ebd, mapping, 1, broadcast=False)
                     if not self.use_loc_mapping
                     else node_ebd
                 )
             else:
-                raise NotImplementedError("Not implemented")
+                assert len(comm_dict) >= 6
+                has_spin = len(comm_dict) >= 7
+                if not has_spin:
+                    n_padding = nall - nloc
+                    # node_ebd = paddle.nn.functional.pad(
+                    #     node_ebd.squeeze(0), [0, 0, 0, n_padding], value=0.0
+                    # )
+                    # [nframes, nloc, tebd_dim]
+                    _shapes = node_ebd.shape[1:]
+                    _shapes[1] = n_padding
+                    node_ebd = paddle.concat(
+                        [node_ebd, paddle.zeros(_shapes, dtype=node_ebd.dtype)],
+                        axis=1,
+                    )
+                    real_nloc = nloc
+                    real_nall = nall
+                else:
+                    # for spin
+                    real_nloc = nloc // 2
+                    real_nall = nall // 2
+                    real_n_padding = real_nall - real_nloc
+                    node_ebd_real, node_ebd_virtual = paddle.split(
+                        node_ebd, [real_nloc, real_nloc], axis=1
+                    )
+                    # mix_node_ebd: nb x real_nloc x (n_dim * 2)
+                    mix_node_ebd = paddle.concat(
+                        [node_ebd_real, node_ebd_virtual], axis=2
+                    )
+                    # nb x real_nall x (n_dim * 2)
+                    node_ebd = paddle.nn.functional.pad(
+                        mix_node_ebd.squeeze(0), (0, 0, 0, real_n_padding), value=0.0
+                    )
+
+                assert len(comm_dict) >= 6
+                # assert "send_list" in comm_dict
+                # assert "send_proc" in comm_dict
+                # assert "recv_proc" in comm_dict
+                # assert "send_num" in comm_dict
+                # assert "recv_num" in comm_dict
+                # assert "communicator" in comm_dict
+                ret = paddle_ops_deepmd_border_op(
+                    comm_dict[0],
+                    comm_dict[1],
+                    comm_dict[2],
+                    comm_dict[3],
+                    comm_dict[4],
+                    node_ebd,
+                    comm_dict[5],
+                    paddle.to_tensor(
+                        real_nloc,
+                        dtype=paddle.int32,
+                        place=paddle.CPUPlace(),
+                    ),  # should be int of c++, placed on cpu
+                    paddle.to_tensor(
+                        real_nall - real_nloc,
+                        dtype=paddle.int32,
+                        place=paddle.CPUPlace(),
+                    ),  # should be int of c++, placed on cpu
+                )
+                node_ebd_ext = ret[0].unsqueeze(0)
+                if has_spin:
+                    node_ebd_real_ext, node_ebd_virtual_ext = paddle.split(
+                        node_ebd_ext, [n_dim, n_dim], axis=2
+                    )
+                    node_ebd_ext = concat_switch_virtual(
+                        node_ebd_real_ext, node_ebd_virtual_ext, real_nloc
+                    )
             node_ebd, edge_ebd, angle_ebd = ll.forward(
                 node_ebd_ext,
                 edge_ebd,

deepmd/pd/model/descriptor/repformer_layer.py

@@ -110,6 +110,7 @@ def _make_nei_g1(
     # index: nb x (nloc x nnei) x ng1
     index = nlist.reshape([nb, nloc * nnei]).unsqueeze(-1).expand([-1, -1, ng1])
     # gg1  : nb x (nloc x nnei) x ng1
+    # print(g1_ext.shape, index.shape)
     gg1 = paddle.take_along_axis(g1_ext, indices=index, axis=1, broadcast=False)
     # gg1  : nb x nloc x nnei x ng1
     gg1 = gg1.reshape([nb, nloc, nnei, ng1])

deepmd/pd/model/descriptor/repformers.py

@@ -10,6 +10,10 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
+from deepmd.pd.cxx_op import (
+    ENABLE_CUSTOMIZED_OP,
+    paddle_ops_deepmd,
+)
 from deepmd.pd.model.descriptor.descriptor import (
     DescriptorBlock,
 )
@@ -31,6 +35,9 @@
 from deepmd.pd.utils.exclude_mask import (
     PairExcludeMask,
 )
+from deepmd.pd.utils.spin import (
+    concat_switch_virtual,
+)
 from deepmd.pd.utils.utils import (
     ActivationFn,
 )
@@ -45,6 +52,29 @@
     RepformerLayer,
 )
 
+if not ENABLE_CUSTOMIZED_OP:
+
+    def border_op(
+        argument0,
+        argument1,
+        argument2,
+        argument3,
+        argument4,
+        argument5,
+        argument6,
+        argument7,
+        argument8,
+    ) -> paddle.Tensor:
+        raise NotImplementedError(
+            "border_op is not available since customized Paddle OP library is not built when freezing the model. "
+            "See documentation for DPA3 for details."
+        )
+
+    # Note: this hack cannot actually save a model that can be run using LAMMPS.
+    paddle_ops_deepmd_border_op = border_op
+else:
+    paddle_ops_deepmd_border_op = paddle_ops_deepmd.border_op
+
 
 @DescriptorBlock.register("se_repformer")
 @DescriptorBlock.register("se_uni")
@@ -380,9 +410,10 @@ def forward(
         type_embedding: Optional[paddle.Tensor] = None,
         comm_dict: Optional[dict[str, paddle.Tensor]] = None,
     ):
-        if comm_dict is None:
-            assert mapping is not None
-            assert extended_atype_embd is not None
+        if comm_dict is None or len(comm_dict) == 0:
+            if paddle.in_dynamic_mode():
+                assert mapping is not None and mapping.numel() > 0
+            assert extended_atype_embd is not None or extended_atype_embd.numel() > 0
         nframes, nloc, nnei = nlist.shape
         nall = extended_coord.reshape([nframes, -1]).shape[1] // 3
         atype = extended_atype[:, :nloc]
@@ -406,7 +437,7 @@ def forward(
         sw = sw.masked_fill(~nlist_mask, 0.0)
 
         # [nframes, nloc, tebd_dim]
-        if comm_dict is None:
+        if comm_dict is None or len(comm_dict) == 0:
             if paddle.in_dynamic_mode():
                 assert isinstance(extended_atype_embd, paddle.Tensor)  # for jit
             atype_embd = extended_atype_embd[:, :nloc, :]
@@ -432,8 +463,9 @@ def forward(
         # if the a neighbor is real or not is indicated by nlist_mask
         nlist[nlist == -1] = 0
         # nb x nall x ng1
-        if comm_dict is None:
-            assert mapping is not None
+        if comm_dict is None or len(comm_dict) == 0:
+            if paddle.in_dynamic_mode():
+                assert mapping is not None and mapping.numel() > 0
             mapping = (
                 mapping.reshape([nframes, nall])
                 .unsqueeze(-1)
@@ -442,13 +474,80 @@ def forward(
         for idx, ll in enumerate(self.layers):
             # g1:     nb x nloc x ng1
             # g1_ext: nb x nall x ng1
-            if comm_dict is None:
-                assert mapping is not None
+            if comm_dict is None or len(comm_dict) == 0:
+                if paddle.in_dynamic_mode():
+                    assert mapping is not None and mapping.numel() > 0
                 g1_ext = paddle.take_along_axis(
                     g1, axis=1, indices=mapping, broadcast=False
                 )
             else:
-                raise NotImplementedError("Not implemented yet")
+                has_spin = len(comm_dict) >= 7
+                if not has_spin:
+                    n_padding = nall - nloc
+                    # g1 = paddle.nn.functional.pad(
+                    #     g1.squeeze(0), (0, 0, 0, n_padding), value=0.0
+                    # )
+                    _shapes = g1.shape[1:]
+                    _shapes[1] = n_padding
+                    g1 = paddle.concat(
+                        [g1.squeeze(0), paddle.zeros(_shapes, dtype=g1.dtype)],
+                        axis=1,
+                    )
+                    real_nloc = nloc
+                    real_nall = nall
+                else:
+                    # for spin
+                    real_nloc = nloc // 2
+                    real_nall = nall // 2
+                    real_n_padding = real_nall - real_nloc
+                    g1_real, g1_virtual = paddle.split(
+                        g1, [real_nloc, real_nloc], axis=1
+                    )
+                    # mix_g1: nb x real_nloc x (ng1 * 2)
+                    mix_g1 = paddle.concat([g1_real, g1_virtual], axis=2)
+                    # nb x real_nall x (ng1 * 2)
+                    g1 = paddle.nn.functional.pad(
+                        mix_g1.squeeze(0), (0, 0, 0, real_n_padding), value=0.0
+                    )
+
+                # assert "send_list" in comm_dict
+                # assert "send_proc" in comm_dict
+                # assert "recv_proc" in comm_dict
+                # assert "send_num" in comm_dict
+                # assert "recv_num" in comm_dict
+                # assert "communicator" in comm_dict
+                # print(f"g1.shape = ", g1.shape)
+                ret = paddle_ops_deepmd_border_op(
+                    comm_dict[0],
+                    comm_dict[1],
+                    comm_dict[2],
+                    comm_dict[3],
+                    comm_dict[4],
+                    g1,
+                    comm_dict[5],
+                    paddle.to_tensor(
+                        real_nloc,
+                        dtype=paddle.int32,
+                        place=paddle.CPUPlace(),
+                    ),  # should be int of c++, placed on cpu
+                    paddle.to_tensor(
+                        real_nall - real_nloc,
+                        dtype=paddle.int32,
+                        place=paddle.CPUPlace(),
+                    ),  # should be int of c++, placed on cpu
+                )
+                # print(f"ret.shape = {ret.shape}")
+                # print(f"ret[0].shape = ", ret[0].shape)
+                g1_ext = ret.unsqueeze(0)
+                # print(f"g1_ext.shape = ", g1_ext.shape)
+                # exit()
+                if has_spin:
+                    g1_real_ext, g1_virtual_ext = paddle.split(
+                        g1_ext, [ng1, ng1], dim=2
+                    )
+                    g1_ext = concat_switch_virtual(
+                        g1_real_ext, g1_virtual_ext, real_nloc
+                    )
 
             g1, g2, h2 = ll.forward(
                 g1_ext,