From 387708026d843ce568d01885f32f61d13d07bc81 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Fri, 29 May 2026 18:49:01 +0000
Subject: [PATCH 1/5] fix(pt): clone inference coords before enabling grad

Avoid enabling gradients in-place on LAMMPS-provided coordinate views.
Clone the extended coordinates before building force graphs and pass that
cloned tensor to derivative generation while keeping auxiliary metadata
private to model output conversion.

Closes #5165

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
---
 .../model/atomic_model/base_atomic_model.py   |  4 ++-
 .../pt/model/atomic_model/dp_atomic_model.py  |  4 ++-
 .../model/atomic_model/linear_atomic_model.py |  4 ++-
 .../atomic_model/pairtab_atomic_model.py      |  5 ++-
 deepmd/pt/model/model/make_model.py           |  5 +--
 source/tests/pt/model/test_dp_atomic_model.py | 28 +++++++++++++++
 source/tests/pt/model/test_dp_model.py        | 35 +++++++++++++++++++
 7 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py
index 8605db9359..4ab995004f 100644
--- a/deepmd/pt/model/atomic_model/base_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/base_atomic_model.py
@@ -383,7 +383,9 @@ def forward_common_atomic(
         if self.atom_excl is not None:
             atom_mask *= self.atom_excl(atype)
 
-        for kk in ret_dict.keys():
+        for kk in list(ret_dict.keys()):
+            if kk.startswith("_"):
+                continue
             out_shape = ret_dict[kk].shape
             out_shape2 = 1
             for ss in out_shape[2:]:
diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 783ee9e766..f7418948a8 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -272,7 +272,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
         if self.do_grad_r() or self.do_grad_c():
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.detach().clone().requires_grad_(True)
 
         # Handle default chg_spin if descriptor supports it
         if self.add_chg_spin_ebd and charge_spin is None:
@@ -302,6 +302,8 @@ def forward_atomic(
             fparam=fparam,
             aparam=aparam,
         )
+        if self.do_grad_r() or self.do_grad_c():
+            fit_ret["_force_coord"] = extended_coord
         if self.enable_eval_fitting_last_layer_hook:
             assert "middle_output" in fit_ret, (
                 "eval_fitting_last_layer not supported for this fitting net!"
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 5c0f616634..8d57a802d5 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -259,7 +259,7 @@ def forward_atomic(
         """
         nframes, nloc, nnei = nlist.shape
         if self.do_grad_r() or self.do_grad_c():
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.detach().clone().requires_grad_(True)
         extended_coord = extended_coord.view(nframes, -1, 3)
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
         nlists = build_multiple_neighbor_list(
@@ -304,6 +304,8 @@ def forward_atomic(
                 dim=0,
             ),
         }  # (nframes, nloc, 1)
+        if self.do_grad_r() or self.do_grad_c():
+            fit_ret["_force_coord"] = extended_coord
         return fit_ret
 
     def apply_out_stat(
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 5750f7cfd1..6b7249d9a3 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -276,7 +276,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         extended_coord = extended_coord.view(nframes, -1, 3)
         if self.do_grad_r() or self.do_grad_c():
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.detach().clone().requires_grad_(True)
 
         # this will mask all -1 in the nlist
         mask = nlist >= 0
@@ -313,6 +313,9 @@ def forward_atomic(
             dim=-1,
         ).unsqueeze(-1)
 
+        if self.do_grad_r() or self.do_grad_c():
+            atomic_energy = atomic_energy + 0.0 * extended_coord.sum()[..., None, None]
+            return {"energy": atomic_energy, "_force_coord": extended_coord}
         return {"energy": atomic_energy}
 
     def _pair_tabulated_inter(
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index 713eab3d8c..ea1df9bb02 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -305,7 +305,7 @@ def forward_common_lower(
             cc_ext, _, fp, ap, input_prec = self._input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
-            del extended_coord, fparam, aparam
+            del fparam, aparam
             atomic_ret = self.atomic_model.forward_common_atomic(
                 cc_ext,
                 extended_atype,
@@ -316,10 +316,11 @@ def forward_common_lower(
                 comm_dict=comm_dict,
                 charge_spin=charge_spin,
             )
+            force_coord = atomic_ret.pop("_force_coord", cc_ext)
             model_predict = fit_output_to_model_output(
                 atomic_ret,
                 self.atomic_output_def(),
-                cc_ext,
+                force_coord,
                 do_atomic_virial=do_atomic_virial,
                 create_graph=self.training,
                 mask=atomic_ret["mask"] if "mask" in atomic_ret else None,
diff --git a/source/tests/pt/model/test_dp_atomic_model.py b/source/tests/pt/model/test_dp_atomic_model.py
index 6d6a22f357..d77c986bb4 100644
--- a/source/tests/pt/model/test_dp_atomic_model.py
+++ b/source/tests/pt/model/test_dp_atomic_model.py
@@ -74,6 +74,34 @@ def test_self_consistency(self) -> None:
                 to_numpy_array(ret1["energy"]),
             )
 
+    def test_forward_common_atomic_accepts_leaf_view_input(self) -> None:
+        ds = DescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        ).to(env.DEVICE)
+        ft = InvarFitting(
+            "energy",
+            self.nt,
+            ds.get_dim_out(),
+            1,
+            mixed_types=ds.mixed_types(),
+        ).to(env.DEVICE)
+        md0 = DPAtomicModel(ds, ft, type_map=["foo", "bar"]).to(env.DEVICE)
+
+        coord = to_torch_tensor(self.coord_ext).requires_grad_(True)
+        coord_view = coord.view(self.nf, self.nall, 3)
+        args = [
+            coord_view,
+            to_torch_tensor(self.atype_ext),
+            to_torch_tensor(self.nlist),
+        ]
+        atomic_ret = md0.forward_atomic(*args)
+        ret = md0.forward_common_atomic(*args)
+
+        self.assertIn("_force_coord", atomic_ret)
+        self.assertIn("energy", ret)
+
     def test_dp_consistency(self) -> None:
         nf, nloc, nnei = self.nlist.shape
         ds = DPDescrptSeA(
diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py
index f4e350869a..a5adeada13 100644
--- a/source/tests/pt/model/test_dp_model.py
+++ b/source/tests/pt/model/test_dp_model.py
@@ -114,6 +114,41 @@ def test_self_consistency(self) -> None:
             atol=self.atol,
         )
 
+    def test_forward_lower_accepts_leaf_view_input(self) -> None:
+        ds = DescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        ).to(env.DEVICE)
+        ft = EnergyFittingNet(
+            self.nt,
+            ds.get_dim_out(),
+            mixed_types=ds.mixed_types(),
+        ).to(env.DEVICE)
+        type_map = ["foo", "bar"]
+        md0 = EnergyModel(ds, ft, type_map=type_map).to(env.DEVICE)
+
+        coord_ext, atype_ext, _ = extend_coord_with_ghosts(
+            to_torch_tensor(self.coord),
+            to_torch_tensor(self.atype),
+            to_torch_tensor(self.cell),
+            self.rcut,
+        )
+        nlist = build_neighbor_list(
+            coord_ext,
+            atype_ext,
+            self.nloc,
+            self.rcut,
+            self.sel,
+            distinguish_types=(not md0.mixed_types()),
+        )
+        coord_view = coord_ext.requires_grad_(True).view(self.nf, -1, 3)
+
+        ret = md0.forward_lower(coord_view, atype_ext, nlist, do_atomic_virial=True)
+
+        self.assertIn("extended_force", ret)
+        self.assertIn("virial", ret)
+
     def test_dp_consistency(self) -> None:
         nf, nloc = self.atype.shape
         nfp, nap = 2, 3

From a101808efb159b4d10d4e0bdb1c87a786686fcd2 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sat, 30 May 2026 02:50:43 +0000
Subject: [PATCH 2/5] fix(pt): preserve existing coordinate grad graph

Only clone extended coordinates when gradients are not already enabled.
This keeps existing autodiff and Hessian graphs intact while still avoiding
in-place requires_grad_ on non-grad coordinate views.

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
---
 deepmd/pt/model/atomic_model/dp_atomic_model.py      | 2 +-
 deepmd/pt/model/atomic_model/linear_atomic_model.py  | 2 +-
 deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index f7418948a8..9955a9e09b 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -271,7 +271,7 @@ def forward_atomic(
         """
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
-        if self.do_grad_r() or self.do_grad_c():
+        if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
             extended_coord = extended_coord.detach().clone().requires_grad_(True)
 
         # Handle default chg_spin if descriptor supports it
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 8d57a802d5..655dad6e9a 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -258,7 +258,7 @@ def forward_atomic(
             the result dict, defined by the fitting net output def.
         """
         nframes, nloc, nnei = nlist.shape
-        if self.do_grad_r() or self.do_grad_c():
+        if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
             extended_coord = extended_coord.detach().clone().requires_grad_(True)
         extended_coord = extended_coord.view(nframes, -1, 3)
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 6b7249d9a3..e1ffe59ea3 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -275,7 +275,7 @@ def forward_atomic(
     ) -> dict[str, torch.Tensor]:
         nframes, nloc, nnei = nlist.shape
         extended_coord = extended_coord.view(nframes, -1, 3)
-        if self.do_grad_r() or self.do_grad_c():
+        if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
             extended_coord = extended_coord.detach().clone().requires_grad_(True)
 
         # this will mask all -1 in the nlist

From d5270996b5e3267d975e212da259f81af8da0103 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sat, 30 May 2026 06:15:24 +0000
Subject: [PATCH 3/5] test(pt): cover grad enabling on coordinate views

---
 source/tests/pt/model/test_dp_atomic_model.py | 2 +-
 source/tests/pt/model/test_dp_model.py        | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/source/tests/pt/model/test_dp_atomic_model.py b/source/tests/pt/model/test_dp_atomic_model.py
index d77c986bb4..c30b937487 100644
--- a/source/tests/pt/model/test_dp_atomic_model.py
+++ b/source/tests/pt/model/test_dp_atomic_model.py
@@ -89,7 +89,7 @@ def test_forward_common_atomic_accepts_leaf_view_input(self) -> None:
         ).to(env.DEVICE)
         md0 = DPAtomicModel(ds, ft, type_map=["foo", "bar"]).to(env.DEVICE)
 
-        coord = to_torch_tensor(self.coord_ext).requires_grad_(True)
+        coord = to_torch_tensor(self.coord_ext)
         coord_view = coord.view(self.nf, self.nall, 3)
         args = [
             coord_view,
diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py
index a5adeada13..37f0fb9caa 100644
--- a/source/tests/pt/model/test_dp_model.py
+++ b/source/tests/pt/model/test_dp_model.py
@@ -142,10 +142,11 @@ def test_forward_lower_accepts_leaf_view_input(self) -> None:
             self.sel,
             distinguish_types=(not md0.mixed_types()),
         )
-        coord_view = coord_ext.requires_grad_(True).view(self.nf, -1, 3)
+        coord_view = coord_ext.view(self.nf, -1, 3)
 
         ret = md0.forward_lower(coord_view, atype_ext, nlist, do_atomic_virial=True)
 
+        self.assertFalse(coord_view.requires_grad)
         self.assertIn("extended_force", ret)
         self.assertIn("virial", ret)
 

From 889eff4b86694dcc1249a09a3b9f8ca933fe280b Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sat, 30 May 2026 10:03:53 +0000
Subject: [PATCH 4/5] fix(pt): prepare force coordinate in model layer

---
 deepmd/pt/model/atomic_model/base_atomic_model.py    | 4 +---
 deepmd/pt/model/atomic_model/dp_atomic_model.py      | 4 +---
 deepmd/pt/model/atomic_model/linear_atomic_model.py  | 4 +---
 deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 5 +----
 deepmd/pt/model/model/make_model.py                  | 9 ++++++---
 source/tests/pt/model/test_dp_atomic_model.py        | 2 --
 6 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py
index 4ab995004f..8605db9359 100644
--- a/deepmd/pt/model/atomic_model/base_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/base_atomic_model.py
@@ -383,9 +383,7 @@ def forward_common_atomic(
         if self.atom_excl is not None:
             atom_mask *= self.atom_excl(atype)
 
-        for kk in list(ret_dict.keys()):
-            if kk.startswith("_"):
-                continue
+        for kk in ret_dict.keys():
             out_shape = ret_dict[kk].shape
             out_shape2 = 1
             for ss in out_shape[2:]:
diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 9955a9e09b..25821fd945 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -272,7 +272,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord = extended_coord.detach().clone().requires_grad_(True)
+            extended_coord.requires_grad_(True)
 
         # Handle default chg_spin if descriptor supports it
         if self.add_chg_spin_ebd and charge_spin is None:
@@ -302,8 +302,6 @@ def forward_atomic(
             fparam=fparam,
             aparam=aparam,
         )
-        if self.do_grad_r() or self.do_grad_c():
-            fit_ret["_force_coord"] = extended_coord
         if self.enable_eval_fitting_last_layer_hook:
             assert "middle_output" in fit_ret, (
                 "eval_fitting_last_layer not supported for this fitting net!"
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 655dad6e9a..41fd49f096 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -259,7 +259,7 @@ def forward_atomic(
         """
         nframes, nloc, nnei = nlist.shape
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord = extended_coord.detach().clone().requires_grad_(True)
+            extended_coord.requires_grad_(True)
         extended_coord = extended_coord.view(nframes, -1, 3)
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
         nlists = build_multiple_neighbor_list(
@@ -304,8 +304,6 @@ def forward_atomic(
                 dim=0,
             ),
         }  # (nframes, nloc, 1)
-        if self.do_grad_r() or self.do_grad_c():
-            fit_ret["_force_coord"] = extended_coord
         return fit_ret
 
     def apply_out_stat(
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index e1ffe59ea3..9277725c23 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -276,7 +276,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         extended_coord = extended_coord.view(nframes, -1, 3)
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord = extended_coord.detach().clone().requires_grad_(True)
+            extended_coord.requires_grad_(True)
 
         # this will mask all -1 in the nlist
         mask = nlist >= 0
@@ -313,9 +313,6 @@ def forward_atomic(
             dim=-1,
         ).unsqueeze(-1)
 
-        if self.do_grad_r() or self.do_grad_c():
-            atomic_energy = atomic_energy + 0.0 * extended_coord.sum()[..., None, None]
-            return {"energy": atomic_energy, "_force_coord": extended_coord}
         return {"energy": atomic_energy}
 
     def _pair_tabulated_inter(
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index ea1df9bb02..78705b153c 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -305,9 +305,13 @@ def forward_common_lower(
             cc_ext, _, fp, ap, input_prec = self._input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
-            del fparam, aparam
+            del extended_coord, fparam, aparam
+            force_coord = cc_ext
+            if self.atomic_model.do_grad_r() or self.atomic_model.do_grad_c():
+                if not force_coord.requires_grad:
+                    force_coord = force_coord.clone().requires_grad_(True)
             atomic_ret = self.atomic_model.forward_common_atomic(
-                cc_ext,
+                force_coord,
                 extended_atype,
                 nlist,
                 mapping=mapping,
@@ -316,7 +320,6 @@ def forward_common_lower(
                 comm_dict=comm_dict,
                 charge_spin=charge_spin,
             )
-            force_coord = atomic_ret.pop("_force_coord", cc_ext)
             model_predict = fit_output_to_model_output(
                 atomic_ret,
                 self.atomic_output_def(),
diff --git a/source/tests/pt/model/test_dp_atomic_model.py b/source/tests/pt/model/test_dp_atomic_model.py
index c30b937487..ddc8f6d5ed 100644
--- a/source/tests/pt/model/test_dp_atomic_model.py
+++ b/source/tests/pt/model/test_dp_atomic_model.py
@@ -96,10 +96,8 @@ def test_forward_common_atomic_accepts_leaf_view_input(self) -> None:
             to_torch_tensor(self.atype_ext),
             to_torch_tensor(self.nlist),
         ]
-        atomic_ret = md0.forward_atomic(*args)
         ret = md0.forward_common_atomic(*args)
 
-        self.assertIn("_force_coord", atomic_ret)
         self.assertIn("energy", ret)
 
     def test_dp_consistency(self) -> None:

From 8e1670d9988a8ecaf1cd9aabc1b4322b800b6401 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sat, 30 May 2026 10:48:26 +0000
Subject: [PATCH 5/5] fix(pt): clone DP atomic coords before enabling grad

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

diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 25821fd945..d59d518cab 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -272,7 +272,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.clone().requires_grad_(True)
 
         # Handle default chg_spin if descriptor supports it
         if self.add_chg_spin_ebd and charge_spin is None: