pd: support dpa3 with paddle backend

.pre-commit-config.yaml

@@ -65,13 +65,13 @@ repos:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|.+\.json$)
   # markdown, yaml, CSS, javascript
-  - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v4.0.0-alpha.8
-    hooks:
-      - id: prettier
-        types_or: [markdown, yaml, css]
-        # workflow files cannot be modified by pre-commit.ci
-        exclude: ^(source/3rdparty|\.github/workflows|\.clang-format)
+  # - repo: https://github.com/pre-commit/mirrors-prettier
+  #   rev: v4.0.0-alpha.8
+  #   hooks:
+  #     - id: prettier
+  #       types_or: [markdown, yaml, css]
+  #       # workflow files cannot be modified by pre-commit.ci
+  #       exclude: ^(source/3rdparty|\.github/workflows|\.clang-format)
   # Shell
   - repo: https://github.com/scop/pre-commit-shfmt
     rev: v3.11.0-1
@@ -83,25 +83,25 @@ repos:
     hooks:
       - id: cmake-format
       #- id: cmake-lint
-  - repo: https://github.com/njzjz/mirrors-bibtex-tidy
-    rev: v1.13.0
-    hooks:
-      - id: bibtex-tidy
-        args:
-          - --curly
-          - --numeric
-          - --align=13
-          - --blank-lines
-          # disable sort: the order of keys and fields has explict meanings
-          #- --sort=key
-          - --duplicates=key,doi,citation,abstract
-          - --merge=combine
-          #- --sort-fields
-          #- --strip-comments
-          - --trailing-commas
-          - --encode-urls
-          - --remove-empty-fields
-          - --wrap=80
+  # - repo: https://github.com/njzjz/mirrors-bibtex-tidy
+  #   rev: v1.13.0
+  #   hooks:
+  #     - id: bibtex-tidy
+  #       args:
+  #         - --curly
+  #         - --numeric
+  #         - --align=13
+  #         - --blank-lines
+  #         # disable sort: the order of keys and fields has explict meanings
+  #         #- --sort=key
+  #         - --duplicates=key,doi,citation,abstract
+  #         - --merge=combine
+  #         #- --sort-fields
+  #         #- --strip-comments
+  #         - --trailing-commas
+  #         - --encode-urls
+  #         - --remove-empty-fields
+  #         - --wrap=80
   # license header
   - repo: https://github.com/Lucas-C/pre-commit-hooks
     rev: v1.5.5

deepmd/pd/loss/__init__.py

@@ -1,12 +1,14 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from .ener import (
+    EnergyHessianStdLoss,
     EnergyStdLoss,
 )
 from .loss import (
     TaskLoss,
 )
 
 __all__ = [
+    "EnergyHessianStdLoss",
     "EnergyStdLoss",
     "TaskLoss",
 ]

deepmd/pd/loss/ener.py

