pd: support dpa3 with paddle backend (#4701)

support dpa3 with paddle backend(eager mode)

### 1. training curve


![pt_vs_pd](https://github.com/user-attachments/assets/22f4681e-a464-41cb-9b1b-16ed20112563)

### 2. accuracy

<details>
  <summary>torch</summary>


![image](https://github.com/user-attachments/assets/c737ed30-0108-43f3-9d0e-7ae289db1498)

</details>


<details>
  <summary>paddle(slightly better than torch)</summary>


![image](https://github.com/user-attachments/assets/5e75abb8-f3be-46ce-adc1-eaf453cedcba)

</details>

### 3. The main modifications in this PR include:

1. Added DPA-3 code and related modules based on the Paddle backend.
2. Added the EnergyHessianStdLoss module based on the Paddle backend.
3. Discovered that Paddle’s ParameterList does not support assignment of
Tensors using the equals sign. Therefore, I added support for this
feature at <PaddlePaddle/Paddle#72190>. However,
considering version compatibility, deepmd still uses paddle.assign for
assignments.
4. Fixed an issue in env_mat_stat.py where the return type was Tensor
instead of float.
5. The SiLUT used APIs from the numpy series that do not support
paddle.Tensor, so I replaced them with Paddle’s native APIs.
Additionally, to temporarily bypass issues with dynamic-to-static
control flow, I changed the if-else branch in SiLUT.forward to a single
branch.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Introduced a new descriptor, DPA3, for advanced molecular simulations,
including its integration and public availability.
- Added support for a new graph-based neural network layer and
descriptor block for RepFlow calculations.
  - Enabled Hessian loss computation for enhanced training capabilities.
  - Added new learning rate utility.

- **Bug Fixes**
- Improved tensor shape handling and assignments for better
compatibility and stability.

- **Tests**
- Added comprehensive tests for the new DPA3 descriptor, including
consistency, JIT, and multitask scenarios.
- Expanded test coverage for model permutation and smoothness with DPA3.
  - Enhanced tests for DPA2 with CINN compiler support.

- **Refactor**
- Standardized tensor shape definitions and updated method signatures
for improved clarity and type safety.

- **Chores**
  - Updated public interfaces to include new features and descriptors.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>
Co-authored-by: Duo <50307526+iProzd@users.noreply.github.com>
Co-authored-by: root <2000011006@stu.pku.edu.cn>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Chenqqian Zhang <100290172+Chengqian-Zhang@users.noreply.github.com>
Co-authored-by: Jia-Xin Zhu <53895049+ChiahsinChu@users.noreply.github.com>
Co-authored-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: 8 people

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

@@ -584,7 +584,7 @@ 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.")
+        raise ValueError("Compression is already enabled.")
 
     def forward(
         self,

deepmd/pd/model/descriptor/dpa2.py

@@ -408,7 +408,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None:
         # shared_level: 1
         # share all parameters in type_embedding
         elif shared_level == 1:
-            self._modules["type_embedding"] = base_class._modules["type_embedding"]
+            self._sub_layers["type_embedding"] = base_class._sub_layers[
+                "type_embedding"
+            ]
         # Other shared levels
         else:
             raise NotImplementedError
@@ -899,4 +901,4 @@ def enable_compression(
             The overflow check frequency
         """
         # do some checks before the mocel compression process
-        raise NotImplementedError("enable_compression is not implemented yet")
+        raise ValueError("Compression is already enabled.")