From 0ee7dfd57ed22d60d84859fb2cbb104ecc76d9ef Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 17 Jul 2024 21:46:22 +0800
Subject: [PATCH 1/8] fix(2024Q1): fix lammps nlist sort with large sel

---
 deepmd/pt/model/model/ener_model.py |  2 ++
 deepmd/pt/model/model/make_model.py | 17 +++++++++++++----
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
index 5217293623..8edc0fa9c7 100644
--- a/deepmd/pt/model/model/ener_model.py
+++ b/deepmd/pt/model/model/ener_model.py
@@ -69,6 +69,7 @@ def forward_lower(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
+        extra_nlist_sort: bool = True,
     ):
         model_ret = self.forward_common_lower(
             extended_coord,
@@ -78,6 +79,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            extra_nlist_sort=extra_nlist_sort,
         )
         if self.get_fitting_net() is not None:
             model_predict = {}
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index 0e89c05b79..3ccae90993 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -216,6 +216,7 @@ def forward_common_lower(
             fparam: Optional[torch.Tensor] = None,
             aparam: Optional[torch.Tensor] = None,
             do_atomic_virial: bool = False,
+            extra_nlist_sort: bool = False,
         ):
             """Return model prediction. Lower interface that takes
             extended atomic coordinates and types, nlist, and mapping
@@ -238,6 +239,8 @@ def forward_common_lower(
                 atomic parameter. nf x nloc x nda
             do_atomic_virial
                 whether calculate atomic virial.
+            extra_nlist_sort
+                whether to forcibly sort the nlist.
 
             Returns
             -------
@@ -247,7 +250,7 @@ def forward_common_lower(
             """
             nframes, nall = extended_atype.shape[:2]
             extended_coord = extended_coord.view(nframes, -1, 3)
-            nlist = self.format_nlist(extended_coord, extended_atype, nlist)
+            nlist = self.format_nlist(extended_coord, extended_atype, nlist, extra_nlist_sort=extra_nlist_sort)
             cc_ext, _, fp, ap, input_prec = self.input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
@@ -347,6 +350,7 @@ def format_nlist(
             extended_coord: torch.Tensor,
             extended_atype: torch.Tensor,
             nlist: torch.Tensor,
+            extra_nlist_sort: bool = False,
         ):
             """Format the neighbor list.
 
@@ -372,6 +376,8 @@ def format_nlist(
                 atomic type in extended region. nf x nall
             nlist
                 neighbor list. nf x nloc x nsel
+            extra_nlist_sort
+                whether to forcibly sort the nlist.
 
             Returns
             -------
@@ -380,7 +386,7 @@ def format_nlist(
 
             """
             mixed_types = self.mixed_types()
-            nlist = self._format_nlist(extended_coord, nlist, sum(self.get_sel()))
+            nlist = self._format_nlist(extended_coord, nlist, sum(self.get_sel()), extra_nlist_sort=extra_nlist_sort)
             if not mixed_types:
                 nlist = nlist_distinguish_types(nlist, extended_atype, self.get_sel())
             return nlist
@@ -390,6 +396,7 @@ def _format_nlist(
             extended_coord: torch.Tensor,
             nlist: torch.Tensor,
             nnei: int,
+            extra_nlist_sort: bool=False,
         ):
             n_nf, n_nloc, n_nnei = nlist.shape
             # nf x nall x 3
@@ -409,7 +416,9 @@ def _format_nlist(
                     ],
                     dim=-1,
                 )
-            elif n_nnei > nnei:
+
+            if n_nnei > nnei or (extra_nlist_sort and n_nnei <= nnei):
+                n_nf, n_nloc, n_nnei = nlist.shape
                 m_real_nei = nlist >= 0
                 nlist = torch.where(m_real_nei, nlist, 0)
                 # nf x nloc x 3