@@ -56,7 +56,7 @@ def __init__(
         use_huber=False,
         huber_delta=0.01,
         **kwargs,
-    ):
+    ) -> None:
         r"""Construct a layer to compute loss on energy, force and virial.
 
         Parameters
@@ -287,9 +287,9 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
                         rmse_f.detach(), find_force
                     )
                 else:
-                    l1_force_loss = F.l1_loss(force_label, force_pred, reduction="none")
+                    l1_force_loss = F.l1_loss(force_label, force_pred, reduction="mean")
                     more_loss["mae_f"] = self.display_if_exist(
-                        l1_force_loss.mean().detach(), find_force
+                        l1_force_loss.detach(), find_force
                     )
                     l1_force_loss = l1_force_loss.sum(-1).mean(-1).sum()
                     loss += (pref_f * l1_force_loss).to(GLOBAL_PD_FLOAT_PRECISION)
@@ -324,20 +324,19 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
                 drdq_reshape = drdq.reshape(
                     [-1, natoms * 3, self.numb_generalized_coord]
                 )
+                gen_force_label = paddle.einsum(
+                    "bij,bi->bj", drdq_reshape, force_label_reshape_nframes
+                )
+                # gen_force_label = (
+                #     drdq_reshape * force_label_reshape_nframes.unsqueeze(-1)
+                # ).sum([-2])
 
-                # gen_force_label = paddle.einsum(
-                #     "bij,bi->bj", drdq_reshape, force_label_reshape_nframes
-                # )
-                gen_force_label = (
-                    drdq_reshape * force_label_reshape_nframes.unsqueeze(-1)
-                ).sum([-2])
-
-                # gen_force = paddle.einsum(
-                #     "bij,bi->bj", drdq_reshape, force_reshape_nframes
-                # )
-                gen_force = (drdq_reshape * force_reshape_nframes.unsqueeze(-1)).sum(
-                    [-2]
+                gen_force = paddle.einsum(
+                    "bij,bi->bj", drdq_reshape, force_reshape_nframes
                 )
+                # gen_force = (drdq_reshape * force_reshape_nframes.unsqueeze(-1)).sum(
+                #     [-2]
+                # )
 
                 diff_gen_force = gen_force_label - gen_force
                 l2_gen_force_loss = paddle.square(diff_gen_force).mean()
@@ -534,3 +533,75 @@ def deserialize(cls, data: dict) -> "TaskLoss":
         check_version_compatibility(data.pop("@version"), 2, 1)
         data.pop("@class")
         return cls(**data)
+
+
+class EnergyHessianStdLoss(EnergyStdLoss):
+    def __init__(
+        self,
+        start_pref_h=0.0,
+        limit_pref_h=0.0,
+        **kwargs,
+    ):
+        r"""Enable the layer to compute loss on hessian.
+
+        Parameters
+        ----------
+        start_pref_h : float
+            The prefactor of hessian loss at the start of the training.
+        limit_pref_h : float
+            The prefactor of hessian loss at the end of the training.
+        **kwargs
+            Other keyword arguments.
+        """
+        super().__init__(**kwargs)
+        self.has_h = (start_pref_h != 0.0 and limit_pref_h != 0.0) or self.inference
+
+        self.start_pref_h = start_pref_h
+        self.limit_pref_h = limit_pref_h
+
+    def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
+        model_pred, loss, more_loss = super().forward(
+            input_dict, model, label, natoms, learning_rate, mae=mae
+        )
+        coef = learning_rate / self.starter_learning_rate
+        pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * coef
+
+        if self.has_h and "hessian" in model_pred and "hessian" in label:
+            find_hessian = label.get("find_hessian", 0.0)
+            pref_h = pref_h * find_hessian
+            diff_h = label["hessian"].reshape(
+                [-1],
+            ) - model_pred["hessian"].reshape(
+                [-1],
+            )
+            l2_hessian_loss = paddle.mean(paddle.square(diff_h))
+            if not self.inference:
+                more_loss["l2_hessian_loss"] = self.display_if_exist(
+                    l2_hessian_loss.detach(), find_hessian
+                )
+            loss += pref_h * l2_hessian_loss
+            rmse_h = l2_hessian_loss.sqrt()
+            more_loss["rmse_h"] = self.display_if_exist(rmse_h.detach(), find_hessian)
+            if mae:
+                mae_h = paddle.mean(paddle.abs(diff_h))
+                more_loss["mae_h"] = self.display_if_exist(mae_h.detach(), find_hessian)
+
+        if not self.inference:
+            more_loss["rmse"] = paddle.sqrt(loss.detach())
+        return model_pred, loss, more_loss
+
+    @property
+    def label_requirement(self) -> list[DataRequirementItem]:
+        """Add hessian label requirement needed for this loss calculation."""
+        label_requirement = super().label_requirement
+        if self.has_h:
+            label_requirement.append(
+                DataRequirementItem(
+                    "hessian",
+                    ndof=1,  # 9=3*3 --> 3N*3N=ndof*natoms*natoms
+                    atomic=True,
+                    must=False,
+                    high_prec=False,
+                )
+            )
+        return label_requirement

deepmd/pd/model/descriptor/__init__.py

@@ -12,6 +12,9 @@
 from .dpa2 import (
     DescrptDPA2,
 )
+from .dpa3 import (
+    DescrptDPA3,
+)
 from .env_mat import (
     prod_env_mat,
 )
@@ -39,6 +42,7 @@
     "DescrptBlockSeTTebd",
     "DescrptDPA1",
     "DescrptDPA2",
+    "DescrptDPA3",
     "DescrptSeA",
     "DescrptSeAttenV2",
     "DescrptSeTTebd",

deepmd/pd/model/descriptor/descriptor.py

@@ -6,12 +6,16 @@
 )
 from typing import (
     Callable,
+    NoReturn,
     Optional,
     Union,
 )
 
 import paddle
 
+from deepmd.pd.model.network.network import (
+    TypeEmbedNet,
+)
 from deepmd.pd.utils import (
     env,
 )
@@ -99,7 +103,7 @@ def compute_input_stats(
         self,
         merged: Union[Callable[[], list[dict]], list[dict]],
         path: Optional[DPPath] = None,
-    ):
+    ) -> NoReturn:
         """
         Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.
 
@@ -122,7 +126,7 @@ def get_stats(self) -> dict[str, StatItem]:
         """Get the statistics of the descriptor."""
         raise NotImplementedError
 
-    def share_params(self, base_class, shared_level, resume=False):
+    def share_params(self, base_class, shared_level, resume=False) -> None:
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
         If not start from checkpoint (resume is False),
@@ -134,7 +138,10 @@ def share_params(self, base_class, shared_level, resume=False):
         if shared_level == 0:
             # link buffers
             if hasattr(self, "mean"):
-                if not resume:
+                if not resume and (
+                    not getattr(self, "set_stddev_constant", False)
+                    or not getattr(self, "set_davg_zero", False)
+                ):
                     # in case of change params during resume
                     base_env = EnvMatStatSe(base_class)
                     base_env.stats = base_class.stats
@@ -172,6 +179,7 @@ def forward(
         extended_atype: paddle.Tensor,
         extended_atype_embd: Optional[paddle.Tensor] = None,
         mapping: Optional[paddle.Tensor] = None,
+        type_embedding: Optional[paddle.Tensor] = None,
     ):
         """Calculate DescriptorBlock."""
         pass
@@ -185,7 +193,15 @@ def need_sorted_nlist_for_lower(self) -> bool:
         """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
 
 
-def extend_descrpt_stat(des, type_map, des_with_stat=None):
+def make_default_type_embedding(
+    ntypes,
+):
+    aux = {}
+    aux["tebd_dim"] = 8
+    return TypeEmbedNet(ntypes, aux["tebd_dim"]), aux
+
+
+def extend_descrpt_stat(des, type_map, des_with_stat=None) -> None:
     r"""
     Extend the statistics of a descriptor block with types from newly provided `type_map`.
 

deepmd/pd/model/descriptor/dpa1.py

@@ -25,9 +25,15 @@
     PRECISION_DICT,
     RESERVED_PRECISION_DICT,
 )
+from deepmd.pd.utils.tabulate import (
+    DPTabulate,
+)
 from deepmd.pd.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.pd.utils.utils import (
+    ActivationFn,
+)
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
 )
@@ -584,7 +590,59 @@ def enable_compression(
             The overflow check frequency
         """
         # do some checks before the mocel compression process
-        raise NotImplementedError("Model compression is not supported in paddle yet.")
+        if self.compress:
+            raise ValueError("Compression is already enabled.")
+        assert not self.se_atten.resnet_dt, (
+            "Model compression error: descriptor resnet_dt must be false!"
+        )
+        for tt in self.se_atten.exclude_types:
+            if (tt[0] not in range(self.se_atten.ntypes)) or (
+                tt[1] not in range(self.se_atten.ntypes)
+            ):
+                raise RuntimeError(
+                    "exclude types"
+                    + str(tt)
+                    + " must within the number of atomic types "
+                    + str(self.se_atten.ntypes)
+                    + "!"
+                )
+        if (
+            self.se_atten.ntypes * self.se_atten.ntypes
+            - len(self.se_atten.exclude_types)
+            == 0
+        ):
+            raise RuntimeError(
+                "Empty embedding-nets are not supported in model compression!"
+            )
+
+        if self.se_atten.attn_layer != 0:
+            raise RuntimeError("Cannot compress model when attention layer is not 0.")
+
+        if self.tebd_input_mode != "strip":
+            raise RuntimeError("Cannot compress model when tebd_input_mode == 'concat'")
+
+        data = self.serialize()
+        self.table = DPTabulate(
+            self,
+            data["neuron"],
+            data["type_one_side"],
+            data["exclude_types"],
+            ActivationFn(data["activation_function"]),
+        )
+        self.table_config = [
+            table_extrapolate,
+            table_stride_1,
+            table_stride_2,
+            check_frequency,
+        ]
+        self.lower, self.upper = self.table.build(
+            min_nbor_dist, table_extrapolate, table_stride_1, table_stride_2
+        )
+
+        self.se_atten.enable_compression(
+            self.table.data, self.table_config, self.lower, self.upper
+        )
+        self.compress = True
 
     def forward(
         self,