feat(pt): add compression support for se_e3_tebd (#4992)

add compression support for se_e3_tebd

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

- **New Features**
- TEBD-style descriptor support with optional runtime compression and a
Python-accessible tabulation operation that supports autograd.

- **Performance Improvements**
- Shared/global tabulation tables and shared embeddings to avoid
redundant table builds; improved device-aware tensor handling for
CPU/GPU.

- **Reliability**
- Validation and guards to prevent repeated compression enabling and to
check input/embedding suitability.

- **Tests**
- New unit and integration tests covering forward/backward CPU and GPU
paths.

- **Documentation**
  - Updated docs to indicate model compression is supported.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: OutisLi

deepmd/pt/model/descriptor/se_t_tebd.py

@@ -7,6 +7,7 @@
 )
 
 import torch
+import torch.nn as nn
 
 from deepmd.dpmodel.utils import EnvMat as DPEnvMat
 from deepmd.dpmodel.utils.seed import (
@@ -39,9 +40,15 @@
 from deepmd.pt.utils.exclude_mask import (
     PairExcludeMask,
 )
+from deepmd.pt.utils.tabulate import (
+    DPTabulate,
+)
 from deepmd.pt.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.pt.utils.utils import (
+    ActivationFn,
+)
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
 )
@@ -181,6 +188,7 @@ def __init__(
         self.tebd_input_mode = tebd_input_mode
         self.concat_output_tebd = concat_output_tebd
         self.trainable = trainable
+        self.compress = False
         # set trainable
         for param in self.parameters():
             param.requires_grad = trainable
@@ -516,6 +524,83 @@ def update_sel(
         local_jdata_cpy["sel"] = sel[0]
         return local_jdata_cpy, min_nbor_dist
 
+    def enable_compression(
+        self,
+        min_nbor_dist: float,
+        table_extrapolate: float = 5,
+        table_stride_1: float = 0.01,
+        table_stride_2: float = 0.1,
+        check_frequency: int = -1,
+    ) -> None:
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
+
+        Parameters
+        ----------
+        min_nbor_dist
+            The nearest distance between atoms
+        table_extrapolate
+            The scale of model extrapolation
+        table_stride_1
+            The uniform stride of the first table
+        table_stride_2
+            The uniform stride of the second table
+        check_frequency
+            The overflow check frequency
+        """
+        # do some checks before the model compression process
+        if self.compress:
+            raise ValueError("Compression is already enabled.")
+        assert not self.se_ttebd.resnet_dt, (
+            "Model compression error: descriptor resnet_dt must be false!"
+        )
+        for tt in self.se_ttebd.exclude_types:
+            if (tt[0] not in range(self.se_ttebd.ntypes)) or (
+                tt[1] not in range(self.se_ttebd.ntypes)
+            ):
+                raise RuntimeError(
+                    "exclude types"
+                    + str(tt)
+                    + " must within the number of atomic types "
+                    + str(self.se_ttebd.ntypes)
+                    + "!"
+                )
+        if (
+            self.se_ttebd.ntypes * self.se_ttebd.ntypes
+            - len(self.se_ttebd.exclude_types)
+            == 0
+        ):
+            raise RuntimeError(
+                "Empty embedding-nets are not supported in model compression!"
+            )
+
+        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"],
+            exclude_types=data["exclude_types"],
+            activation_fn=ActivationFn(data["activation_function"]),
+        )
+        # Scale the stride values for SE_T descriptor
+        stride_1_scaled = table_stride_1 * 10
+        stride_2_scaled = table_stride_2 * 10
+        self.table_config = [
+            table_extrapolate,
+            stride_1_scaled,
+            stride_2_scaled,
+            check_frequency,
+        ]
+        self.lower, self.upper = self.table.build(
+            min_nbor_dist, table_extrapolate, stride_1_scaled, stride_2_scaled
+        )
+
+        self.se_ttebd.enable_compression(
+            self.table.data, self.table_config, self.lower, self.upper
+        )
+        self.compress = True
+
 
 @DescriptorBlock.register("se_ttebd")
 class DescrptBlockSeTTebd(DescriptorBlock):
@@ -607,6 +692,14 @@ def __init__(
             )
             self.filter_layers_strip = filter_layers_strip
         self.stats = None
+        # compression related variables
+        self.compress = False
+        self.compress_info = nn.ParameterList(
+            [nn.Parameter(torch.zeros(0, dtype=self.prec, device="cpu"))]
+        )
+        self.compress_data = nn.ParameterList(
+            [nn.Parameter(torch.zeros(0, dtype=self.prec, device=env.DEVICE))]
+        )
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
@@ -811,6 +904,7 @@ def forward(
             self.rcut_smth,
             protection=self.env_protection,
         )
+        # dmatrix: [1/r, dx/r^2, dy/r^2, dz/r^2], sw: distance weighting
         # nb x nloc x nnei
         exclude_mask = self.emask(nlist, extended_atype)
         nlist = torch.where(exclude_mask != 0, nlist, -1)
@@ -831,11 +925,13 @@ def forward(
         rr = dmatrix
         rr = rr * exclude_mask[:, :, None]
 
-        # nfnl x nt_i x 3
+        # nfnl x nt_i x 3: direction vectors
+        # nt_i = nnei
+        # nt_j = nnei
         rr_i = rr[:, :, 1:]
         # nfnl x nt_j x 3
         rr_j = rr[:, :, 1:]
-        # nfnl x nt_i x nt_j
+        # nfnl x nt_i x nt_j: three-body angular correlations (cos theta_ij)
         env_ij = torch.einsum("ijm,ikm->ijk", rr_i, rr_j)
         # nfnl x nt_i x nt_j x 1
         ss = env_ij.unsqueeze(-1)
@@ -857,8 +953,24 @@ def forward(
             # nfnl x nt_i x nt_j x ng
             gg = self.filter_layers.networks[0](ss)
         elif self.tebd_input_mode in ["strip"]:
-            # nfnl x nt_i x nt_j x ng
-            gg_s = self.filter_layers.networks[0](ss)
+            if self.compress:
+                # Tabulated geometric embedding from angular features
+                # using SE_T_TEBD specific function
+                ebd_env_ij = env_ij.view(-1, 1)
+                gg_s = torch.ops.deepmd.tabulate_fusion_se_t_tebd(
+                    self.compress_data[0].contiguous(),
+                    self.compress_info[0].cpu().contiguous(),
+                    ebd_env_ij.contiguous(),
+                    env_ij.contiguous(),
+                    self.filter_neuron[-1],
+                )[0]
+                # SE_T_TEBD tabulation preserves the full neighbor structure
+                # nfnl x nt_i x nt_j x ng
+                gg_s = gg_s.view(nfnl, nnei, nnei, self.filter_neuron[-1])
+            else:
+                # nfnl x nt_i x nt_j x ng
+                gg_s = self.filter_layers.networks[0](ss)
+
             assert self.filter_layers_strip is not None
             assert type_embedding is not None
             ng = self.filter_neuron[-1]
@@ -902,16 +1014,19 @@ def forward(
             # (nfnl x nt_i x nt_j) x ng
             gg_t = gg_t.reshape(nfnl, nnei, nnei, ng)
             if self.smooth:
+                # Apply distance weighting to type features
                 gg_t = (
                     gg_t
                     * sw.reshape(nfnl, self.nnei, 1, 1)
                     * sw.reshape(nfnl, 1, self.nnei, 1)
                 )
+            # Combine geometric and type embeddings: gg_s * (1 + gg_t)
             # nfnl x nt_i x nt_j x ng
             gg = gg_s * gg_t + gg_s
         else:
             raise NotImplementedError
 
+        # Contract angular correlations with learned features
         # nfnl x ng
         res_ij = torch.einsum("ijk,ijkm->im", env_ij, gg)
         res_ij = res_ij * (1.0 / float(self.nnei) / float(self.nnei))
@@ -925,6 +1040,46 @@ def forward(
             sw,
         )
 
+    def enable_compression(
+        self,
+        table_data: dict,
+        table_config: dict,
+        lower: dict,
+        upper: dict,
+    ) -> None:
+        """Enable compression for the SE_T_TEBD descriptor block.
+
+        Parameters
+        ----------
+        table_data : dict
+            The tabulated data from DPTabulate
+        table_config : dict
+            Configuration for table compression
+        lower : dict
+            Lower bounds for compression
+        upper : dict
+            Upper bounds for compression
+        """
+        # Compress the main geometric embedding network (self.filter_layers)
+        net_key = "filter_net"
+        self.compress_info[0] = torch.as_tensor(
+            [
+                lower[net_key],
+                upper[net_key],
+                upper[net_key] * table_config[0],
+                table_config[1],
+                table_config[2],
+                table_config[3],
+            ],
+            dtype=self.prec,
+            device="cpu",
+        )
+        self.compress_data[0] = table_data[net_key].to(
+            device=env.DEVICE, dtype=self.prec
+        )
+
+        self.compress = True
+
     def has_message_passing(self) -> bool:
         """Returns whether the descriptor block has message passing."""
         return False

deepmd/pt/utils/tabulate.py

@@ -66,12 +66,7 @@ def __init__(
         )
         self.descrpt_type = self._get_descrpt_type()
 
-        supported_descrpt_type = (
-            "Atten",
-            "A",
-            "T",
-            "R",
-        )
+        supported_descrpt_type = ("Atten", "A", "T", "T_TEBD", "R")
 
         if self.descrpt_type in supported_descrpt_type:
             self.sel_a = self.descrpt.get_sel()
@@ -156,7 +151,7 @@ def _make_data(self, xx: np.ndarray, idx: int) -> Any:
                         self.matrix["layer_" + str(layer + 1)][idx],
                         xbar,
                         self.functype,
-                    ) + torch.ones((1, 1), dtype=yy.dtype)  # pylint: disable=no-explicit-device
+                    ) + torch.ones((1, 1), dtype=yy.dtype, device=yy.device)
                     dy2 = unaggregated_dy2_dx_s(
                         yy - xx,
                         dy,
@@ -175,7 +170,7 @@ def _make_data(self, xx: np.ndarray, idx: int) -> Any:
                         self.matrix["layer_" + str(layer + 1)][idx],
                         xbar,
                         self.functype,
-                    ) + torch.ones((1, 2), dtype=yy.dtype)  # pylint: disable=no-explicit-device
+                    ) + torch.ones((1, 2), dtype=yy.dtype, device=yy.device)
                     dy2 = unaggregated_dy2_dx_s(
                         yy - tt,
                         dy,
@@ -311,6 +306,8 @@ def _get_descrpt_type(self) -> str:
             return "R"
         elif isinstance(self.descrpt, deepmd.pt.model.descriptor.DescrptSeT):
             return "T"
+        elif isinstance(self.descrpt, deepmd.pt.model.descriptor.DescrptSeTTebd):
+            return "T_TEBD"
         raise RuntimeError(f"Unsupported descriptor {self.descrpt}")
 
     def _get_layer_size(self) -> int:
@@ -325,7 +322,7 @@ def _get_layer_size(self) -> int:
                 * len(self.embedding_net_nodes[0])
                 * len(self.neuron)
             )
-        if self.descrpt_type == "Atten":
+        if self.descrpt_type in ("Atten", "T_TEBD"):
             layer_size = len(self.embedding_net_nodes[0]["layers"])
         elif self.descrpt_type == "A":
             layer_size = len(self.embedding_net_nodes[0]["layers"])
@@ -394,6 +391,13 @@ def _get_network_variable(self, var_name: str) -> dict:
                             "layers"
                         ][layer - 1]["@variables"][var_name]
                         result["layer_" + str(layer)].append(node)
+            elif self.descrpt_type == "T_TEBD":
+                # For the se_e3_tebd descriptor, a single,
+                # shared embedding network is used for all type pairs
+                node = self.embedding_net_nodes[0]["layers"][layer - 1]["@variables"][
+                    var_name
+                ]
+                result["layer_" + str(layer)].append(node)
             elif self.descrpt_type == "R":
                 if self.type_one_side:
                     for ii in range(0, self.ntypes):