@@ -426,7 +435,7 @@ def _format_nlist(
                 nlist = torch.gather(nlist, 2, nlist_mapping)
                 nlist = torch.where(rr > rcut, -1, nlist)
                 nlist = nlist[..., :nnei]
-            else:  # n_nnei == nnei:
+            else:  # not extra_nlist_sort and n_nnei <= nnei:
                 pass  # great!
             assert nlist.shape[-1] == nnei
             return nlist

From 90c0b3d5def8d7641e6cb40f5a1118105ef78281 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 18 Jul 2024 09:13:39 +0000
Subject: [PATCH 2/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/model/make_model.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index 3ccae90993..f051f9cccd 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -250,7 +250,9 @@ def forward_common_lower(
             """
             nframes, nall = extended_atype.shape[:2]
             extended_coord = extended_coord.view(nframes, -1, 3)
-            nlist = self.format_nlist(extended_coord, extended_atype, nlist, extra_nlist_sort=extra_nlist_sort)
+            nlist = self.format_nlist(
+                extended_coord, extended_atype, nlist, extra_nlist_sort=extra_nlist_sort
+            )
             cc_ext, _, fp, ap, input_prec = self.input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
@@ -386,7 +388,12 @@ def format_nlist(
 
             """
             mixed_types = self.mixed_types()
-            nlist = self._format_nlist(extended_coord, nlist, sum(self.get_sel()), extra_nlist_sort=extra_nlist_sort)
+            nlist = self._format_nlist(
+                extended_coord,
+                nlist,
+                sum(self.get_sel()),
+                extra_nlist_sort=extra_nlist_sort,
+            )
             if not mixed_types:
                 nlist = nlist_distinguish_types(nlist, extended_atype, self.get_sel())
             return nlist
@@ -396,7 +403,7 @@ def _format_nlist(
             extended_coord: torch.Tensor,
             nlist: torch.Tensor,
             nnei: int,
-            extra_nlist_sort: bool=False,
+            extra_nlist_sort: bool = False,
         ):
             n_nf, n_nloc, n_nnei = nlist.shape
             # nf x nall x 3

From 8d9470af74c95336d020f90dff4a85920486b132 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Thu, 18 Jul 2024 17:40:11 +0800
Subject: [PATCH 3/8] Update make_model.py

---
 deepmd/pt/model/model/make_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index f051f9cccd..cdbc3627a4 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -424,7 +424,7 @@ def _format_nlist(
                     dim=-1,
                 )
 
-            if n_nnei > nnei or (extra_nlist_sort and n_nnei <= nnei):
+            if n_nnei > nnei or extra_nlist_sort:
                 n_nf, n_nloc, n_nnei = nlist.shape
                 m_real_nei = nlist >= 0
                 nlist = torch.where(m_real_nei, nlist, 0)

From c52a0e02d549b4e3315db75f395ef6e9450a237e Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Mon, 29 Jul 2024 15:22:00 +0800
Subject: [PATCH 4/8] cherry-pick #3993

---
 .../dpmodel/atomic_model/dp_atomic_model.py   |  4 +++
 .../atomic_model/linear_atomic_model.py       |  4 +++
 .../atomic_model/make_base_atomic_model.py    |  4 +++
 .../atomic_model/pairtab_atomic_model.py      |  4 +++
 deepmd/dpmodel/descriptor/hybrid.py           |  4 +++
 .../descriptor/make_base_descriptor.py        |  4 +++
 deepmd/dpmodel/descriptor/se_e2_a.py          |  4 +++
 deepmd/dpmodel/descriptor/se_r.py             |  4 +++
 deepmd/dpmodel/model/make_model.py            | 33 +++++++++++++++----
 .../pt/model/atomic_model/dp_atomic_model.py  |  4 +++
 .../model/atomic_model/linear_atomic_model.py |  4 +++
 .../atomic_model/pairtab_atomic_model.py      |  4 +++
 deepmd/pt/model/descriptor/descriptor.py      |  4 +++
 deepmd/pt/model/descriptor/dpa1.py            |  4 +++
 deepmd/pt/model/descriptor/dpa2.py            |  4 +++
 deepmd/pt/model/descriptor/hybrid.py          |  4 +++
 deepmd/pt/model/descriptor/repformers.py      |  4 +++
 deepmd/pt/model/descriptor/se_a.py            |  9 +++++
 deepmd/pt/model/descriptor/se_atten.py        |  4 +++
 deepmd/pt/model/descriptor/se_r.py            |  4 +++
 deepmd/pt/model/model/dipole_model.py         |  1 +
 deepmd/pt/model/model/dos_model.py            |  1 +
 deepmd/pt/model/model/dp_zbl_model.py         |  1 +
 deepmd/pt/model/model/ener_model.py           |  3 +-
 deepmd/pt/model/model/frozen.py               |  4 +++
 deepmd/pt/model/model/make_model.py           |  4 +++
 deepmd/pt/model/model/polar_model.py          |  1 +
 deepmd/pt/model/model/spin_model.py           |  7 ++++
 28 files changed, 128 insertions(+), 8 deletions(-)

diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py
index 8a40f8d238..8d8bbf90c4 100644
--- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py
@@ -104,6 +104,10 @@ def get_out_bias(self) -> np.ndarray:
         """Return the output bias of the atomic model."""
         return self.fitting["bias_atom_e"]
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
+        return self.descriptor.need_sorted_nlist_for_lower()
+
     def forward_atomic(
         self,
         extended_coord: np.ndarray,
diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
index 93a885f3ab..ef2c629aad 100644
--- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py
@@ -111,6 +111,10 @@ def _sort_rcuts_sels(self) -> Tuple[List[float], List[int]]:
         )
         return [p[0] for p in zipped], [p[1] for p in zipped]
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
+        return True
+
     def forward_atomic(
         self,
         extended_coord,
diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
index 3e02a5d076..9daf38a82f 100644
--- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py
@@ -134,6 +134,10 @@ def mixed_types(self) -> bool:
             """
             pass
 
+        @abstractmethod
+        def need_sorted_nlist_for_lower(self) -> bool:
+            """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+
         @abstractmethod
         def fwd(
             self,
diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
index 30ab58928b..3660c97742 100644
--- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
@@ -177,6 +177,10 @@ def deserialize(cls, data) -> "PairTabAtomicModel":
         tab_model.tab_data = tab_model.tab.tab_data.reshape(ntypes, ntypes, nspline, 4)
         return tab_model
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
+        return False
+
     def forward_atomic(
         self,
         extended_coord,
diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py
index 96640d75c8..2219897661 100644
--- a/deepmd/dpmodel/descriptor/hybrid.py
+++ b/deepmd/dpmodel/descriptor/hybrid.py
@@ -140,6 +140,10 @@ def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None)
         for descrpt in self.descrpt_list:
             descrpt.compute_input_stats(merged, path)
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return True
+
     def call(
         self,
         coord_ext,
diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py
index 940bd0cd27..b98de7fa1f 100644
--- a/deepmd/dpmodel/descriptor/make_base_descriptor.py
+++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py
@@ -103,6 +103,10 @@ def compute_input_stats(
             """Update mean and stddev for descriptor elements."""
             raise NotImplementedError
 
+        @abstractmethod
+        def need_sorted_nlist_for_lower(self) -> bool:
+            """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+
         @abstractmethod
         def fwd(
             self,
diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py
index 531aa09f3a..2ba0bc9882 100644
--- a/deepmd/dpmodel/descriptor/se_e2_a.py
+++ b/deepmd/dpmodel/descriptor/se_e2_a.py
@@ -283,6 +283,10 @@ def reinit_exclude(
         self.exclude_types = exclude_types
         self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types)
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return False
+
     def call(
         self,
         coord_ext,
diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py
index 3128a28493..f27db3e975 100644
--- a/deepmd/dpmodel/descriptor/se_r.py
+++ b/deepmd/dpmodel/descriptor/se_r.py
@@ -232,6 +232,10 @@ def cal_g(
         gg = self.embeddings[(ll,)].call(ss)
         return gg
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return False
+
     def call(
         self,
         coord_ext,
diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py
index 68889ad331..43a91650c8 100644
--- a/deepmd/dpmodel/model/make_model.py
+++ b/deepmd/dpmodel/model/make_model.py
@@ -214,7 +214,12 @@ def call_lower(
             """
             nframes, nall = extended_atype.shape[:2]
             extended_coord = extended_coord.reshape(nframes, -1, 3)
-            nlist = self.format_nlist(extended_coord, extended_atype, nlist)
+            nlist = self.format_nlist(
+                extended_coord,
+                extended_atype,
+                nlist,
+                extra_nlist_sort=self.need_sorted_nlist_for_lower(),
+            )
             cc_ext, _, fp, ap, input_prec = self.input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
@@ -309,6 +314,7 @@ def format_nlist(
             extended_coord: np.ndarray,
             extended_atype: np.ndarray,
             nlist: np.ndarray,
+            extra_nlist_sort: bool = False,
         ):
             """Format the neighbor list.
 
@@ -334,6 +340,8 @@ def format_nlist(
                 atomic type in extended region. nf x nall
             nlist
                 neighbor list. nf x nloc x nsel
+            extra_nlist_sort
+                whether to forcibly sort the nlist.
 
             Returns
             -------
@@ -343,7 +351,12 @@ def format_nlist(
             """
             n_nf, n_nloc, n_nnei = nlist.shape
             mixed_types = self.mixed_types()
-            ret = self._format_nlist(extended_coord, nlist, sum(self.get_sel()))
+            ret = self._format_nlist(
+                extended_coord,
+                nlist,
+                sum(self.get_sel()),
+                extra_nlist_sort=extra_nlist_sort,
+            )
             if not mixed_types:
                 ret = nlist_distinguish_types(ret, extended_atype, self.get_sel())
             return ret
@@ -353,6 +366,7 @@ def _format_nlist(
             extended_coord: np.ndarray,
             nlist: np.ndarray,
             nnei: int,
+            extra_nlist_sort: bool = False,
         ):
             n_nf, n_nloc, n_nnei = nlist.shape
             extended_coord = extended_coord.reshape([n_nf, -1, 3])
@@ -368,8 +382,9 @@ def _format_nlist(
                     ],
                     axis=-1,
                 )
-            elif n_nnei > nnei:
-                # make a copy before revise
+
+            if n_nnei > nnei or extra_nlist_sort:
+                n_nf, n_nloc, n_nnei = nlist.shape
                 m_real_nei = nlist >= 0
                 ret = np.where(m_real_nei, nlist, 0)
                 coord0 = extended_coord[:, :n_nloc, :]
@@ -382,9 +397,11 @@ def _format_nlist(
                 ret = np.take_along_axis(ret, ret_mapping, axis=2)
                 ret = np.where(rr > rcut, -1, ret)
                 ret = ret[..., :nnei]
-            else:  # n_nnei == nnei:
-                # copy anyway...
+            # not extra_nlist_sort and n_nnei <= nnei:
+            elif n_nnei == nnei:
                 ret = nlist
+            else:
+                pass
             assert ret.shape[-1] == nnei
             return ret
 
@@ -469,6 +486,10 @@ def mixed_types(self) -> bool:
             """
             return self.atomic_model.mixed_types()
 
+        def need_sorted_nlist_for_lower(self) -> bool:
+            """Returns whether the model needs sorted nlist when using `forward_lower`."""
+            return self.atomic_model.need_sorted_nlist_for_lower()
+
         def atomic_output_def(self) -> FittingOutputDef:
             """Get the output def of the atomic model."""
             return self.atomic_model.atomic_output_def()
diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 13b8f09a79..2205bd5b05 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -100,6 +100,10 @@ def mixed_types(self) -> bool:
         """
         return self.descriptor.mixed_types()
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
+        return self.descriptor.need_sorted_nlist_for_lower()
+
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index f599399e66..2edf17a046 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -132,6 +132,10 @@ def _sort_rcuts_sels(self) -> Tuple[List[float], List[int]]:
         sorted_sels: List[int] = outer_sorted[:, 1].to(torch.int64).tolist()
         return sorted_rcuts, sorted_sels
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
+        return True
+
     def forward_atomic(
         self,
         extended_coord: torch.Tensor,
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index c20abf6a12..cff286e793 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -164,6 +164,10 @@ def mixed_types(self) -> bool:
         # to match DPA1 and DPA2.
         return True
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
+        return False
+
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(
diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py
index 5aae848aa4..68c06efa52 100644
--- a/deepmd/pt/model/descriptor/descriptor.py
+++ b/deepmd/pt/model/descriptor/descriptor.py
@@ -159,6 +159,10 @@ def forward(
         """Calculate DescriptorBlock."""
         pass
 
+    @abstractmethod
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
+
 
 def make_default_type_embedding(
     ntypes,
diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py
index 21275317dc..afdaa8ea4b 100644
--- a/deepmd/pt/model/descriptor/dpa1.py
+++ b/deepmd/pt/model/descriptor/dpa1.py
@@ -150,6 +150,10 @@ def mixed_types(self) -> bool:
         """
         return self.se_atten.mixed_types()
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return self.se_atten.need_sorted_nlist_for_lower()
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py
index fb792a51e2..4976a45aa0 100644
--- a/deepmd/pt/model/descriptor/dpa2.py
+++ b/deepmd/pt/model/descriptor/dpa2.py
@@ -307,6 +307,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return True
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py
index 204ca7589d..eb75466314 100644
--- a/deepmd/pt/model/descriptor/hybrid.py
+++ b/deepmd/pt/model/descriptor/hybrid.py
@@ -140,6 +140,10 @@ def mixed_types(self):
         """
         return any(descrpt.mixed_types() for descrpt in self.descrpt_list)
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return True
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py
index 16a38052b1..7a157def0d 100644
--- a/deepmd/pt/model/descriptor/repformers.py
+++ b/deepmd/pt/model/descriptor/repformers.py
@@ -205,6 +205,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
+        return False
+
     @property
     def dim_out(self):
         """Returns the output dimension of this descriptor."""
diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
index e17b7c5d54..48bf781cc4 100644
--- a/deepmd/pt/model/descriptor/se_a.py
+++ b/deepmd/pt/model/descriptor/se_a.py
@@ -132,6 +132,11 @@ def mixed_types(self):
         """
         return self.sea.mixed_types()
 
+
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return self.sea.need_sorted_nlist_for_lower()
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
@@ -433,6 +438,10 @@ def mixed_types(self) -> bool:
         """
         return False
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
+        return False
+
     @property
     def dim_out(self):
         """Returns the output dimension of this descriptor."""
diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py
index 051c66385c..ef4e32ddde 100644
--- a/deepmd/pt/model/descriptor/se_atten.py
+++ b/deepmd/pt/model/descriptor/se_atten.py
@@ -196,6 +196,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
+        return False
+
     @property
     def dim_out(self):
         """Returns the output dimension of this descriptor."""
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index ff922e0649..a06b08bf13 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -160,6 +160,10 @@ def mixed_types(self) -> bool:
         """
         return False
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return False
+
     def share_params(self, base_class, shared_level, resume=False):
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py
index 45b120771b..2e6eeef641 100644
--- a/deepmd/pt/model/model/dipole_model.py
+++ b/deepmd/pt/model/model/dipole_model.py
@@ -76,6 +76,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
         if self.get_fitting_net() is not None:
             model_predict = {}
diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py
index e043700bee..b921ee5b10 100644
--- a/deepmd/pt/model/model/dos_model.py
+++ b/deepmd/pt/model/model/dos_model.py
@@ -37,6 +37,7 @@ def forward(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
         if self.get_fitting_net() is not None:
             model_predict = {}
diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py
index bbc82b8d77..0f771a58f6 100644
--- a/deepmd/pt/model/model/dp_zbl_model.py
+++ b/deepmd/pt/model/model/dp_zbl_model.py
@@ -50,6 +50,7 @@ def forward(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
 
         model_predict = {}
diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
index 8edc0fa9c7..35acef0f1b 100644
--- a/deepmd/pt/model/model/ener_model.py
+++ b/deepmd/pt/model/model/ener_model.py
@@ -69,7 +69,6 @@ def forward_lower(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
-        extra_nlist_sort: bool = True,
     ):
         model_ret = self.forward_common_lower(
             extended_coord,
@@ -79,7 +78,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
-            extra_nlist_sort=extra_nlist_sort,
+            extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
         if self.get_fitting_net() is not None:
             model_predict = {}
diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py
index e3dcd389bb..6a6283b12e 100644
--- a/deepmd/pt/model/model/frozen.py
+++ b/deepmd/pt/model/model/frozen.py
@@ -102,6 +102,10 @@ def mixed_types(self) -> bool:
         """
         return self.model.mixed_types()
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the model needs sorted nlist when using `forward_lower`."""
+        return self.model.need_sorted_nlist_for_lower()
+
     @torch.jit.export
     def forward(
         self,
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index cdbc3627a4..cff10edfdf 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -548,4 +548,8 @@ def mixed_types(self) -> bool:
             """
             return self.atomic_model.mixed_types()
 
+        def need_sorted_nlist_for_lower(self) -> bool:
+            """Returns whether the model needs sorted nlist when using `forward_lower`."""
+            return self.atomic_model.need_sorted_nlist_for_lower()
+
     return CM
diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py
index 403058aa47..a7c97cdd39 100644
--- a/deepmd/pt/model/model/polar_model.py
+++ b/deepmd/pt/model/model/polar_model.py
@@ -68,6 +68,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
         if self.get_fitting_net() is not None:
             model_predict = {}
diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py
index df2f48e2e4..2fdef4fb8a 100644
--- a/deepmd/pt/model/model/spin_model.py
+++ b/deepmd/pt/model/model/spin_model.py
@@ -305,6 +305,10 @@ def has_spin(self) -> bool:
         """Returns whether it has spin input and output."""
         return True
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the model needs sorted nlist when using `forward_lower`."""
+        return self.backbone_model.need_sorted_nlist_for_lower()
+
     def __getattr__(self, name):
         """Get attribute from the wrapped model."""
         if (
@@ -418,6 +422,7 @@ def forward_common_lower(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
+        extra_nlist_sort: bool = False,
     ):
         nframes, nloc = nlist.shape[:2]
         (
@@ -436,6 +441,7 @@ def forward_common_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            extra_nlist_sort=extra_nlist_sort,
         )
         model_output_type = self.backbone_model.model_output_type()
         if "mask" in model_output_type:
@@ -546,6 +552,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            extra_nlist_sort=self.backbone_model.need_sorted_nlist_for_lower(),
         )
         model_predict = {}
         model_predict["atom_energy"] = model_ret["energy"]

From c3ce4bf05432b3ed1aaa5cb46e886eff038d9abb Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 29 Jul 2024 07:23:36 +0000
Subject: [PATCH 5/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/descriptor/se_a.py        |  1 -
 deepmd/pt/model/model/transform_output.py | 24 +++++++++++++++++++----
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
index 48bf781cc4..ef32c42a10 100644
--- a/deepmd/pt/model/descriptor/se_a.py
+++ b/deepmd/pt/model/descriptor/se_a.py
@@ -132,7 +132,6 @@ def mixed_types(self):
         """
         return self.sea.mixed_types()
 
-
     def need_sorted_nlist_for_lower(self) -> bool:
         """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
         return self.sea.need_sorted_nlist_for_lower()
diff --git a/deepmd/pt/model/model/transform_output.py b/deepmd/pt/model/model/transform_output.py
index e8a99fb5fc..204d656482 100644
--- a/deepmd/pt/model/model/transform_output.py
+++ b/deepmd/pt/model/model/transform_output.py
@@ -33,15 +33,27 @@ def atomic_virial_corr(
     faked_grad = torch.ones_like(sumce0)
     lst = torch.jit.annotate(List[Optional[torch.Tensor]], [faked_grad])
     extended_virial_corr0 = torch.autograd.grad(
-        [sumce0], [extended_coord], grad_outputs=lst, create_graph=False, retain_graph=True,
+        [sumce0],
+        [extended_coord],
+        grad_outputs=lst,
+        create_graph=False,
+        retain_graph=True,
     )[0]
     assert extended_virial_corr0 is not None
     extended_virial_corr1 = torch.autograd.grad(
-        [sumce1], [extended_coord], grad_outputs=lst, create_graph=False, retain_graph=True,
+        [sumce1],
+        [extended_coord],
+        grad_outputs=lst,
+        create_graph=False,
+        retain_graph=True,
     )[0]
     assert extended_virial_corr1 is not None
     extended_virial_corr2 = torch.autograd.grad(
-        [sumce2], [extended_coord], grad_outputs=lst, create_graph=False, retain_graph=True,
+        [sumce2],
+        [extended_coord],
+        grad_outputs=lst,
+        create_graph=False,
+        retain_graph=True,
     )[0]
     assert extended_virial_corr2 is not None
     extended_virial_corr = torch.concat(
@@ -66,7 +78,11 @@ def task_deriv_one(
     faked_grad = torch.ones_like(energy)
     lst = torch.jit.annotate(List[Optional[torch.Tensor]], [faked_grad])
     extended_force = torch.autograd.grad(
-        [energy], [extended_coord], grad_outputs=lst, create_graph=create_graph, retain_graph=True,
+        [energy],
+        [extended_coord],
+        grad_outputs=lst,
+        create_graph=create_graph,
+        retain_graph=True,
     )[0]
     assert extended_force is not None
     extended_force = -extended_force

From 333c41321a85b95aeb086523a49a6d43e6a1f443 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Tue, 30 Jul 2024 03:36:18 +0800
Subject: [PATCH 6/8] fix ut

---
 deepmd/pt/model/model/dos_model.py    | 2 +-
 deepmd/pt/model/model/dp_zbl_model.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py
index b921ee5b10..510b14d31b 100644
--- a/deepmd/pt/model/model/dos_model.py
+++ b/deepmd/pt/model/model/dos_model.py
@@ -37,7 +37,6 @@ def forward(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
-            extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
         if self.get_fitting_net() is not None:
             model_predict = {}
@@ -75,6 +74,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
         if self.get_fitting_net() is not None:
             model_predict = {}
diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py
index 0f771a58f6..d63327d1fd 100644
--- a/deepmd/pt/model/model/dp_zbl_model.py
+++ b/deepmd/pt/model/model/dp_zbl_model.py
@@ -50,7 +50,6 @@ def forward(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
-            extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
 
         model_predict = {}
@@ -87,6 +86,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            extra_nlist_sort=self.need_sorted_nlist_for_lower(),
         )
 
         model_predict = {}

From 0ade671fc4844f548fc9deea52f4b5fec4492109 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Tue, 30 Jul 2024 11:02:55 +0800
Subject: [PATCH 7/8] Update hybrid.py

---
 deepmd/pt/model/descriptor/hybrid.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py
index eb75466314..fb202cb3eb 100644
--- a/deepmd/pt/model/descriptor/hybrid.py
+++ b/deepmd/pt/model/descriptor/hybrid.py
@@ -373,6 +373,10 @@ def mixed_types(self) -> bool:
         """
         return all(descriptor.mixed_types() for descriptor in self.descriptor_list)
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
+        return True
+
     @property
     def dim_out(self):
         """Returns the output dimension of this descriptor."""

From 0e2e0f6c9560ecc3fea61964035d16c1aa03645b Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 31 Jul 2024 15:05:12 +0800
Subject: [PATCH 8/8] update from devel

---
 .github/workflows/build_wheel.yml       | 2 ++
 .github/workflows/test_cuda.yml         | 2 +-
 source/lmp/tests/run_mpi_pair_deepmd.py | 2 +-
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml
index 18fd7a1ac1..33ff8c6600 100644
--- a/.github/workflows/build_wheel.yml
+++ b/.github/workflows/build_wheel.yml
@@ -151,6 +151,8 @@ jobs:
           - variant: "_cu11"
             cuda_version: "11"
     steps:
+      - name: Delete huge unnecessary tools folder
+        run: rm -rf /opt/hostedtoolcache
       - uses: actions/checkout@v4
       - uses: actions/download-artifact@v4
         with:
diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml
index 15a20d889a..b024431d65 100644
--- a/.github/workflows/test_cuda.yml
+++ b/.github/workflows/test_cuda.yml
@@ -47,7 +47,7 @@ jobs:
          && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3
       if: false  # skip as we use nvidia image
     - run: python -m pip install -U "pip>=21.3.1,!=23.0.0"
-    - run: python -m pip install "tensorflow>=2.15.0rc0" "torch>=2.2.0"
+    - run: python -m pip install "tensorflow>=2.15.0rc0" "torch==2.3.1.*"
     - run: python -m pip install -v -e .[gpu,test,lmp,cu12,torch] mpi4py "ase @ https://gitlab.com/ase/ase/-/archive/8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f/ase-8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f.tar.gz"
       env:
         DP_VARIANT: cuda
diff --git a/source/lmp/tests/run_mpi_pair_deepmd.py b/source/lmp/tests/run_mpi_pair_deepmd.py
index b27774ce11..0c4291ab3a 100644
--- a/source/lmp/tests/run_mpi_pair_deepmd.py
+++ b/source/lmp/tests/run_mpi_pair_deepmd.py
@@ -54,8 +54,8 @@
 )
 lammps.pair_coeff("* *")
 lammps.run(0)
-pe = lammps.eval("pe")
 if rank == 0:
+    pe = lammps.eval("pe")
     arr = [pe]
     np.savetxt(output, np.array(arr))
 MPI.Finalize()