From ae453b010110c98eb1813deb819c1f0f8d6a4fa8 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Date: Tue, 30 Jun 2026 02:09:21 +0800
Subject: [PATCH 01/12] fix(jax): add hessian energy loss

---
 deepmd/dpmodel/loss/ener.py                   | 39 ++++++++++++++
 deepmd/jax/train/trainer.py                   |  6 +++
 source/tests/common/dpmodel/test_loss_ener.py | 54 ++++++++++++++++++-
 3 files changed, 97 insertions(+), 2 deletions(-)

diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index 7515f19b9a..0a8d57fc8c 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -71,6 +71,10 @@ class EnergyLoss(Loss):
         The prefactor of generalized force loss at the end of the training.
     numb_generalized_coord : int
         The dimension of generalized coordinates.
+    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.
     use_default_pf : bool
         If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided.
         This allows using the prefactor force loss (pf) without requiring atom_pref.npy files.
@@ -123,6 +127,8 @@ def __init__(
         start_pref_gf: float = 0.0,
         limit_pref_gf: float = 0.0,
         numb_generalized_coord: int = 0,
+        start_pref_h: float = 0.0,
+        limit_pref_h: float = 0.0,
         use_huber: bool = False,
         huber_delta: float | list[float] = 0.01,
         loss_func: str = "mse",
@@ -155,12 +161,15 @@ def __init__(
         self.start_pref_gf = start_pref_gf
         self.limit_pref_gf = limit_pref_gf
         self.numb_generalized_coord = numb_generalized_coord
+        self.start_pref_h = start_pref_h
+        self.limit_pref_h = limit_pref_h
         self.has_e = self.start_pref_e != 0.0 or self.limit_pref_e != 0.0
         self.has_f = self.start_pref_f != 0.0 or self.limit_pref_f != 0.0
         self.has_v = self.start_pref_v != 0.0 or self.limit_pref_v != 0.0
         self.has_ae = self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0
         self.has_pf = self.start_pref_pf != 0.0 or self.limit_pref_pf != 0.0
         self.has_gf = self.start_pref_gf != 0.0 or self.limit_pref_gf != 0.0
+        self.has_h = self.start_pref_h != 0.0 or self.limit_pref_h != 0.0
         if self.has_gf and self.numb_generalized_coord < 1:
             raise RuntimeError(
                 "When generalized force loss is used, the dimension of generalized coordinates should be larger than 0"
@@ -270,6 +279,7 @@ def call(
         pref_pf = find_atom_pref * (
             self.limit_pref_pf + (self.start_pref_pf - self.limit_pref_pf) * lr_ratio
         )
+        pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * lr_ratio
 
         loss = 0
         more_loss = {}
@@ -457,6 +467,23 @@ def call(
             more_loss["rmse_gf"] = self.display_if_exist(
                 xp.sqrt(l2_gen_force_loss), find_drdq
             )
+        hessian = model_dict.get(
+            "hessian", model_dict.get("energy_derv_r_derv_r", None)
+        )
+        if self.has_h and hessian is not None and "hessian" in label_dict:
+            find_hessian = label_dict.get("find_hessian", 0.0)
+            diff_h = xp.reshape(label_dict["hessian"], (-1,)) - xp.reshape(
+                hessian,
+                (-1,),
+            )
+            l2_hessian_loss = xp.mean(xp.square(diff_h))
+            loss += pref_h * find_hessian * l2_hessian_loss
+            more_loss["rmse_h"] = self.display_if_exist(
+                xp.sqrt(l2_hessian_loss), find_hessian
+            )
+            if mae:
+                mae_h = xp.mean(xp.abs(diff_h))
+                more_loss["mae_h"] = self.display_if_exist(mae_h, find_hessian)
 
         self.l2_l = loss
         more_loss["rmse"] = xp.sqrt(loss)
@@ -535,6 +562,16 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     default=1.0,
                 )
             )
+        if self.has_h:
+            label_requirement.append(
+                DataRequirementItem(
+                    "hessian",
+                    ndof=1,
+                    atomic=True,
+                    must=False,
+                    high_prec=False,
+                )
+            )
         return label_requirement
 
     def serialize(self) -> dict:
@@ -564,6 +601,8 @@ def serialize(self) -> dict:
             "start_pref_gf": self.start_pref_gf,
             "limit_pref_gf": self.limit_pref_gf,
             "numb_generalized_coord": self.numb_generalized_coord,
+            "start_pref_h": self.start_pref_h,
+            "limit_pref_h": self.limit_pref_h,
             "use_huber": self.use_huber,
             "huber_delta": self.huber_delta,
             "loss_func": self.loss_func,
diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py
index 180249eaef..46d8fc581e 100644
--- a/deepmd/jax/train/trainer.py
+++ b/deepmd/jax/train/trainer.py
@@ -117,6 +117,8 @@ def get_lr_and_coef(lr_param: dict) -> LearningRateExp:
         loss_type = loss_param.get("type", "ener")
         if loss_type == "ener":
             self.loss = EnergyLoss.get_loss(loss_param)
+            if getattr(self.loss, "has_h", False):
+                self.model.enable_hessian()
         else:
             raise RuntimeError("unknown loss type " + loss_type)
 
@@ -211,6 +213,8 @@ def loss_fn(
             model_dict["energy"] = model_dict["energy_redu"]
             model_dict["force"] = model_dict["energy_derv_r"].squeeze(-2)
             model_dict["virial"] = model_dict["energy_derv_c_redu"].squeeze(-2)
+            if model_dict.get("energy_derv_r_derv_r") is not None:
+                model_dict["hessian"] = model_dict["energy_derv_r_derv_r"].squeeze(-3)
             loss, more_loss = self.loss(
                 learning_rate=lr,
                 natoms=label_dict["type"].shape[1],
@@ -249,6 +253,8 @@ def loss_fn_more_loss(
             model_dict["energy"] = model_dict["energy_redu"]
             model_dict["force"] = model_dict["energy_derv_r"].squeeze(-2)
             model_dict["virial"] = model_dict["energy_derv_c_redu"].squeeze(-2)
+            if model_dict.get("energy_derv_r_derv_r") is not None:
+                model_dict["hessian"] = model_dict["energy_derv_r_derv_r"].squeeze(-3)
             loss, more_loss = self.loss(
                 learning_rate=lr,
                 natoms=label_dict["type"].shape[1],
diff --git a/source/tests/common/dpmodel/test_loss_ener.py b/source/tests/common/dpmodel/test_loss_ener.py
index ebf9ba0a64..cd6803a4ab 100644
--- a/source/tests/common/dpmodel/test_loss_ener.py
+++ b/source/tests/common/dpmodel/test_loss_ener.py
@@ -15,7 +15,14 @@
 class TestEnergyLossBase(unittest.TestCase):
     """Base class providing common setup for dpmodel EnergyLoss tests."""
 
-    def _make_data(self, natoms=5, nframes=2, numb_generalized_coord=0):
+    def _make_data(
+        self,
+        natoms=5,
+        nframes=2,
+        numb_generalized_coord=0,
+        hessian=False,
+        hessian_key="hessian",
+    ):
         """Generate fake model predictions and labels."""
         rng = np.random.default_rng(GLOBAL_SEED)
         model_dict = {
@@ -43,6 +50,10 @@ def _make_data(self, natoms=5, nframes=2, numb_generalized_coord=0):
             label_dict["find_drdq"] = 1.0
         if hasattr(self, "enable_atom_ener_coeff") and self.enable_atom_ener_coeff:
             label_dict["atom_ener_coeff"] = rng.random((nframes, natoms, 1))
+        if hessian:
+            model_dict[hessian_key] = rng.random((nframes, 3 * natoms, 3 * natoms))
+            label_dict["hessian"] = rng.random((nframes, 3 * natoms, 3 * natoms))
+            label_dict["find_hessian"] = 1.0
         return model_dict, label_dict, natoms
 
 
@@ -145,6 +156,38 @@ def test_forward(self) -> None:
         self.assertIsNotNone(loss)
 
 
+class TestEnergyLossHessian(TestEnergyLossBase):
+    """Test Hessian loss inside the dpmodel energy loss."""
+
+    def test_forward_hessian(self) -> None:
+        loss_fn = EnergyLoss(
+            starter_learning_rate=1.0,
+            start_pref_e=0.0,
+            limit_pref_e=0.0,
+            start_pref_f=0.0,
+            limit_pref_f=0.0,
+            start_pref_v=0.0,
+            limit_pref_v=0.0,
+            start_pref_h=2.0,
+            limit_pref_h=1.0,
+        )
+        model_dict, label_dict, natoms = self._make_data(
+            hessian=True,
+            hessian_key="energy_derv_r_derv_r",
+        )
+        loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
+        diff_h = label_dict["hessian"].reshape(-1) - model_dict[
+            "energy_derv_r_derv_r"
+        ].reshape(-1)
+        l2_hessian_loss = np.mean(np.square(diff_h))
+        np.testing.assert_allclose(loss, 2.0 * l2_hessian_loss)
+        np.testing.assert_allclose(more_loss["rmse_h"], np.sqrt(l2_hessian_loss))
+        self.assertIn(
+            "hessian",
+            {item.key for item in loss_fn.label_requirement},
+        )
+
+
 class TestEnergyLossSerialize(TestEnergyLossBase):
     """Test serialize/deserialize round-trip."""
 
@@ -160,10 +203,17 @@ def test_serialize_deserialize(self) -> None:
             start_pref_gf=1.0,
             limit_pref_gf=0.5,
             numb_generalized_coord=2,
+            start_pref_h=2.0,
+            limit_pref_h=0.5,
         )
         data = loss_fn.serialize()
+        self.assertEqual(data["start_pref_h"], 2.0)
+        self.assertEqual(data["limit_pref_h"], 0.5)
         loss_fn2 = EnergyLoss.deserialize(data)
-        model_dict, label_dict, natoms = self._make_data(numb_generalized_coord=2)
+        model_dict, label_dict, natoms = self._make_data(
+            numb_generalized_coord=2,
+            hessian=True,
+        )
         loss1, more1 = loss_fn.call(1.0, natoms, model_dict, label_dict)
         loss2, more2 = loss_fn2.call(1.0, natoms, model_dict, label_dict)
         np.testing.assert_allclose(loss1, loss2)

From fcfc2ece7482a70c5d3d68bf30c9afa93a3ca841 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Date: Tue, 7 Jul 2026 23:40:37 +0800
Subject: [PATCH 02/12] fix(data): declare hessian label shape

---
 deepmd/dpmodel/loss/ener.py                   |  3 +-
 deepmd/utils/data.py                          | 95 +++++++++++--------
 deepmd/utils/data_system.py                   |  5 +
 source/tests/common/dpmodel/test_loss_ener.py | 35 +++++++
 4 files changed, 97 insertions(+), 41 deletions(-)

diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index 0a8d57fc8c..996787a043 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -567,9 +567,10 @@ def label_requirement(self) -> list[DataRequirementItem]:
                 DataRequirementItem(
                     "hessian",
                     ndof=1,
-                    atomic=True,
+                    atomic=False,
                     must=False,
                     high_prec=False,
+                    special_shape="hessian",
                 )
             )
         return label_requirement
diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py
index 38db47809d..ddc29ce8cf 100644
--- a/deepmd/utils/data.py
+++ b/deepmd/utils/data.py
@@ -159,6 +159,7 @@ def add(
         default: float = 0.0,
         dtype: np.dtype | None = None,
         output_natoms_for_type_sel: bool = False,
+        special_shape: str | None = None,
     ) -> "DeepmdData":
         """Add a data item that to be loaded.
 
@@ -187,6 +188,9 @@ def add(
             the dtype of data, overwrites `high_prec` if provided
         output_natoms_for_type_sel : bool, optional
             if True and type_sel is True, the atomic dimension will be natoms instead of nsel
+        special_shape : str, optional
+            Name of a loader-defined non-standard shape contract. ``"hessian"``
+            stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame.
         """
         # normalize key: "atomic_" prefix -> "atom_", same convention as _load_set output
         if key.startswith("atomic_"):
@@ -202,6 +206,7 @@ def add(
             "default": default,
             "dtype": dtype,
             "output_natoms_for_type_sel": output_natoms_for_type_sel,
+            "special_shape": special_shape,
         }
         return self
 
@@ -732,6 +737,7 @@ def _load_set(self, set_name: DPPath) -> dict[str, Any]:
                     output_natoms_for_type_sel=self.data_dict[kk][
                         "output_natoms_for_type_sel"
                     ],
+                    special_shape=self.data_dict[kk].get("special_shape"),
                 )
         for kk in self.data_dict.keys():
             if self.data_dict[kk]["reduce"] is not None:
@@ -808,7 +814,9 @@ def _load_data(
         default: float = 0.0,
         dtype: np.dtype | None = None,
         output_natoms_for_type_sel: bool = False,
+        special_shape: str | None = None,
     ) -> np.ndarray:
+        is_hessian = special_shape == "hessian" or key == "hessian"
         if atomic:
             natoms = self.natoms
             idx_map = self.idx_map
@@ -836,7 +844,19 @@ def _load_data(
         if path.is_file():
             data = path.load_numpy().astype(dtype)
             try:  # YWolfeee: deal with data shape error
-                if atomic:
+                if is_hessian:
+                    natoms = self.natoms
+                    idx_map = self.idx_map
+                    data = data.reshape(nframes, 3 * natoms, 3 * natoms)
+                    num_chunks, chunk_size = len(idx_map), 3
+                    idx_map_hess = np.arange(num_chunks * chunk_size)  # pylint: disable=no-explicit-dtype
+                    idx_map_hess = idx_map_hess.reshape(num_chunks, chunk_size)
+                    idx_map_hess = idx_map_hess[idx_map].flatten()
+                    data = data[:, idx_map_hess, :]
+                    data = data[:, :, idx_map_hess]
+                    data = data.reshape([nframes, -1])
+                    ndof = 9 * natoms * natoms
+                elif atomic:
                     if type_sel is not None:
                         # check the data shape is nsel or natoms
                         if data.size == nframes * natoms_sel * ndof_:
@@ -868,24 +888,9 @@ def _load_data(
                                 f"({nframes}, {natoms_sel}, {ndof_}) or"
                                 f"({nframes}, {natoms}, {ndof_})"
                             )
-                    if key == "hessian":
-                        data = data.reshape(nframes, 3 * natoms, 3 * natoms)
-                        # get idx_map for hessian
-                        num_chunks, chunk_size = len(idx_map), 3
-                        idx_map_hess = np.arange(num_chunks * chunk_size)  # pylint: disable=no-explicit-dtype
-                        idx_map_hess = idx_map_hess.reshape(num_chunks, chunk_size)
-                        idx_map_hess = idx_map_hess[idx_map]
-                        idx_map_hess = idx_map_hess.flatten()
-                        data = data[:, idx_map_hess, :]
-                        data = data[:, :, idx_map_hess]
-                        data = data.reshape([nframes, -1])
-                        ndof = (
-                            3 * ndof * 3 * ndof
-                        )  # size of hessian is 3Natoms * 3Natoms
-                    else:
-                        data = data.reshape([nframes, natoms, -1])
-                        data = data[:, idx_map, :]
-                        data = data.reshape([nframes, -1])
+                    data = data.reshape([nframes, natoms, -1])
+                    data = data[:, idx_map, :]
+                    data = data.reshape([nframes, -1])
                 data = np.reshape(data, [nframes, ndof])
             except ValueError as err_message:
                 explanation = "This error may occur when your label mismatch its name, i.e. you might store global tensor in `atomic_tensor.npy` or atomic tensor in `tensor.npy`."
@@ -898,7 +903,9 @@ def _load_data(
         elif must:
             raise RuntimeError(f"{path} not found!")
         else:
-            if atomic and type_sel is not None and not output_natoms_for_type_sel:
+            if is_hessian:
+                ndof = 9 * self.natoms * self.natoms
+            elif atomic and type_sel is not None and not output_natoms_for_type_sel:
                 ndof = ndof_ * natoms_sel
             data = np.full([nframes, ndof], default, dtype=dtype)
             if repeat != 1:
@@ -925,8 +932,12 @@ def _load_single_data(
         """
         vv = self.data_dict[key]
         path = self._get_data_path(set_dir, key)
+        is_hessian = vv.get("special_shape") == "hessian" or key == "hessian"
 
-        if vv["atomic"]:
+        if is_hessian:
+            natoms = self.natoms
+            idx_map = self.idx_map
+        elif vv["atomic"]:
             natoms = self.natoms
             idx_map = self.idx_map
             # if type_sel, then revise natoms and idx_map
@@ -959,7 +970,9 @@ def _load_single_data(
                 raise RuntimeError(f"{path} not found!")
 
             # Create a default array based on requirements
-            if vv["atomic"]:
+            if is_hessian:
+                data = np.full([9 * natoms * natoms], vv["default"], dtype=dtype)
+            elif vv["atomic"]:
                 if vv["type_sel"] is not None and not vv["output_natoms_for_type_sel"]:
                     natoms = natoms_sel
                 data = np.full([natoms, ndof], vv["default"], dtype=dtype)
@@ -983,7 +996,17 @@ def _load_single_data(
         data = mmap_obj[frame_idx].copy().astype(dtype, copy=False)
 
         try:
-            if vv["atomic"]:
+            if is_hessian:
+                data = data.reshape(3 * natoms, 3 * natoms)
+                num_chunks, chunk_size = len(idx_map), 3
+                idx_map_hess = np.arange(num_chunks * chunk_size, dtype=int).reshape(
+                    num_chunks, chunk_size
+                )
+                idx_map_hess = idx_map_hess[idx_map].flatten()
+                data = data[idx_map_hess, :]
+                data = data[:, idx_map_hess]
+                data = data.reshape(-1)
+            elif vv["atomic"]:
                 # Handle type_sel logic
                 if vv["type_sel"] is not None:
                     if mmap_obj.shape[1] == natoms_sel * ndof:
@@ -1008,23 +1031,9 @@ def _load_single_data(
                             f"The shape of the data {key} in {set_dir} has width {mmap_obj.shape[1]}, which doesn't match either ({natoms_sel * ndof}) or ({natoms * ndof})"
                         )
 
-                # Handle special case for Hessian
-                if key == "hessian":
-                    data = data.reshape(3 * natoms, 3 * natoms)
-                    num_chunks, chunk_size = len(idx_map), 3
-                    idx_map_hess = np.arange(
-                        num_chunks * chunk_size, dtype=int
-                    ).reshape(num_chunks, chunk_size)
-                    idx_map_hess = idx_map_hess[idx_map].flatten()
-                    data = data[idx_map_hess, :]
-                    data = data[:, idx_map_hess]
-                    data = data.reshape(-1)
-                    # size of hessian is 3Natoms * 3Natoms
-                    # ndof = 3 * ndof * 3 * ndof
-                else:
-                    # data should be 2D here: [natoms, ndof]
-                    data = data.reshape([natoms, -1])
-                    data = data[idx_map, :]
+                # data should be 2D here: [natoms, ndof]
+                data = data.reshape([natoms, -1])
+                data = data[idx_map, :]
             else:
                 data = data.reshape([ndof])
 
@@ -1137,6 +1146,9 @@ class DataRequirementItem:
         the dtype of data, overwrites `high_prec` if provided
     output_natoms_for_type_sel : bool, optional
         if True and type_sel is True, the atomic dimension will be natoms instead of nsel
+    special_shape : str, optional
+        Name of a loader-defined non-standard shape contract. ``"hessian"``
+        stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame.
     """
 
     def __init__(
@@ -1151,6 +1163,7 @@ def __init__(
         default: float = 0.0,
         dtype: np.dtype | None = None,
         output_natoms_for_type_sel: bool = False,
+        special_shape: str | None = None,
     ) -> None:
         self.key = key
         self.ndof = ndof
@@ -1162,6 +1175,7 @@ def __init__(
         self.default = default
         self.dtype = dtype
         self.output_natoms_for_type_sel = output_natoms_for_type_sel
+        self.special_shape = special_shape
         self.dict = self.to_dict()
 
     def to_dict(self) -> dict:
@@ -1176,6 +1190,7 @@ def to_dict(self) -> dict:
             "default": self.default,
             "dtype": self.dtype,
             "output_natoms_for_type_sel": self.output_natoms_for_type_sel,
+            "special_shape": self.special_shape,
         }
 
     def __getitem__(self, key: str) -> np.ndarray:
diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py
index 9d13cb4699..f020c5cf38 100644
--- a/deepmd/utils/data_system.py
+++ b/deepmd/utils/data_system.py
@@ -333,6 +333,7 @@ def add_dict(self, adict: dict[str, dict[str, Any]]) -> None:
                 output_natoms_for_type_sel=adict[kk].get(
                     "output_natoms_for_type_sel", False
                 ),
+                special_shape=adict[kk].get("special_shape"),
             )
 
     def add_data_requirements(
@@ -353,6 +354,7 @@ def add(
         default: float = 0.0,
         dtype: np.dtype | None = None,
         output_natoms_for_type_sel: bool = False,
+        special_shape: str | None = None,
     ) -> None:
         """Add a data item that to be loaded.
 
@@ -381,6 +383,8 @@ def add(
             The dtype of data, overwrites `high_prec` if provided
         output_natoms_for_type_sel : bool
             If True and type_sel is True, the atomic dimension will be natoms instead of nsel
+        special_shape : str, optional
+            Name of a loader-defined non-standard shape contract.
         """
         for ii in self.data_systems:
             ii.add(
@@ -394,6 +398,7 @@ def add(
                 default=default,
                 dtype=dtype,
                 output_natoms_for_type_sel=output_natoms_for_type_sel,
+                special_shape=special_shape,
             )
 
     def reduce(self, key_out: str, key_in: str) -> None:
diff --git a/source/tests/common/dpmodel/test_loss_ener.py b/source/tests/common/dpmodel/test_loss_ener.py
index cd6803a4ab..154688e5b0 100644
--- a/source/tests/common/dpmodel/test_loss_ener.py
+++ b/source/tests/common/dpmodel/test_loss_ener.py
@@ -1,11 +1,17 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import unittest
+from pathlib import (
+    Path,
+)
 
 import numpy as np
 
 from deepmd.dpmodel.loss.ener import (
     EnergyLoss,
 )
+from deepmd.utils.data import (
+    DeepmdData,
+)
 
 from ...seed import (
     GLOBAL_SEED,
@@ -187,6 +193,35 @@ def test_forward_hessian(self) -> None:
             {item.key for item in loss_fn.label_requirement},
         )
 
+    def test_hessian_label_requirement_loads_full_matrix(self) -> None:
+        loss_fn = EnergyLoss(
+            starter_learning_rate=1.0,
+            start_pref_h=1.0,
+            limit_pref_h=1.0,
+        )
+        hessian_req = next(
+            item for item in loss_fn.label_requirement if item.key == "hessian"
+        )
+        self.assertFalse(hessian_req.atomic)
+        self.assertEqual(hessian_req.special_shape, "hessian")
+
+        system = (
+            Path(__file__).resolve().parents[2] / "pt" / "hessian" / "data" / "H8C4N2O"
+        )
+        data = DeepmdData(str(system), type_map=["C", "H", "N", "O"])
+        data.add(
+            hessian_req.key,
+            hessian_req.ndof,
+            atomic=hessian_req.atomic,
+            must=hessian_req.must,
+            high_prec=hessian_req.high_prec,
+            special_shape=hessian_req.special_shape,
+        )
+
+        batch = data.get_batch(2)
+
+        self.assertEqual(batch["hessian"].shape, (2, (3 * data.natoms) ** 2))
+
 
 class TestEnergyLossSerialize(TestEnergyLossBase):
     """Test serialize/deserialize round-trip."""

From a1d8f5694b2e0a825adef8ec583453f48d83e1c5 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Tue, 21 Jul 2026 11:54:59 +0800
Subject: [PATCH 03/12] fix(dpmodel): mask Hessian padding in energy loss

Exclude Hessian rows and columns associated with placeholder atoms from the training MSE and the MAE/RMSE denominators while preserving the unmasked reduction.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 deepmd/dpmodel/loss/ener.py                   |  25 +++-
 deepmd/dpmodel/loss/reduction.py              |  39 +++++
 .../tests/common/dpmodel/test_loss_padding.py | 134 ++++++++++++++++++
 3 files changed, 192 insertions(+), 6 deletions(-)

diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index 9d69946ad0..79c43a1646 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -13,6 +13,7 @@
 )
 from deepmd.dpmodel.loss.reduction import (
     masked_atom_mean,
+    masked_pair_mean,
     per_frame_component_mean,
 )
 from deepmd.utils.data import (
@@ -754,17 +755,29 @@ def call(
         hessian = model_dict.get("hessian", model_dict.get("energy_derv_r_derv_r"))
         if self.has_h and hessian is not None and "hessian" in label_dict:
             find_hessian = label_dict.get("find_hessian", 0.0)
-            diff_h = xp.reshape(label_dict["hessian"], (-1,)) - xp.reshape(
-                hessian,
-                (-1,),
-            )
-            l2_hessian_loss = xp.mean(xp.square(diff_h))
+            if maskf is not None:
+                hessian_shape = (_nf, _nloc * 3, _nloc * 3)
+                diff_h = xp.reshape(label_dict["hessian"], hessian_shape) - xp.reshape(
+                    hessian, hessian_shape
+                )
+                # A Hessian element couples two Cartesian atom components, so
+                # it is valid only when both corresponding atoms are real.
+                l2_hessian_loss = masked_pair_mean(xp.square(diff_h), maskf, ncomp=3)
+            else:
+                diff_h = xp.reshape(label_dict["hessian"], (-1,)) - xp.reshape(
+                    hessian,
+                    (-1,),
+                )
+                l2_hessian_loss = xp.mean(xp.square(diff_h))
             loss += pref_h * find_hessian * l2_hessian_loss
             more_loss["rmse_h"] = self.display_if_exist(
                 xp.sqrt(l2_hessian_loss), find_hessian
             )
             if mae:
-                mae_h = xp.mean(xp.abs(diff_h))
+                if maskf is not None:
+                    mae_h = masked_pair_mean(xp.abs(diff_h), maskf, ncomp=3)
+                else:
+                    mae_h = xp.mean(xp.abs(diff_h))
                 more_loss["mae_h"] = self.display_if_exist(mae_h, find_hessian)
 
         self.l2_l = loss
diff --git a/deepmd/dpmodel/loss/reduction.py b/deepmd/dpmodel/loss/reduction.py
index 068c11e2bc..2568913538 100644
--- a/deepmd/dpmodel/loss/reduction.py
+++ b/deepmd/dpmodel/loss/reduction.py
@@ -67,6 +67,45 @@ def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array:
     return xp.mean(per_frame)
 
 
+def masked_pair_mean(elem: Array, maskf: Array, ncomp: int) -> Array:
+    """Return a per-frame mean over valid atom-pair components.
+
+    Parameters
+    ----------
+    elem : Array
+        Non-negative pair contribution of shape
+        ``[nf, nloc * ncomp, nloc * ncomp]``. The contribution has already
+        been squared or converted to an absolute value, but is not masked.
+    maskf : Array
+        Per-atom real/placeholder mask of shape ``[nf, nloc]``.
+    ncomp : int
+        Number of components per atom on each pair axis. A Cartesian Hessian
+        uses three components on both axes.
+
+    Returns
+    -------
+    Array
+        ``mean_over_frames(sum(valid_pair_elem) / (real_natoms*ncomp)**2)``.
+        A pair is valid only when both atom indices are real. An all-padding
+        frame contributes a neutral zero.
+    """
+    xp = array_api_compat.array_namespace(elem, maskf)
+    nf, nloc = maskf.shape
+    component_mask = xp.reshape(
+        xp.broadcast_to(maskf[:, :, None], (nf, nloc, ncomp)),
+        (nf, nloc * ncomp),
+    )
+    masked = elem * component_mask[:, :, None] * component_mask[:, None, :]
+    per_frame_sum = xp.sum(xp.reshape(masked, (nf, -1)), axis=-1)
+    per_frame_dof = xp.square(xp.sum(component_mask, axis=-1))
+    has_dof = per_frame_dof > 0
+    safe_dof = xp.where(has_dof, per_frame_dof, xp.ones_like(per_frame_dof))
+    per_frame = xp.where(
+        has_dof, per_frame_sum / safe_dof, xp.zeros_like(per_frame_sum)
+    )
+    return xp.mean(per_frame)
+
+
 def per_frame_component_mean(err: Array) -> Array:
     """Idiom 2 primitive: per-frame mean over the flattened component axis.
 
diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py
index 1109c21131..c345023ce7 100644
--- a/source/tests/common/dpmodel/test_loss_padding.py
+++ b/source/tests/common/dpmodel/test_loss_padding.py
@@ -490,6 +490,13 @@ def _padded_force(f_A, f_B):
     return np.stack([f_A_pad, f_B], axis=0)  # [2, NP, 3]
 
 
+def _padded_hessian(h_A, h_B):
+    """Pad both Cartesian axes of frame A's Hessian to the batch width."""
+    h_A_pad = np.zeros((3 * NP, 3 * NP), dtype=np.float64)
+    h_A_pad[: 3 * NA, : 3 * NA] = h_A
+    return np.stack([h_A_pad, h_B], axis=0)
+
+
 def _padded_atom(arr_A, arr_B, ncomp):
     """Pad arr_A from [NA, ncomp] to [NP, ncomp] with zeros, stack with arr_B."""
     pad = np.zeros((NP, ncomp), dtype=np.float64)
@@ -749,6 +756,133 @@ def test_no_op_for_non_mixed(self):
         )
 
 
+class TestDPModelEnergyLossHessianGradAccum:
+    """Hessian reductions exclude placeholder rows and columns."""
+
+    def _make_loss(self):
+        return EnergyLoss(
+            starter_learning_rate=1.0,
+            start_pref_e=0.0,
+            limit_pref_e=0.0,
+            start_pref_f=0.0,
+            limit_pref_f=0.0,
+            start_pref_v=0.0,
+            limit_pref_v=0.0,
+            start_pref_ae=0.0,
+            limit_pref_ae=0.0,
+            start_pref_pf=0.0,
+            limit_pref_pf=0.0,
+            start_pref_h=1.0,
+            limit_pref_h=1.0,
+        )
+
+    @staticmethod
+    def _set_hessian(model_pred, label, pred_h, label_h):
+        model_pred["hessian"] = pred_h
+        label["hessian"] = label_h
+        label["find_hessian"] = 1.0
+        return model_pred, label
+
+    def _loss_fn(self, model_pred, label, natoms):
+        loss, _ = self._make_loss().call(1.0, natoms, model_pred, label)
+        return float(loss)
+
+    def test_mse_grad_accum(self):
+        """Padded Hessian MSE equals the mean of per-frame Hessian MSEs."""
+        h_A = _rnd(3 * NA, 3 * NA)
+        h_A_hat = _rnd(3 * NA, 3 * NA)
+        h_B = _rnd(3 * NB, 3 * NB)
+        h_B_hat = _rnd(3 * NB, 3 * NB)
+
+        def make_A():
+            pred, label = _full_ener_dicts(
+                1,
+                NA,
+                np.zeros((1, 1)),
+                np.zeros((1, 1)),
+                mask=np.ones((1, NA)),
+            )
+            pred, label = self._set_hessian(pred, label, h_A[None], h_A_hat[None])
+            return pred, label, NA
+
+        def make_B():
+            pred, label = _full_ener_dicts(
+                1,
+                NB,
+                np.zeros((1, 1)),
+                np.zeros((1, 1)),
+                mask=np.ones((1, NB)),
+            )
+            pred, label = self._set_hessian(pred, label, h_B[None], h_B_hat[None])
+            return pred, label, NB
+
+        def make_padded():
+            pred, label = _full_ener_dicts(
+                2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD
+            )
+            pred, label = self._set_hessian(
+                pred,
+                label,
+                _padded_hessian(h_A, h_B),
+                _padded_hessian(h_A_hat, h_B_hat),
+            )
+            return pred, label, NP
+
+        assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded)
+
+    def test_placeholder_pairs_are_excluded_from_metrics(self):
+        """Ghost-touching matrix entries do not dilute or pollute MAE/RMSE."""
+        pred, label = _full_ener_dicts(
+            1,
+            NP,
+            np.zeros((1, 1)),
+            np.zeros((1, 1)),
+            mask=np.array([[1.0] * NA + [0.0] * (NP - NA)]),
+        )
+        pred_h = np.zeros((1, 3 * NP, 3 * NP), dtype=np.float64)
+        label_h = np.full_like(pred_h, 100.0)
+        label_h[:, : 3 * NA, : 3 * NA] = 2.0
+        pred, label = self._set_hessian(pred, label, pred_h, label_h)
+
+        loss, more_loss = self._make_loss().call(1.0, NP, pred, label, mae=True)
+
+        np.testing.assert_allclose(loss, 4.0)
+        np.testing.assert_allclose(more_loss["rmse_h"], 2.0)
+        np.testing.assert_allclose(more_loss["mae_h"], 2.0)
+
+    def test_no_op_for_non_mixed(self):
+        """An all-ones mask preserves the established Hessian reduction."""
+        pred_h = _rnd(1, 3 * NP, 3 * NP)
+        label_h = _rnd(1, 3 * NP, 3 * NP)
+        pred_mask, label_mask = _full_ener_dicts(
+            1,
+            NP,
+            np.zeros((1, 1)),
+            np.zeros((1, 1)),
+            mask=np.ones((1, NP)),
+        )
+        pred_plain, label_plain = _full_ener_dicts(
+            1, NP, np.zeros((1, 1)), np.zeros((1, 1))
+        )
+        pred_mask, label_mask = self._set_hessian(
+            pred_mask, label_mask, pred_h, label_h
+        )
+        pred_plain, label_plain = self._set_hessian(
+            pred_plain, label_plain, pred_h, label_h
+        )
+
+        loss_mask, metrics_mask = self._make_loss().call(
+            1.0, NP, pred_mask, label_mask, mae=True
+        )
+        loss_plain, metrics_plain = self._make_loss().call(
+            1.0, NP, pred_plain, label_plain, mae=True
+        )
+
+        np.testing.assert_allclose(loss_mask, loss_plain)
+        np.testing.assert_allclose(metrics_mask["rmse_h"], metrics_plain["rmse_h"])
+        np.testing.assert_allclose(metrics_mask["mae_h"], metrics_plain["mae_h"])
+
+
 class TestDPModelEnergyLossVirialGradAccum:
     """Idiom 2 (extensive, k=9) for the virial (has_v) term.
 

From 484fd6dfed1f1f4a8b201af998411b48bf14fed4 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Tue, 21 Jul 2026 12:06:13 +0800
Subject: [PATCH 04/12] fix(argcheck): alias Hessian loss to energy

Keep the energy loss as the canonical schema for Hessian prefactors, normalize the legacy ener_hess type to ener, and update documentation links to the canonical loss fields.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 deepmd/utils/argcheck.py                  |  6 +++---
 doc/model/train-energy-hessian.md         |  8 ++++---
 source/tests/common/test_argcheck_loss.py | 26 +++++++++++++++++++++++
 3 files changed, 34 insertions(+), 6 deletions(-)
 create mode 100644 source/tests/common/test_argcheck_loss.py

diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 91dfab4824..2e4185cee9 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -4276,7 +4276,7 @@ def limit_pref(item: str) -> str:
 loss_args_plugin = ArgsPlugin()
 
 
-@loss_args_plugin.register("ener")
+@loss_args_plugin.register("ener", alias=["ener_hess"])
 def loss_ener() -> list[Argument]:
     doc_start_pref_e = start_pref("energy", abbr="e")
     doc_limit_pref_e = limit_pref("energy")
@@ -4971,7 +4971,7 @@ def loss_tensor() -> list[Argument]:
 
 
 def loss_variant_type_args() -> Variant:
-    doc_loss = "The type of the loss. When the fitting type is `ener`, the loss type should be set to `ener`, `dens` (Only DPA4/SeZM supported), or left unset. When the fitting type is `property`, the loss type should be set to `property`. When the fitting type is `dipole` or `polar`, the loss type should be set to `tensor`."
+    doc_loss = "The type of the loss. When the fitting type is `ener`, the loss type should be set to `ener`, its legacy alias `ener_hess`, `dens` (Only DPA4/SeZM supported), or left unset. Hessian supervision is configured through `start_pref_h` and `limit_pref_h` on the `ener` loss. When the fitting type is `property`, the loss type should be set to `property`. When the fitting type is `dipole` or `polar`, the loss type should be set to `tensor`."
 
     return Variant(
         "type",
@@ -4983,7 +4983,7 @@ def loss_variant_type_args() -> Variant:
 
 
 def loss_args() -> list[Argument]:
-    doc_loss = "The definition of loss function. The loss type should be set to `tensor`, `property`, `ener`, `dens` or left unset."
+    doc_loss = "The definition of loss function. The loss type should be set to `tensor`, `property`, `ener`, `dens` or left unset. The legacy `ener_hess` type is normalized to `ener`."
     ca = Argument(
         "loss", dict, [], [loss_variant_type_args()], optional=True, doc=doc_loss
     )
diff --git a/doc/model/train-energy-hessian.md b/doc/model/train-energy-hessian.md
index 092f8d30e3..85aae38c80 100644
--- a/doc/model/train-energy-hessian.md
+++ b/doc/model/train-energy-hessian.md
@@ -8,7 +8,7 @@ To train a model that takes Hessian matrices, i.e., the second order derivatives
 
 ## Energy Hessian Loss
 
-If you want to train with Hessians, you are expected to add the start and limit prefactors of Hessians, i.e., {ref}`start_pref_h <loss[ener_hess]/start_pref_h>` and {ref}`limit_pref_h <loss[ener_hess]/limit_pref_h>` to the {ref}`loss <loss>` section in the `input.json`:
+If you want to train with Hessians, add the start and limit prefactors of Hessians, i.e., {ref}`start_pref_h <loss[ener]/start_pref_h>` and {ref}`limit_pref_h <loss[ener]/limit_pref_h>` to the standard energy {ref}`loss <loss>` section in the `input.json`:
 
 ```json
    "loss": {
@@ -24,9 +24,11 @@ If you want to train with Hessians, you are expected to add the start and limit
    },
 ```
 
-The options {ref}`start_pref_e <loss[ener_hess]/start_pref_e>`, {ref}`limit_pref_e <loss[ener_hess]/limit_pref_e>`, {ref}`start_pref_f <loss[ener_hess]/start_pref_f>`, {ref}`limit_pref_f <loss[ener_hess]/limit_pref_f>`, {ref}`start_pref_v <loss[ener_hess]/start_pref_v>` and {ref}`limit_pref_v <loss[ener_hess]/limit_pref_v>` determine the start and limit prefactors of energy, force, and virial, respectively. The calculation and definition of Hessian loss are the same as for the other terms.
+The legacy loss type `"ener_hess"` remains accepted as an alias of `"ener"`, but new input files should use the canonical `"ener"` type.
 
-If one does not want to train with virial, then he/she may set the virial prefactors {ref}`start_pref_v <loss[ener_hess]/start_pref_v>` and {ref}`limit_pref_v <loss[ener_hess]/limit_pref_v>` to 0.
+The options {ref}`start_pref_e <loss[ener]/start_pref_e>`, {ref}`limit_pref_e <loss[ener]/limit_pref_e>`, {ref}`start_pref_f <loss[ener]/start_pref_f>`, {ref}`limit_pref_f <loss[ener]/limit_pref_f>`, {ref}`start_pref_v <loss[ener]/start_pref_v>` and {ref}`limit_pref_v <loss[ener]/limit_pref_v>` determine the start and limit prefactors of energy, force, and virial, respectively. The calculation and definition of Hessian loss are the same as for the other terms.
+
+If one does not want to train with virial, set the virial prefactors {ref}`start_pref_v <loss[ener]/start_pref_v>` and {ref}`limit_pref_v <loss[ener]/limit_pref_v>` to 0.
 
 ## Hessian Format in PyTorch
 
diff --git a/source/tests/common/test_argcheck_loss.py b/source/tests/common/test_argcheck_loss.py
new file mode 100644
index 0000000000..b623e1ab4e
--- /dev/null
+++ b/source/tests/common/test_argcheck_loss.py
@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.utils.argcheck import (
+    loss_args,
+)
+
+
+def test_energy_hessian_loss_type_is_energy_alias() -> None:
+    """The legacy Hessian loss type must normalize to the energy schema."""
+    canonical = loss_args().normalize_value(
+        {
+            "type": "ener",
+            "start_pref_h": 2.0,
+            "limit_pref_h": 1.0,
+        }
+    )
+    legacy = loss_args().normalize_value(
+        {
+            "type": "ener_hess",
+            "start_pref_h": 2.0,
+            "limit_pref_h": 1.0,
+        }
+    )
+
+    loss_args().check_value(legacy, strict=True)
+    assert legacy == canonical
+    assert legacy["type"] == "ener"

From d64c14fc1df8b70dbb4ed7a2ee10719243beecc4 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Tue, 21 Jul 2026 12:16:37 +0800
Subject: [PATCH 05/12] fix(pt): align Hessian with energy loss

Move Hessian supervision into EnergyStdLoss, keep EnergyHessianStdLoss and ener_hess as compatibility aliases, enable Hessian when either prefactor is active, and apply the shared placeholder-pair mask to PT metrics and training loss.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 deepmd/pt/loss/ener.py                        | 137 ++++++++----------
 deepmd/pt/train/training.py                   |  16 +-
 .../tests/common/dpmodel/test_loss_padding.py |  42 ++++++
 source/tests/pt/test_loss.py                  |  38 +++++
 4 files changed, 148 insertions(+), 85 deletions(-)

diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py
index bff23e2ba2..8d885ea73a 100644
--- a/deepmd/pt/loss/ener.py
+++ b/deepmd/pt/loss/ener.py
@@ -8,6 +8,7 @@
 
 from deepmd.dpmodel.loss.reduction import (
     masked_atom_mean,
+    masked_pair_mean,
     per_frame_component_mean,
 )
 from deepmd.pt.loss.loss import (
@@ -60,6 +61,8 @@ def __init__(
         start_pref_gf: float = 0.0,
         limit_pref_gf: float = 0.0,
         numb_generalized_coord: int = 0,
+        start_pref_h: float = 0.0,
+        limit_pref_h: float = 0.0,
         loss_func: str = "mse",
         inference: bool = False,
         use_huber: bool = False,
@@ -107,6 +110,10 @@ def __init__(
             The prefactor of generalized force loss at the end of the training.
         numb_generalized_coord : int
             The dimension of generalized coordinates.
+        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.
         loss_func : str
             Loss function type. Options: 'mse' (Mean Squared Error, L2 loss, default) or 'mae' (Mean Absolute Error, L1 loss).
             MAE loss is less sensitive to outliers compared to MSE loss.
@@ -156,6 +163,7 @@ def __init__(
         self.has_ae = (start_pref_ae != 0.0 and limit_pref_ae != 0.0) or inference
         self.has_pf = (start_pref_pf != 0.0 and limit_pref_pf != 0.0) or inference
         self.has_gf = start_pref_gf != 0.0 and limit_pref_gf != 0.0
+        self.has_h = (start_pref_h != 0.0 or limit_pref_h != 0.0) or inference
 
         self.start_pref_e = start_pref_e
         self.limit_pref_e = limit_pref_e
@@ -169,6 +177,8 @@ def __init__(
         self.limit_pref_pf = limit_pref_pf
         self.start_pref_gf = start_pref_gf
         self.limit_pref_gf = limit_pref_gf
+        self.start_pref_h = start_pref_h
+        self.limit_pref_h = limit_pref_h
         self.use_default_pf = use_default_pf
         self.relative_f = relative_f
         self.enable_atom_ener_coeff = enable_atom_ener_coeff
@@ -237,6 +247,7 @@ def forward(
         pref_ae = self.limit_pref_ae + (self.start_pref_ae - self.limit_pref_ae) * coef
         pref_pf = self.limit_pref_pf + (self.start_pref_pf - self.limit_pref_pf) * coef
         pref_gf = self.limit_pref_gf + (self.start_pref_gf - self.limit_pref_gf) * coef
+        pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * coef
 
         loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0]
         more_loss = {}
@@ -775,6 +786,37 @@ def forward(
                     f"Loss type {self.loss_func} is not implemented for atomic energy loss."
                 )
 
+        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
+            if maskf is not None:
+                hessian_shape = (_nf, _nloc * 3, _nloc * 3)
+                diff_h = label["hessian"].reshape(hessian_shape) - model_pred[
+                    "hessian"
+                ].reshape(hessian_shape)
+                # Both Cartesian axes must refer to real atoms for a Hessian
+                # element to contribute to the loss or display denominator.
+                l2_hessian_loss = masked_pair_mean(torch.square(diff_h), maskf, ncomp=3)
+            else:
+                diff_h = label["hessian"].reshape(-1) - model_pred["hessian"].reshape(
+                    -1
+                )
+                l2_hessian_loss = torch.mean(torch.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
+            more_loss["rmse_h"] = self.display_if_exist(
+                l2_hessian_loss.sqrt().detach(), find_hessian
+            )
+            if mae:
+                if maskf is not None:
+                    mae_h = masked_pair_mean(torch.abs(diff_h), maskf, ncomp=3)
+                else:
+                    mae_h = torch.mean(torch.abs(diff_h))
+                more_loss["mae_h"] = self.display_if_exist(mae_h.detach(), find_hessian)
+
         if not self.inference:
             more_loss["rmse"] = torch.sqrt(loss.detach())
         return model_pred, loss, more_loss
@@ -856,6 +898,17 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     default=1.0,
                 )
             )
+        if self.has_h:
+            label_requirement.append(
+                DataRequirementItem(
+                    "hessian",
+                    ndof=1,
+                    atomic=False,
+                    must=False,
+                    high_prec=False,
+                    special_shape="hessian",
+                )
+            )
         return label_requirement
 
     def serialize(self) -> dict:
@@ -866,7 +919,7 @@ def serialize(self) -> dict:
         dict
             The serialized loss module
         """
-        return {
+        data = {
             "@class": "EnergyLoss",
             "@version": 4,
             "starter_learning_rate": self.starter_learning_rate,
@@ -892,6 +945,10 @@ def serialize(self) -> dict:
             "use_default_pf": self.use_default_pf,
             "intensive_ener_virial": self.intensive_ener_virial,
         }
+        if self.start_pref_h != 0.0 or self.limit_pref_h != 0.0:
+            data["start_pref_h"] = self.start_pref_h
+            data["limit_pref_h"] = self.limit_pref_h
+        return data
 
     @classmethod
     def deserialize(cls, data: dict) -> "TaskLoss":
@@ -918,80 +975,4 @@ def deserialize(cls, data: dict) -> "TaskLoss":
 
 
 class EnergyHessianStdLoss(EnergyStdLoss):
-    def __init__(
-        self,
-        start_pref_h: float = 0.0,
-        limit_pref_h: float = 0.0,
-        **kwargs: Any,
-    ) -> None:
-        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: dict[str, torch.Tensor],
-        model: torch.nn.Module,
-        label: dict[str, torch.Tensor],
-        natoms: int,
-        learning_rate: float,
-        mae: bool = False,
-    ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]:
-        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 = torch.mean(torch.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 = torch.mean(torch.abs(diff_h))
-                more_loss["mae_h"] = self.display_if_exist(mae_h.detach(), find_hessian)
-
-        if not self.inference:
-            more_loss["rmse"] = torch.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
+    """Backward-compatible name for the unified energy loss."""
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index 8dd68c8cdb..a0b098ee10 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -44,7 +44,6 @@
     DenoiseLoss,
     DeNSLoss,
     DOSLoss,
-    EnergyHessianStdLoss,
     EnergySpinLoss,
     EnergyStdLoss,
     PopulationLoss,
@@ -2484,7 +2483,10 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
 
 def whether_hessian(loss_params: dict[str, Any]) -> bool:
     loss_type = loss_params.get("type", "ener")
-    return loss_type == "ener" and loss_params.get("start_pref_h", 0.0) > 0.0
+    return loss_type in {"ener", "ener_hess"} and (
+        loss_params.get("start_pref_h", 0.0) != 0.0
+        or loss_params.get("limit_pref_h", 0.0) != 0.0
+    )
 
 
 def prepare_model_for_loss(
@@ -2501,17 +2503,17 @@ def prepare_model_for_loss(
                 prepare_model_for_loss(sub_model, sub_loss)
         return
     if hasattr(model, "set_active_mode_from_loss"):
-        model.set_active_mode_from_loss(loss_params.get("type", "ener"))
+        loss_type = loss_params.get("type", "ener")
+        model.set_active_mode_from_loss(
+            "ener" if loss_type == "ener_hess" else loss_type
+        )
 
 
 def get_loss(
     loss_params: dict[str, Any], start_lr: float, _ntypes: int, _model: Any
 ) -> TaskLoss:
     loss_type = loss_params.get("type", "ener")
-    if whether_hessian(loss_params):
-        loss_params["starter_learning_rate"] = start_lr
-        return EnergyHessianStdLoss(**loss_params)
-    elif loss_type == "ener":
+    if loss_type in {"ener", "ener_hess"}:
         loss_params["starter_learning_rate"] = start_lr
         return EnergyStdLoss(**loss_params)
     elif loss_type == "dens":
diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py
index c345023ce7..8789584192 100644
--- a/source/tests/common/dpmodel/test_loss_padding.py
+++ b/source/tests/common/dpmodel/test_loss_padding.py
@@ -850,6 +850,48 @@ def test_placeholder_pairs_are_excluded_from_metrics(self):
         np.testing.assert_allclose(more_loss["rmse_h"], 2.0)
         np.testing.assert_allclose(more_loss["mae_h"], 2.0)
 
+    def test_pt_placeholder_pairs_are_excluded_from_metrics(self):
+        """The PyTorch energy loss uses the same Hessian pair mask."""
+        import pytest
+
+        torch = pytest.importorskip("torch")
+        from deepmd.pt.loss.ener import EnergyStdLoss as PTEnergyStdLoss
+
+        pred, label = _full_ener_dicts(
+            1,
+            NP,
+            np.zeros((1, 1)),
+            np.zeros((1, 1)),
+            mask=np.array([[1.0] * NA + [0.0] * (NP - NA)]),
+        )
+        pred_h = np.zeros((1, 3 * NP, 3 * NP), dtype=np.float64)
+        label_h = np.full_like(pred_h, 100.0)
+        label_h[:, : 3 * NA, : 3 * NA] = 2.0
+        pred, label = self._set_hessian(pred, label, pred_h, label_h)
+        pt_pred = {
+            key: torch.as_tensor(value, device="cpu")
+            if isinstance(value, np.ndarray)
+            else value
+            for key, value in pred.items()
+        }
+        pt_label = {
+            key: torch.as_tensor(value, device="cpu")
+            if isinstance(value, np.ndarray)
+            else value
+            for key, value in label.items()
+        }
+        loss_obj = PTEnergyStdLoss(
+            starter_learning_rate=1.0,
+            start_pref_h=1.0,
+            limit_pref_h=1.0,
+        )
+
+        _, loss, more_loss = loss_obj({}, lambda: pt_pred, pt_label, NP, 1.0, mae=True)
+
+        torch.testing.assert_close(loss, loss.new_tensor(4.0))
+        torch.testing.assert_close(more_loss["rmse_h"], loss.new_tensor(2.0))
+        torch.testing.assert_close(more_loss["mae_h"], loss.new_tensor(2.0))
+
     def test_no_op_for_non_mixed(self):
         """An all-ones mask preserves the established Hessian reduction."""
         pred_h = _rnd(1, 3 * NP, 3 * NP)
diff --git a/source/tests/pt/test_loss.py b/source/tests/pt/test_loss.py
index 2519111357..a37447ebb8 100644
--- a/source/tests/pt/test_loss.py
+++ b/source/tests/pt/test_loss.py
@@ -708,6 +708,44 @@ def fake_model():
             self.assertTrue(np.isnan(pt_more_loss_h_absent[f"l2_{key}_loss"]))
 
 
+class TestEnergyHessianLossCompatibility(unittest.TestCase):
+    """Hessian supervision is canonical on EnergyStdLoss in PyTorch."""
+
+    def test_legacy_class_matches_energy_loss(self) -> None:
+        kwargs = {
+            "starter_learning_rate": 1.0,
+            "start_pref_h": 2.0,
+            "limit_pref_h": 1.0,
+        }
+        canonical = EnergyStdLoss(**kwargs)
+        legacy = EnergyHessianStdLoss(**kwargs)
+
+        self.assertEqual(canonical.serialize(), legacy.serialize())
+        self.assertEqual(canonical.label_requirement, legacy.label_requirement)
+        hessian_req = next(
+            item for item in canonical.label_requirement if item.key == "hessian"
+        )
+        self.assertFalse(hessian_req.atomic)
+        self.assertEqual(hessian_req.special_shape, "hessian")
+
+    def test_legacy_loss_type_dispatches_to_energy_loss(self) -> None:
+        from deepmd.pt.train.training import (
+            get_loss,
+            whether_hessian,
+        )
+
+        for loss_type in ("ener", "ener_hess"):
+            params = {
+                "type": loss_type,
+                "start_pref_h": 0.0,
+                "limit_pref_h": 1.0,
+            }
+            self.assertTrue(whether_hessian(params))
+            loss = get_loss(params, start_lr=1.0, _ntypes=0, _model=None)
+            self.assertIs(type(loss), EnergyStdLoss)
+            self.assertTrue(loss.has_h)
+
+
 class TestEnerSpinLoss(LossCommonTest):
     def setUp(self) -> None:
         self.start_lr = 1.1

From d57fe96edcc8219d2d4c50602c9f3b6e6c753e1c Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Tue, 21 Jul 2026 12:44:25 +0800
Subject: [PATCH 06/12] refactor(train): derive Hessian mode from data
 requirements

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 deepmd/dpmodel/model/ener_model.py            |  3 +
 deepmd/jax/train/trainer.py                   | 22 ++++--
 deepmd/pt/model/model/ener_model.py           |  3 +
 deepmd/pt/train/training.py                   | 78 ++++++++++++++-----
 deepmd/utils/data.py                          | 12 +++
 source/tests/jax/test_make_hessian_model.py   | 10 +++
 source/tests/jax/test_training.py             | 71 +++++++++++++++++
 .../tests/pt/model/test_make_hessian_model.py | 10 +++
 source/tests/pt/model/test_sezm_model.py      |  2 +-
 source/tests/pt/test_change_bias.py           |  4 -
 source/tests/pt/test_loss.py                  | 68 +++++++++++++++-
 11 files changed, 250 insertions(+), 33 deletions(-)

diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py
index 09b50a6f17..0aa791f5ea 100644
--- a/deepmd/dpmodel/model/ener_model.py
+++ b/deepmd/dpmodel/model/ener_model.py
@@ -75,6 +75,9 @@ def __init__(
         self.hess_fitting_def = None
 
     def enable_hessian(self) -> None:
+        """Enable Hessian outputs without changing an already-enabled model."""
+        if self._enable_hessian:
+            return
         self.hess_fitting_def = deepcopy(self.atomic_output_def())
         self.hess_fitting_def["energy"].r_hessian = True
         self._enable_hessian = True
diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py
index 81f15b0d54..2f162a1aad 100644
--- a/deepmd/jax/train/trainer.py
+++ b/deepmd/jax/train/trainer.py
@@ -94,6 +94,7 @@
 )
 from deepmd.utils.data import (
     DataRequirementItem,
+    has_data_requirement,
 )
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
@@ -195,17 +196,24 @@ def __init__(
         learning_rate_param = jdata["learning_rate"]
         self.lr = self._get_lr_and_coef(learning_rate_param)
         self.losses = self._build_losses(jdata, learning_rate_param)
-        for model_key, loss_obj in self.losses.items():
-            # Hessians are expensive and are only exposed by models whose task
-            # loss explicitly requests them. Enable each multi-task branch
-            # independently so unrelated branches keep their normal outputs.
-            if getattr(loss_obj, "has_h", False):
-                self.models[model_key].enable_hessian()
-        self.loss = self.losses if self.multi_task else self.losses[DEFAULT_TASK_KEY]
         self.data_requirements_by_task = {
             model_key: list(self.losses[model_key].label_requirement)
             for model_key in self.model_keys
         }
+        for model_key, requirements in self.data_requirements_by_task.items():
+            # Hessians are expensive and are only exposed by models whose task
+            # loss requests the corresponding label. Enable each multi-task
+            # branch independently so unrelated branches keep normal outputs.
+            if has_data_requirement(requirements, "hessian"):
+                enable_hessian = getattr(self.models[model_key], "enable_hessian", None)
+                if not callable(enable_hessian):
+                    raise RuntimeError(
+                        f"Model {type(self.models[model_key]).__name__} does not "
+                        "support Hessian supervision."
+                    )
+                enable_hessian()
+                self.model_params_by_task[model_key]["hessian_mode"] = True
+        self.loss = self.losses if self.multi_task else self.losses[DEFAULT_TASK_KEY]
 
         self.valid_numb_batch_by_task = self._valid_numb_batch_by_task()
         self.valid_numb_batch = (
diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
index 28387553fb..bc858b6419 100644
--- a/deepmd/pt/model/model/ener_model.py
+++ b/deepmd/pt/model/model/ener_model.py
@@ -39,6 +39,9 @@ def __init__(
         self._hessian_enabled = False
 
     def enable_hessian(self) -> None:
+        """Enable Hessian outputs without changing an already-enabled model."""
+        if self._hessian_enabled:
+            return
         self.__class__ = make_hessian_model(type(self))
         self.hess_fitting_def = super(type(self), self).atomic_output_def()
         self.requires_hessian("energy")
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index a0b098ee10..d74b8313e0 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -120,6 +120,7 @@
 )
 from deepmd.utils.data import (
     DataRequirementItem,
+    has_data_requirement,
 )
 from deepmd.utils.finetune import (
     warn_configuration_mismatch_during_finetune,
@@ -449,12 +450,12 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
         if self.zero_stage > 0 and self.opt_type == "LKF":
             raise ValueError("training.zero_stage does not support LKF optimizer.")
 
-        # loss_param_tmp for Hessian activation
-        loss_param_tmp = None
+        # Loss parameters are also used to select SeZM/DeNS execution modes.
+        loss_params_for_model = None
         if not self.multi_task:
-            loss_param_tmp = config["loss"]
+            loss_params_for_model = config["loss"]
         else:
-            loss_param_tmp = {
+            loss_params_for_model = {
                 model_key: config["loss_dict"][model_key]
                 for model_key in self.model_keys
             }
@@ -466,10 +467,9 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
             self.model = get_model_for_wrapper(
                 model_params,
                 resuming=resuming,
-                _loss_params=loss_param_tmp,
             )
         # SeZM specific process for DeNS training
-        prepare_model_for_loss(self.model, loss_param_tmp)
+        prepare_model_for_loss(self.model, loss_params_for_model)
 
         # Loss
         if not self.multi_task:
@@ -489,6 +489,22 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                     loss_param, lr_param, ntypes, self.model[model_key]
                 )
 
+        # Losses own the interpretation of their prefactors. The trainer only
+        # consumes their data contract when selecting expensive model outputs.
+        loss_data_requirements = (
+            {
+                model_key: self.loss[model_key].label_requirement
+                for model_key in self.model_keys
+            }
+            if self.multi_task
+            else self.loss.label_requirement
+        )
+        prepare_model_for_data_requirements(
+            self.model,
+            loss_data_requirements,
+            model_params,
+        )
+
         # Data
         if not self.multi_task:
             # add data requirement for labels
@@ -2481,14 +2497,6 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
     return additional_data_requirement
 
 
-def whether_hessian(loss_params: dict[str, Any]) -> bool:
-    loss_type = loss_params.get("type", "ener")
-    return loss_type in {"ener", "ener_hess"} and (
-        loss_params.get("start_pref_h", 0.0) != 0.0
-        or loss_params.get("limit_pref_h", 0.0) != 0.0
-    )
-
-
 def prepare_model_for_loss(
     model: Any,
     loss_params: dict[str, Any] | None,
@@ -2509,6 +2517,43 @@ def prepare_model_for_loss(
         )
 
 
+def prepare_model_for_data_requirements(
+    model: Any,
+    data_requirements: list[DataRequirementItem] | dict[str, list[DataRequirementItem]],
+    model_params: dict[str, Any],
+) -> None:
+    """Enable model outputs requested by a loss's data requirements.
+
+    Hessian prefactors are deliberately not inspected here. The presence of a
+    ``hessian`` requirement is the loss-to-trainer contract for enabling that
+    expensive output. The mode is persisted in the model definition so reload
+    and freeze paths reconstruct the same model interface.
+    """
+    if isinstance(model, dict):
+        if not isinstance(data_requirements, dict):
+            raise TypeError("Multi-task models require per-task data requirements.")
+        for model_key, sub_model in model.items():
+            prepare_model_for_data_requirements(
+                sub_model,
+                data_requirements[model_key],
+                model_params["model_dict"][model_key],
+            )
+        return
+    if isinstance(data_requirements, dict):
+        raise TypeError("Single-task models require a list of data requirements.")
+    if not has_data_requirement(data_requirements, "hessian"):
+        return
+    enable_hessian = getattr(model, "enable_hessian", None)
+    if not callable(enable_hessian):
+        raise RuntimeError(
+            f"Model {type(model).__name__} does not support Hessian supervision."
+        )
+    enable_hessian()
+    model_params["hessian_mode"] = True
+    if hasattr(model, "model_def_script"):
+        model.model_def_script = json.dumps(model_params)
+
+
 def get_loss(
     loss_params: dict[str, Any], start_lr: float, _ntypes: int, _model: Any
 ) -> TaskLoss:
@@ -2570,11 +2615,8 @@ def get_single_model(
 def get_model_for_wrapper(
     _model_params: dict[str, Any],
     resuming: bool = False,
-    _loss_params: dict[str, Any] | None = None,
 ) -> Any:
     if "model_dict" not in _model_params:
-        if _loss_params is not None and whether_hessian(_loss_params):
-            _model_params["hessian_mode"] = True
         _model = get_single_model(
             _model_params,
         )
@@ -2583,8 +2625,6 @@ def get_model_for_wrapper(
         model_keys = list(_model_params["model_dict"])
         do_case_embd, case_embd_index = get_case_embd_config(_model_params)
         for _model_key in model_keys:
-            if _loss_params is not None and whether_hessian(_loss_params[_model_key]):
-                _model_params["model_dict"][_model_key]["hessian_mode"] = True
             _model[_model_key] = get_single_model(
                 _model_params["model_dict"][_model_key],
             )
diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py
index c54969f582..c7a682c921 100644
--- a/deepmd/utils/data.py
+++ b/deepmd/utils/data.py
@@ -5,6 +5,9 @@
 import copy
 import functools
 import logging
+from collections.abc import (
+    Iterable,
+)
 from concurrent.futures import (
     ThreadPoolExecutor,
     as_completed,
@@ -1214,3 +1217,12 @@ def __eq__(self, value: object, /) -> bool:
 
     def __repr__(self) -> str:
         return f"DataRequirementItem({self.dict})"
+
+
+def has_data_requirement(requirements: Iterable[DataRequirementItem], key: str) -> bool:
+    """Return whether a collection requests data identified by ``key``.
+
+    Consumers should use requirement items as the contract with losses instead
+    of reinterpreting loss-specific configuration such as prefactors.
+    """
+    return any(item.key == key for item in requirements)
diff --git a/source/tests/jax/test_make_hessian_model.py b/source/tests/jax/test_make_hessian_model.py
index bb25bd67ca..e9fe727139 100644
--- a/source/tests/jax/test_make_hessian_model.py
+++ b/source/tests/jax/test_make_hessian_model.py
@@ -174,3 +174,13 @@ def test_output_def(self) -> None:
             self.model_hess.model_output_def()["energy_derv_r_derv_r"].category,
             OutputVariableCategory.DERV_R_DERV_R,
         )
+
+    def test_enable_hessian_is_idempotent(self) -> None:
+        """Restored Hessian models may be enabled again by loss requirements."""
+        self.model_valu.enable_hessian()
+        enabled_def = self.model_valu.hess_fitting_def
+
+        self.model_valu.enable_hessian()
+
+        self.assertIs(self.model_valu.hess_fitting_def, enabled_def)
+        self.assertTrue(self.model_valu.model_output_def()["energy"].r_hessian)
diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py
index b720a71789..4fc6f54788 100644
--- a/source/tests/jax/test_training.py
+++ b/source/tests/jax/test_training.py
@@ -71,6 +71,9 @@
 from deepmd.utils.compat import (
     convert_optimizer_v31_to_v32,
 )
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
 from deepmd.utils.env_mat_stat import (
     StatItem,
 )
@@ -200,6 +203,74 @@ def _minimal_jax_multitask_config(model_params: dict) -> dict:
     }
 
 
+class _RequirementModel:
+    """Minimal model double for trainer requirement-driven output tests."""
+
+    def __init__(self) -> None:
+        self.hessian_enable_calls = 0
+
+    def enable_hessian(self) -> None:
+        self.hessian_enable_calls += 1
+
+    def get_dim_fparam(self) -> int:
+        return 0
+
+
+@patch("deepmd.jax.train.trainer.DPTrainer._build_losses")
+@patch("deepmd.jax.train.trainer.get_model")
+def test_jax_hessian_mode_follows_loss_data_requirement(
+    get_model,
+    build_losses,
+) -> None:
+    """The trainer consumes the loss requirement instead of its prefactors."""
+    model = _RequirementModel()
+    get_model.return_value = model
+    build_losses.return_value = {
+        "Default": SimpleNamespace(
+            label_requirement=[DataRequirementItem("hessian", ndof=1)]
+        )
+    }
+    model_params = {"type_map": ["O"], "descriptor": {}}
+
+    trainer = DPTrainer(_minimal_jax_config(model_params))
+
+    assert model.hessian_enable_calls == 1
+    assert trainer.model_def_script["hessian_mode"] is True
+
+
+@patch("deepmd.jax.train.trainer.DPTrainer._build_losses")
+@patch("deepmd.jax.train.trainer.get_model")
+def test_jax_multitask_hessian_only_enables_requesting_branch(
+    get_model,
+    build_losses,
+) -> None:
+    """Each multi-task branch follows its own loss data requirements."""
+    model_a = _RequirementModel()
+    model_b = _RequirementModel()
+    get_model.side_effect = [model_a, model_b]
+    build_losses.return_value = {
+        "task_a": SimpleNamespace(
+            label_requirement=[DataRequirementItem("hessian", ndof=1)]
+        ),
+        "task_b": SimpleNamespace(
+            label_requirement=[DataRequirementItem("energy", ndof=1)]
+        ),
+    }
+    model_params = {
+        "model_dict": {
+            "task_a": {"type_map": ["O"], "descriptor": {}},
+            "task_b": {"type_map": ["O"], "descriptor": {}},
+        }
+    }
+
+    trainer = DPTrainer(_minimal_jax_multitask_config(model_params))
+
+    assert model_a.hessian_enable_calls == 1
+    assert model_b.hessian_enable_calls == 0
+    assert trainer.model_def_script["model_dict"]["task_a"]["hessian_mode"] is True
+    assert "hessian_mode" not in trainer.model_def_script["model_dict"]["task_b"]
+
+
 def _shared_jax_model_config(*, share_fitting: bool = True) -> dict:
     shared_dict: dict = {
         "shared_type_map": ["O", "H", "B"],
diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py
index f2387ca8b0..2bb1539f71 100644
--- a/source/tests/pt/model/test_make_hessian_model.py
+++ b/source/tests/pt/model/test_make_hessian_model.py
@@ -182,3 +182,13 @@ def test_output_def(self) -> None:
             self.model_hess.model_output_def()["energy_derv_r_derv_r"].category,
             OutputVariableCategory.DERV_R_DERV_R,
         )
+
+    def test_enable_hessian_is_idempotent(self) -> None:
+        """Restored Hessian models may be enabled again by loss requirements."""
+        self.model_valu.enable_hessian()
+        enabled_type = type(self.model_valu)
+
+        self.model_valu.enable_hessian()
+
+        self.assertIs(type(self.model_valu), enabled_type)
+        self.assertTrue(self.model_valu.model_output_def()["energy"].r_hessian)
diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py
index 494338a2da..21210a6260 100644
--- a/source/tests/pt/model/test_sezm_model.py
+++ b/source/tests/pt/model/test_sezm_model.py
@@ -1005,7 +1005,7 @@ def _build_wrapper(use_compile: bool) -> ModelWrapper:
                 "water_1": {"type": "ener"},
                 "water_2": {"type": "ener"},
             }
-            models = get_model_for_wrapper(mt_cfg, _loss_params=loss_params)
+            models = get_model_for_wrapper(mt_cfg)
             prepare_model_for_loss(models, loss_params)
             wrapper = ModelWrapper(models)
             wrapper.share_params(shared_links, {"water_1": 0.5, "water_2": 0.5})
diff --git a/source/tests/pt/test_change_bias.py b/source/tests/pt/test_change_bias.py
index a883449d22..746656ac8c 100644
--- a/source/tests/pt/test_change_bias.py
+++ b/source/tests/pt/test_change_bias.py
@@ -96,7 +96,6 @@ def setUp(self) -> None:
         self.model_path_user_bias = Path(current_path) / (
             model_name + "user_bias" + ".pt"
         )
-        self.loss_params = self.config["loss"]
 
     def test_change_bias_with_data(self) -> None:
         run_dp(
@@ -108,7 +107,6 @@ def test_change_bias_with_data(self) -> None:
         model_params = state_dict["model"]["_extra_state"]["model_params"]
         model_for_wrapper = get_model_for_wrapper(
             model_params,
-            _loss_params=self.loss_params,
         )
         wrapper = ModelWrapper(model_for_wrapper)
         wrapper.load_state_dict(state_dict["model"])
@@ -134,7 +132,6 @@ def test_change_bias_with_data_sys_file(self) -> None:
         model_params = state_dict["model"]["_extra_state"]["model_params"]
         model_for_wrapper = get_model_for_wrapper(
             model_params,
-            _loss_params=self.loss_params,
         )
         wrapper = ModelWrapper(model_for_wrapper)
         wrapper.load_state_dict(state_dict["model"])
@@ -158,7 +155,6 @@ def test_change_bias_with_user_defined(self) -> None:
         model_params = state_dict["model"]["_extra_state"]["model_params"]
         model_for_wrapper = get_model_for_wrapper(
             model_params,
-            _loss_params=self.loss_params,
         )
         wrapper = ModelWrapper(model_for_wrapper)
         wrapper.load_state_dict(state_dict["model"])
diff --git a/source/tests/pt/test_loss.py b/source/tests/pt/test_loss.py
index a37447ebb8..b4c10ae690 100644
--- a/source/tests/pt/test_loss.py
+++ b/source/tests/pt/test_loss.py
@@ -731,7 +731,6 @@ def test_legacy_class_matches_energy_loss(self) -> None:
     def test_legacy_loss_type_dispatches_to_energy_loss(self) -> None:
         from deepmd.pt.train.training import (
             get_loss,
-            whether_hessian,
         )
 
         for loss_type in ("ener", "ener_hess"):
@@ -740,10 +739,75 @@ def test_legacy_loss_type_dispatches_to_energy_loss(self) -> None:
                 "start_pref_h": 0.0,
                 "limit_pref_h": 1.0,
             }
-            self.assertTrue(whether_hessian(params))
             loss = get_loss(params, start_lr=1.0, _ntypes=0, _model=None)
             self.assertIs(type(loss), EnergyStdLoss)
             self.assertTrue(loss.has_h)
+            self.assertTrue(
+                any(item.key == "hessian" for item in loss.label_requirement)
+            )
+
+    def test_model_hessian_mode_follows_loss_data_requirement(self) -> None:
+        from deepmd.pt.train.training import (
+            prepare_model_for_data_requirements,
+        )
+
+        class HessianCapableModel:
+            def __init__(self) -> None:
+                self.enable_calls = 0
+                self.model_def_script = ""
+
+            def enable_hessian(self) -> None:
+                self.enable_calls += 1
+
+        model = HessianCapableModel()
+        model_params: dict = {}
+        requirements = [
+            DataRequirementItem(
+                "hessian",
+                ndof=1,
+                special_shape="hessian",
+            )
+        ]
+
+        prepare_model_for_data_requirements(model, requirements, model_params)
+
+        self.assertEqual(model.enable_calls, 1)
+        self.assertTrue(model_params["hessian_mode"])
+        self.assertEqual(model.model_def_script, '{"hessian_mode": true}')
+
+    def test_multitask_hessian_mode_only_enables_requesting_branch(self) -> None:
+        from deepmd.pt.train.training import (
+            prepare_model_for_data_requirements,
+        )
+
+        class HessianCapableModel:
+            def __init__(self) -> None:
+                self.enable_calls = 0
+                self.model_def_script = ""
+
+            def enable_hessian(self) -> None:
+                self.enable_calls += 1
+
+        models = {
+            "with_hessian": HessianCapableModel(),
+            "without_hessian": object(),
+        }
+        requirements = {
+            "with_hessian": [DataRequirementItem("hessian", ndof=1)],
+            "without_hessian": [DataRequirementItem("energy", ndof=1)],
+        }
+        model_params = {
+            "model_dict": {
+                "with_hessian": {},
+                "without_hessian": {},
+            }
+        }
+
+        prepare_model_for_data_requirements(models, requirements, model_params)
+
+        self.assertEqual(models["with_hessian"].enable_calls, 1)
+        self.assertTrue(model_params["model_dict"]["with_hessian"]["hessian_mode"])
+        self.assertNotIn("hessian_mode", model_params["model_dict"]["without_hessian"])
 
 
 class TestEnerSpinLoss(LossCommonTest):

From a663f92429b818aa118df4142f1bdebf663678d3 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Mon, 27 Jul 2026 14:17:33 +0800
Subject: [PATCH 07/12] fix: dispatch the hessian loss on loss_func and pad
 mixed batches

The hessian term hard-coded squared error, so loss_func='mae' trained
energy/force/virial on L1 while the hessian minimised L2 inside the same
scalar; mae_h was display-only. Add the mse/mae/NotImplementedError
branch in both backends and reject use_huber together with has_h.

_merge_batch_data could not pad a hessian either: a (1, 9*N**2) row per
frame made the non-atomic concatenate raise for batch_size 'mixed:N',
which the JAX trainer reaches through DeepmdDataSystem.get_batch().
Embed each frame's block in the top-left of the padded square instead.
---
 deepmd/dpmodel/loss/ener.py                   |  21 +++-
 deepmd/pt/loss/ener.py                        |  21 +++-
 deepmd/utils/data_system.py                   |  20 ++-
 .../tests/common/dpmodel/test_loss_padding.py | 118 ++++++++++++++++++
 .../tests/common/test_data_system_hessian.py  | 107 ++++++++++++++++
 5 files changed, 276 insertions(+), 11 deletions(-)
 create mode 100644 source/tests/common/test_data_system_hessian.py

diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index 79c43a1646..d3b8cc2ed8 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -231,6 +231,8 @@ def __init__(
             raise RuntimeError(
                 "Huber loss is not implemented for force with atom_pref, generalized force and relative force. "
             )
+        if self.use_huber and self.has_h:
+            raise RuntimeError("Huber loss is not implemented for hessian. ")
 
     def call(
         self,
@@ -769,15 +771,24 @@ def call(
                     (-1,),
                 )
                 l2_hessian_loss = xp.mean(xp.square(diff_h))
-            loss += pref_h * find_hessian * l2_hessian_loss
-            more_loss["rmse_h"] = self.display_if_exist(
-                xp.sqrt(l2_hessian_loss), find_hessian
-            )
-            if mae:
+            mae_h = None
+            if self.loss_func == "mae" or mae:
                 if maskf is not None:
                     mae_h = masked_pair_mean(xp.abs(diff_h), maskf, ncomp=3)
                 else:
                     mae_h = xp.mean(xp.abs(diff_h))
+            if self.loss_func == "mse":
+                loss += pref_h * find_hessian * l2_hessian_loss
+            elif self.loss_func == "mae":
+                loss += pref_h * find_hessian * mae_h
+            else:
+                raise NotImplementedError(
+                    f"Loss type {self.loss_func} is not implemented for hessian loss."
+                )
+            more_loss["rmse_h"] = self.display_if_exist(
+                xp.sqrt(l2_hessian_loss), find_hessian
+            )
+            if mae:
                 more_loss["mae_h"] = self.display_if_exist(mae_h, find_hessian)
 
         self.l2_l = loss
diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py
index 8d885ea73a..c4d8731bd4 100644
--- a/deepmd/pt/loss/ener.py
+++ b/deepmd/pt/loss/ener.py
@@ -207,6 +207,8 @@ def __init__(
             raise RuntimeError(
                 "Huber loss is not implemented for force with atom_pref, generalized force and relative force. "
             )
+        if self.use_huber and self.has_h:
+            raise RuntimeError("Huber loss is not implemented for hessian. ")
 
     def forward(
         self,
@@ -806,15 +808,24 @@ def forward(
                 more_loss["l2_hessian_loss"] = self.display_if_exist(
                     l2_hessian_loss.detach(), find_hessian
                 )
-            loss += pref_h * l2_hessian_loss
-            more_loss["rmse_h"] = self.display_if_exist(
-                l2_hessian_loss.sqrt().detach(), find_hessian
-            )
-            if mae:
+            mae_h = None
+            if self.loss_func == "mae" or mae:
                 if maskf is not None:
                     mae_h = masked_pair_mean(torch.abs(diff_h), maskf, ncomp=3)
                 else:
                     mae_h = torch.mean(torch.abs(diff_h))
+            if self.loss_func == "mse":
+                loss += pref_h * l2_hessian_loss
+            elif self.loss_func == "mae":
+                loss += pref_h * mae_h
+            else:
+                raise NotImplementedError(
+                    f"Loss type {self.loss_func} is not implemented for hessian loss."
+                )
+            more_loss["rmse_h"] = self.display_if_exist(
+                l2_hessian_loss.sqrt().detach(), find_hessian
+            )
+            if mae:
                 more_loss["mae_h"] = self.display_if_exist(mae_h.detach(), find_hessian)
 
         if not self.inference:
diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py
index bd8e2b64c7..dfc4076a9f 100644
--- a/deepmd/utils/data_system.py
+++ b/deepmd/utils/data_system.py
@@ -554,7 +554,25 @@ def _merge_batch_data(self, batch_data: list[dict]) -> dict:
             if kk not in batch_data[0]:
                 continue
             b_data["find_" + kk] = batch_data[0]["find_" + kk]
-            if not vv["atomic"]:
+            if vv.get("special_shape") == "hessian" or kk == "hessian":
+                # A Hessian is a (3 * natoms, 3 * natoms) matrix, so neither
+                # branch below pads it correctly: concatenating raises on
+                # ragged systems and copying a flat prefix would scatter the
+                # rows. Embed each frame's block in the top-left corner of the
+                # padded square instead; the padded rows and columns stay zero
+                # and are dropped by the loss mask.
+                padded_dof = max_natoms * 3
+                merged = np.zeros(
+                    (len(batch_data), padded_dof, padded_dof),
+                    dtype=batch_data[0][kk].dtype,
+                )
+                for ii, bb in enumerate(batch_data):
+                    frame_dof = bb["natoms_vec"][0] * 3
+                    merged[ii, :frame_dof, :frame_dof] = bb[kk][0].reshape(
+                        frame_dof, frame_dof
+                    )
+                b_data[kk] = merged.reshape(len(batch_data), -1)
+            elif not vv["atomic"]:
                 b_data[kk] = np.concatenate([bb[kk] for bb in batch_data], axis=0)
             else:
                 b_data[kk] = np.zeros(
diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py
index 8789584192..ec2cedac15 100644
--- a/source/tests/common/dpmodel/test_loss_padding.py
+++ b/source/tests/common/dpmodel/test_loss_padding.py
@@ -924,6 +924,124 @@ def test_no_op_for_non_mixed(self):
         np.testing.assert_allclose(metrics_mask["rmse_h"], metrics_plain["rmse_h"])
         np.testing.assert_allclose(metrics_mask["mae_h"], metrics_plain["mae_h"])
 
+    def _make_mae_loss(self):
+        return EnergyLoss(
+            starter_learning_rate=1.0,
+            start_pref_e=0.0,
+            limit_pref_e=0.0,
+            start_pref_f=0.0,
+            limit_pref_f=0.0,
+            start_pref_v=0.0,
+            limit_pref_v=0.0,
+            start_pref_ae=0.0,
+            limit_pref_ae=0.0,
+            start_pref_pf=0.0,
+            limit_pref_pf=0.0,
+            start_pref_h=1.0,
+            limit_pref_h=1.0,
+            loss_func="mae",
+        )
+
+    def _constant_residual_dicts(self, residual):
+        """One frame whose Hessian residual is *residual* everywhere."""
+        pred, label = _full_ener_dicts(
+            1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP))
+        )
+        pred_h = np.zeros((1, 3 * NP, 3 * NP), dtype=np.float64)
+        label_h = np.full_like(pred_h, residual)
+        return self._set_hessian(pred, label, pred_h, label_h)
+
+    def test_mae_loss_func_uses_l1_for_the_hessian(self):
+        """``loss_func='mae'`` must not train the Hessian term on squared error."""
+        residual = 10.0
+        pred, label = self._constant_residual_dicts(residual)
+
+        loss, more_loss = self._make_mae_loss().call(1.0, NP, pred, label, mae=True)
+
+        # L2 would give residual ** 2; the displays stay MSE-derived.
+        np.testing.assert_allclose(loss, residual)
+        np.testing.assert_allclose(more_loss["mae_h"], residual)
+        np.testing.assert_allclose(more_loss["rmse_h"], residual)
+
+        mse_pred, mse_label = self._constant_residual_dicts(residual)
+        mse_loss, _ = self._make_loss().call(1.0, NP, mse_pred, mse_label)
+        np.testing.assert_allclose(mse_loss, residual**2)
+
+    def test_unknown_loss_func_is_rejected_for_the_hessian(self):
+        """Any other loss_func must fail rather than silently train on MSE."""
+        import pytest
+
+        loss_obj = self._make_loss()
+        # loss_func is validated in __init__, so set it directly to reach the
+        # Hessian branch's dispatch.
+        loss_obj.loss_func = "huber"
+        pred, label = self._constant_residual_dicts(1.0)
+        with pytest.raises(NotImplementedError, match="hessian loss"):
+            loss_obj.call(1.0, NP, pred, label)
+
+    def test_huber_rejects_hessian_supervision(self):
+        """use_huber must not silently fall back to a raw MSE Hessian term."""
+        import pytest
+
+        with pytest.raises(RuntimeError, match="Huber loss is not implemented"):
+            EnergyLoss(
+                starter_learning_rate=1.0,
+                start_pref_h=1.0,
+                limit_pref_h=1.0,
+                use_huber=True,
+            )
+
+    def test_pt_mae_loss_func_uses_l1_for_the_hessian(self):
+        """The PyTorch loss must dispatch on loss_func for the Hessian too."""
+        import pytest
+
+        torch = pytest.importorskip("torch")
+        from deepmd.pt.loss.ener import EnergyStdLoss as PTEnergyStdLoss
+
+        residual = 10.0
+        pred, label = self._constant_residual_dicts(residual)
+        pt_pred = {
+            key: torch.as_tensor(value, device="cpu")
+            if isinstance(value, np.ndarray)
+            else value
+            for key, value in pred.items()
+        }
+        pt_label = {
+            key: torch.as_tensor(value, device="cpu")
+            if isinstance(value, np.ndarray)
+            else value
+            for key, value in label.items()
+        }
+
+        def make_loss(loss_func):
+            return PTEnergyStdLoss(
+                starter_learning_rate=1.0,
+                start_pref_e=0.0,
+                limit_pref_e=0.0,
+                start_pref_f=0.0,
+                limit_pref_f=0.0,
+                start_pref_h=1.0,
+                limit_pref_h=1.0,
+                loss_func=loss_func,
+            )
+
+        _, mae_loss, more_loss = make_loss("mae")(
+            {}, lambda: pt_pred, pt_label, NP, 1.0, mae=True
+        )
+        _, mse_loss, _ = make_loss("mse")({}, lambda: pt_pred, pt_label, NP, 1.0)
+
+        torch.testing.assert_close(mae_loss, mae_loss.new_tensor(residual))
+        torch.testing.assert_close(mse_loss, mse_loss.new_tensor(residual**2))
+        torch.testing.assert_close(more_loss["mae_h"], mae_loss.new_tensor(residual))
+
+        with pytest.raises(RuntimeError, match="Huber loss is not implemented"):
+            PTEnergyStdLoss(
+                starter_learning_rate=1.0,
+                start_pref_h=1.0,
+                limit_pref_h=1.0,
+                use_huber=True,
+            )
+
 
 class TestDPModelEnergyLossVirialGradAccum:
     """Idiom 2 (extensive, k=9) for the virial (has_v) term.
diff --git a/source/tests/common/test_data_system_hessian.py b/source/tests/common/test_data_system_hessian.py
new file mode 100644
index 0000000000..193126e08b
--- /dev/null
+++ b/source/tests/common/test_data_system_hessian.py
@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Mixed-system batching must pad Hessian labels as square blocks.
+
+``DeepmdDataSystem.get_batch()`` merges frames from systems with different
+atom counts.  A Hessian is a ``(3 * natoms, 3 * natoms)`` matrix, so neither
+the non-atomic concatenate nor the atomic flat-prefix copy can pad it: the
+former raises on ragged systems and the latter would scatter the rows.  The
+JAX trainer calls ``get_batch()`` directly, so this is reachable from a
+``batch_size: "mixed:N"`` run with Hessian supervision.
+"""
+
+import shutil
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
+
+NATOMS = (2, 3)
+NFRAMES = 2
+# One constant per system, so a mis-scattered block is visible as a stray value.
+HESSIAN_VALUES = (1.0, 2.0)
+
+
+def _write_system(root: Path, natoms: int, value: float) -> None:
+    """Write a one-set system whose Hessian is a constant square block."""
+    root.mkdir(parents=True)
+    (root / "type.raw").write_text("\n".join(["0"] * natoms) + "\n")
+    (root / "type_map.raw").write_text("A\n")
+    set_dir = root / "set.000"
+    set_dir.mkdir()
+    rng = np.random.default_rng(natoms)
+    np.save(set_dir / "coord.npy", rng.random((NFRAMES, natoms * 3)))
+    np.save(set_dir / "box.npy", np.tile(np.eye(3).reshape(9) * 20.0, (NFRAMES, 1)))
+    np.save(set_dir / "energy.npy", rng.random((NFRAMES, 1)))
+    dof = natoms * 3
+    np.save(
+        set_dir / "hessian.npy",
+        np.full((NFRAMES, dof * dof), value, dtype=np.float64),
+    )
+
+
+class TestMixedBatchHessian(unittest.TestCase):
+    def setUp(self) -> None:
+        self.tmpdir = Path(tempfile.mkdtemp())
+        systems = []
+        for natoms, value in zip(NATOMS, HESSIAN_VALUES, strict=True):
+            path = self.tmpdir / f"sys_{natoms}"
+            _write_system(path, natoms, value)
+            systems.append(str(path))
+        self.ds = DeepmdDataSystem(systems, "mixed:2", 1, 2.0)
+        self.ds.add("energy", 1, atomic=False, must=True, high_prec=True)
+        self.ds.add("hessian", 1, atomic=False, must=True, special_shape="hessian")
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.tmpdir, ignore_errors=True)
+
+    def _ragged_batch(self) -> dict:
+        """Merge exactly one frame from each system.
+
+        ``get_batch_mixed`` picks systems at random, so it does not reliably
+        produce the ragged case this covers.
+        """
+        batch_data = []
+        for sys_idx in range(self.ds.nsystems):
+            bb_data = self.ds.data_systems[sys_idx].get_batch(1)
+            bb_data["natoms_vec"] = self.ds.natoms_vec[sys_idx]
+            bb_data["default_mesh"] = self.ds.default_mesh[sys_idx]
+            batch_data.append(bb_data)
+        return self.ds._merge_batch_data(batch_data)
+
+    def test_mixed_batch_embeds_each_hessian_block(self) -> None:
+        """Each frame's block sits in the top-left of the padded square."""
+        batch = self._ragged_batch()
+        max_natoms = int(batch["natoms_vec"][0])
+        self.assertEqual(max_natoms, max(NATOMS))
+        padded_dof = max_natoms * 3
+
+        self.assertEqual(batch["hessian"].shape, (len(NATOMS), padded_dof * padded_dof))
+        hessian = batch["hessian"].reshape(-1, padded_dof, padded_dof)
+
+        for frame, (natoms, value) in enumerate(
+            zip(NATOMS, HESSIAN_VALUES, strict=True)
+        ):
+            frame_dof = natoms * 3
+            self.assertEqual(int(batch["real_natoms_vec"][frame, 0]), natoms)
+            np.testing.assert_array_equal(hessian[frame, :frame_dof, :frame_dof], value)
+            # Everything outside the block is padding and must stay zero.
+            np.testing.assert_array_equal(hessian[frame, frame_dof:, :], 0.0)
+            np.testing.assert_array_equal(hessian[frame, :, frame_dof:], 0.0)
+
+    def test_get_batch_accepts_a_mixed_hessian_batch(self) -> None:
+        """The public entry point must not raise for any sampled combination."""
+        for _ in range(20):
+            batch = self.ds.get_batch()
+            padded_dof = int(batch["natoms_vec"][0]) * 3
+            self.assertEqual(batch["hessian"].shape[1], padded_dof * padded_dof)
+
+
+if __name__ == "__main__":
+    unittest.main()

From a1ed51d2a55427223e6c4c44ec791d60ce08db94 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Thu, 30 Jul 2026 11:35:18 +0800
Subject: [PATCH 08/12] style(loss): remove trailing spaces from errors

Address the outstanding requested-change review comments.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 deepmd/dpmodel/loss/ener.py | 4 ++--
 deepmd/pt/loss/ener.py      | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index d3b8cc2ed8..7d4fb6d99b 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -229,10 +229,10 @@ def __init__(
             self.has_pf or self.has_gf or self.relative_f is not None
         ):
             raise RuntimeError(
-                "Huber loss is not implemented for force with atom_pref, generalized force and relative force. "
+                "Huber loss is not implemented for force with atom_pref, generalized force and relative force."
             )
         if self.use_huber and self.has_h:
-            raise RuntimeError("Huber loss is not implemented for hessian. ")
+            raise RuntimeError("Huber loss is not implemented for hessian.")
 
     def call(
         self,
diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py
index c4d8731bd4..5c85b053f2 100644
--- a/deepmd/pt/loss/ener.py
+++ b/deepmd/pt/loss/ener.py
@@ -205,10 +205,10 @@ def __init__(
             self.has_pf or self.has_gf or self.relative_f is not None
         ):
             raise RuntimeError(
-                "Huber loss is not implemented for force with atom_pref, generalized force and relative force. "
+                "Huber loss is not implemented for force with atom_pref, generalized force and relative force."
             )
         if self.use_huber and self.has_h:
-            raise RuntimeError("Huber loss is not implemented for hessian. ")
+            raise RuntimeError("Huber loss is not implemented for hessian.")
 
     def forward(
         self,

From 4d667b290cffb804b748e5cda56601cce21d05da Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Thu, 30 Jul 2026 22:04:14 +0800
Subject: [PATCH 09/12] fix(loss): resolve hessian review notes

Address Hessian serialization compatibility, avoid quadratic change-bias labels, align Paddle activation, and document JAX support.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 deepmd/dpmodel/loss/ener.py                   | 13 ++++-
 deepmd/pd/loss/ener.py                        |  9 ++--
 deepmd/pd/train/training.py                   |  6 ++-
 deepmd/pt/loss/ener.py                        | 16 ++++--
 doc/model/train-energy-hessian.md             | 52 +++++++++++++++----
 source/tests/common/dpmodel/test_loss_ener.py | 12 +++++
 source/tests/pd/test_loss.py                  | 41 +++++++++++++++
 source/tests/pt/test_loss.py                  | 22 ++++++++
 8 files changed, 152 insertions(+), 19 deletions(-)

diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index 7d4fb6d99b..34f7858b83 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -891,7 +891,10 @@ def serialize(self) -> dict:
         """
         data = {
             "@class": "EnergyLoss",
-            "@version": 4,
+            # Version 5 identifies the opt-in Hessian fields. Keep ordinary
+            # energy losses at version 4 so readers that already support the
+            # standard schema remain interoperable across backends.
+            "@version": 5 if self.has_h else 4,
             "starter_learning_rate": self.starter_learning_rate,
             "start_pref_e": self.start_pref_e,
             "limit_pref_e": self.limit_pref_e,
@@ -938,9 +941,15 @@ def deserialize(cls, data: dict) -> "Loss":
         """
         data = data.copy()
         version = data.pop("@version")
-        check_version_compatibility(version, 4, 1)
+        check_version_compatibility(version, 5, 1)
         data.pop("@class")
         # Backward compatibility: version 1-2 used legacy normalization
         if version < 3:
             data.setdefault("intensive_ener_virial", False)
+        # Version 5 introduced explicit Hessian prefactors. Older payloads
+        # represent an ordinary energy loss unless these development fields
+        # were already present.
+        if version < 5:
+            data.setdefault("start_pref_h", 0.0)
+            data.setdefault("limit_pref_h", 0.0)
         return cls(**data)
diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py
index 4ab9d71ab9..1601238370 100644
--- a/deepmd/pd/loss/ener.py
+++ b/deepmd/pd/loss/ener.py
@@ -649,7 +649,9 @@ def __init__(
             Other keyword arguments.
         """
         super().__init__(**kwargs)
-        self.has_h = (start_pref_h != 0.0 and limit_pref_h != 0.0) or self.inference
+        # A scheduled term is active when either endpoint is nonzero. Requiring
+        # both endpoints silently disabled valid ramp-up and ramp-down inputs.
+        self.has_h = (start_pref_h != 0.0 or limit_pref_h != 0.0) or self.inference
 
         self.start_pref_h = start_pref_h
         self.limit_pref_h = limit_pref_h
@@ -701,10 +703,11 @@ def label_requirement(self) -> list[DataRequirementItem]:
             label_requirement.append(
                 DataRequirementItem(
                     "hessian",
-                    ndof=1,  # 9=3*3 --> 3N*3N=ndof*natoms*natoms
-                    atomic=True,
+                    ndof=1,
+                    atomic=False,
                     must=False,
                     high_prec=False,
+                    special_shape="hessian",
                 )
             )
         return label_requirement
diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py
index 92225feeda..d703d8f8c5 100644
--- a/deepmd/pd/train/training.py
+++ b/deepmd/pd/train/training.py
@@ -1356,8 +1356,12 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
 
 
 def whether_hessian(loss_params: dict[str, Any]) -> bool:
+    """Return whether either Hessian schedule endpoint enables supervision."""
     loss_type = loss_params.get("type", "ener")
-    return loss_type == "ener" and loss_params.get("start_pref_h", 0.0) > 0.0
+    return loss_type in {"ener", "ener_hess"} and (
+        loss_params.get("start_pref_h", 0.0) != 0.0
+        or loss_params.get("limit_pref_h", 0.0) != 0.0
+    )
 
 
 def get_loss(
diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py
index 5c85b053f2..468b95f1eb 100644
--- a/deepmd/pt/loss/ener.py
+++ b/deepmd/pt/loss/ener.py
@@ -163,7 +163,10 @@ def __init__(
         self.has_ae = (start_pref_ae != 0.0 and limit_pref_ae != 0.0) or inference
         self.has_pf = (start_pref_pf != 0.0 and limit_pref_pf != 0.0) or inference
         self.has_gf = start_pref_gf != 0.0 and limit_pref_gf != 0.0
-        self.has_h = (start_pref_h != 0.0 or limit_pref_h != 0.0) or inference
+        # Hessian labels scale quadratically with the atom count. Unlike the
+        # linear-size labels above, do not request them merely because a mock
+        # inference loss is collecting requirements for ``dp change-bias``.
+        self.has_h = start_pref_h != 0.0 or limit_pref_h != 0.0
 
         self.start_pref_e = start_pref_e
         self.limit_pref_e = limit_pref_e
@@ -932,7 +935,9 @@ def serialize(self) -> dict:
         """
         data = {
             "@class": "EnergyLoss",
-            "@version": 4,
+            # Only Hessian-bearing payloads need the version-5 schema. Keeping
+            # ordinary energy losses at version 4 preserves existing readers.
+            "@version": 5 if self.has_h else 4,
             "starter_learning_rate": self.starter_learning_rate,
             "start_pref_e": self.start_pref_e,
             "limit_pref_e": self.limit_pref_e,
@@ -977,11 +982,16 @@ def deserialize(cls, data: dict) -> "TaskLoss":
         """
         data = data.copy()
         version = data.pop("@version")
-        check_version_compatibility(version, 4, 1)
+        check_version_compatibility(version, 5, 1)
         data.pop("@class")
         # Handle backward compatibility for older versions without intensive_ener_virial
         if version < 3:
             data.setdefault("intensive_ener_virial", False)
+        # Version 5 introduced explicit Hessian prefactors. Version 1-4
+        # payloads therefore default to the standard non-Hessian loss.
+        if version < 5:
+            data.setdefault("start_pref_h", 0.0)
+            data.setdefault("limit_pref_h", 0.0)
         return cls(**data)
 
 
diff --git a/doc/model/train-energy-hessian.md b/doc/model/train-energy-hessian.md
index b2ab0f307b..d84fc754fe 100644
--- a/doc/model/train-energy-hessian.md
+++ b/doc/model/train-energy-hessian.md
@@ -1,7 +1,7 @@
-# Fit energy Hessian {{ pytorch_icon }}
+# Fit energy Hessian {{ pytorch_icon }} {{ jax_icon }}
 
 > [!NOTE]
-> **Supported backends**: PyTorch {{ pytorch_icon }}
+> **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}
 
 To train a model that takes Hessian matrices, i.e., the second order derivatives of energies w.r.t coordinates as input, you only need to prepare full Hessian matrices and modify the `loss` section to define the Hessian-specific settings, keeping other sections the same as the normal energy model's input script.
 
@@ -25,13 +25,15 @@ If you want to train with Hessians, add the start and limit prefactors of Hessia
 
 The legacy loss type `"ener_hess"` remains accepted as an alias of `"ener"`, but new input files should use the canonical `"ener"` type.
 
+Setting either Hessian prefactor to a nonzero value enables Hessian supervision, so schedules may ramp the term up from zero or down to zero. Both prefactors must be zero to disable the term. Earlier releases could silently disable the term on some backends when exactly one endpoint was zero.
+
 The options {ref}`start_pref_e <loss[ener]/start_pref_e>`, {ref}`limit_pref_e <loss[ener]/limit_pref_e>`, {ref}`start_pref_f <loss[ener]/start_pref_f>`, {ref}`limit_pref_f <loss[ener]/limit_pref_f>`, {ref}`start_pref_v <loss[ener]/start_pref_v>` and {ref}`limit_pref_v <loss[ener]/limit_pref_v>` determine the start and limit prefactors of energy, force, and virial, respectively. The calculation and definition of Hessian loss are the same as for the other terms.
 
 If one does not want to train with virial, set the virial prefactors {ref}`start_pref_v <loss[ener]/start_pref_v>` and {ref}`limit_pref_v <loss[ener]/limit_pref_v>` to 0.
 
-## Hessian Format in PyTorch
+## Hessian Data Format
 
-In the PyTorch backend, Hessian matrices are listed in `hessian.npy` files, and the data format may contain the following files:
+In the PyTorch and JAX backends, Hessian matrices are listed in `hessian.npy` files, and the data format may contain the following files:
 
 ```
 type.raw
@@ -52,7 +54,7 @@ Note that the `hessian.npy` should contain the **full** Hessian matrices with sh
 
 ## Train the Model
 
-There are two approaches to training a Hessian model. The first method involves training the model from scratch using the same command as in the `ener` mode within the PyTorch backend:
+There are two approaches to training a Hessian model. The first method involves training the model from scratch using the same command as in the `ener` mode:
 
 ::::{tab-set}
 
@@ -63,9 +65,16 @@ dp --pt train input.json
 ```
 :::
 
+:::{tab-item} JAX {{ jax_icon }}
+
+```bash
+dp --jax train input.json
+```
+:::
+
 ::::
 
-The second approach is to train a Hessian model from a pretrained energy model, following the same command as the `finetune` strategy within the PyTorch backend:
+The second approach is to train a Hessian model from a pretrained energy model, following the same command as the `finetune` strategy:
 
 ::::{tab-set}
 
@@ -76,6 +85,13 @@ dp --pt train input.json --finetune pretrained_energy.pt
 ```
 :::
 
+:::{tab-item} JAX {{ jax_icon }}
+
+```bash
+dp --jax train input.json --finetune pretrained_energy.jax
+```
+:::
+
 ::::
 
 The detailed loss can be found in `lcurve.out`:
@@ -93,9 +109,9 @@ The detailed loss can be found in `lcurve.out`:
 ## Test the Model
 
 > [!WARNING]
-> A model trained with Hessian cannot be frozen. If freezing is enforced, the model will be treated as a standard energy model, and the frozen one will no longer be able to output Hessian predictions.
+> A PyTorch model trained with Hessian cannot currently be frozen with its Hessian output. If freezing is enforced, the frozen model is treated as a standard energy model.
 
-If one do freeze and test a Hessian model using the commands:
+PyTorch can test such a frozen model as a standard energy model. JAX can preserve the Hessian output during freezing by passing `--hessian`:
 
 ::::{tab-set}
 
@@ -109,6 +125,15 @@ dp --pt test -m frozen_model.pth -s test_system -d ${output_prefix} -a -n 1
 ```
 :::
 
+:::{tab-item} JAX {{ jax_icon }}
+
+```bash
+dp --jax freeze -c . -o frozen_model.hlo --hessian
+
+dp --jax test -m frozen_model.hlo -s test_system -d ${output_prefix} -a -n 1
+```
+:::
+
 ::::
 
 If `dp --pt test -d ${output_prefix} -a` is specified, the output files will be the same as those in the `ener` mode, i.e.,
@@ -118,7 +143,7 @@ ${output_prefix}.e.out   ${output_prefix}.e_peratom.out  ${output_prefix}.f.out
 ${output_prefix}.v.out   ${output_prefix}.v_peratom.out
 ```
 
-If one intends to use the trained model for Hessian predictions, then he/she is supposed to test the model directly without performing a freezing operation:
+The backend checkpoints can also be tested directly without freezing:
 
 ::::{tab-set}
 
@@ -130,9 +155,16 @@ dp --pt test -m model.pt -s test_system -d ${output_prefix} -a -n 1
 ```
 :::
 
+:::{tab-item} JAX {{ jax_icon }}
+
+```bash
+dp --jax test -m model.ckpt.jax -s test_system -d ${output_prefix} -a -n 1
+```
+:::
+
 ::::
 
-If `dp --pt test -d ${output_prefix} -a` is specified, the predicted Hessian for each frame are output in an additional file in the working directory:
+When either backend tests a Hessian-capable checkpoint with `-d ${output_prefix} -a`, the predicted Hessian for each frame is written to an additional file in the working directory:
 
 ```
 ${output_prefix}.h.out
diff --git a/source/tests/common/dpmodel/test_loss_ener.py b/source/tests/common/dpmodel/test_loss_ener.py
index 73ea9a928f..761bc39768 100644
--- a/source/tests/common/dpmodel/test_loss_ener.py
+++ b/source/tests/common/dpmodel/test_loss_ener.py
@@ -171,6 +171,7 @@ def test_disabled_hessian_preserves_standard_serialization(self) -> None:
 
         data = loss_fn.serialize()
 
+        self.assertEqual(data["@version"], 4)
         self.assertNotIn("start_pref_h", data)
         self.assertNotIn("limit_pref_h", data)
         self.assertNotIn(
@@ -261,6 +262,7 @@ def test_serialize_deserialize(self) -> None:
             limit_pref_h=0.5,
         )
         data = loss_fn.serialize()
+        self.assertEqual(data["@version"], 5)
         self.assertEqual(data["start_pref_h"], 2.0)
         self.assertEqual(data["limit_pref_h"], 0.5)
         loss_fn2 = EnergyLoss.deserialize(data)
@@ -274,6 +276,16 @@ def test_serialize_deserialize(self) -> None:
         for key in more1:
             np.testing.assert_allclose(more1[key], more2[key])
 
+    def test_version_four_defaults_to_non_hessian_loss(self) -> None:
+        """Version-4 payloads deserialize without inventing Hessian supervision."""
+        data = EnergyLoss(starter_learning_rate=1.0).serialize()
+
+        loss_fn = EnergyLoss.deserialize(data)
+
+        self.assertFalse(loss_fn.has_h)
+        self.assertEqual(loss_fn.start_pref_h, 0.0)
+        self.assertEqual(loss_fn.limit_pref_h, 0.0)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pd/test_loss.py b/source/tests/pd/test_loss.py
index a7b8109e10..3c4e21cfb6 100644
--- a/source/tests/pd/test_loss.py
+++ b/source/tests/pd/test_loss.py
@@ -12,8 +12,13 @@
 )
 
 from deepmd.pd.loss import (
+    EnergyHessianStdLoss,
     EnergyStdLoss,
 )
+from deepmd.pd.train.training import (
+    get_loss,
+    whether_hessian,
+)
 from deepmd.pd.utils.dataset import (
     DeepmdDataSetForLoader,
 )
@@ -581,5 +586,41 @@ def fake_model():
             self.assertTrue(np.isnan(pd_more_loss_absent[f"l2_{key}_loss"].numpy()))
 
 
+class TestEnergyHessianLossCompatibility(unittest.TestCase):
+    """Paddle uses the same Hessian activation and data-shape contract."""
+
+    def test_either_schedule_endpoint_enables_hessian(self) -> None:
+        """Ramp-up and ramp-down schedules must both activate supervision."""
+        for loss_type in ("ener", "ener_hess"):
+            for start_pref_h, limit_pref_h in ((1.0, 0.0), (0.0, 1.0)):
+                params = {
+                    "type": loss_type,
+                    "start_pref_h": start_pref_h,
+                    "limit_pref_h": limit_pref_h,
+                }
+
+                self.assertTrue(whether_hessian(params))
+                loss = get_loss(params, start_lr=1.0, _ntypes=0, _model=None)
+                self.assertIsInstance(loss, EnergyHessianStdLoss)
+                self.assertTrue(loss.has_h)
+                hessian_req = next(
+                    item for item in loss.label_requirement if item.key == "hessian"
+                )
+                self.assertFalse(hessian_req.atomic)
+                self.assertEqual(hessian_req.special_shape, "hessian")
+
+    def test_zero_schedule_keeps_standard_energy_loss(self) -> None:
+        """Two zero endpoints must leave Hessian supervision disabled."""
+        params = {
+            "type": "ener",
+            "start_pref_h": 0.0,
+            "limit_pref_h": 0.0,
+        }
+
+        self.assertFalse(whether_hessian(params))
+        loss = get_loss(params, start_lr=1.0, _ntypes=0, _model=None)
+        self.assertIs(type(loss), EnergyStdLoss)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt/test_loss.py b/source/tests/pt/test_loss.py
index b4c10ae690..c5176f1fea 100644
--- a/source/tests/pt/test_loss.py
+++ b/source/tests/pt/test_loss.py
@@ -721,6 +721,7 @@ def test_legacy_class_matches_energy_loss(self) -> None:
         legacy = EnergyHessianStdLoss(**kwargs)
 
         self.assertEqual(canonical.serialize(), legacy.serialize())
+        self.assertEqual(canonical.serialize()["@version"], 5)
         self.assertEqual(canonical.label_requirement, legacy.label_requirement)
         hessian_req = next(
             item for item in canonical.label_requirement if item.key == "hessian"
@@ -746,6 +747,27 @@ def test_legacy_loss_type_dispatches_to_energy_loss(self) -> None:
                 any(item.key == "hessian" for item in loss.label_requirement)
             )
 
+    def test_inference_only_does_not_request_hessian(self) -> None:
+        """The change-bias mock loss must not allocate unused quadratic labels."""
+        loss = EnergyStdLoss(starter_learning_rate=1.0, inference=True)
+
+        self.assertFalse(loss.has_h)
+        self.assertNotIn(
+            "hessian",
+            {item.key for item in loss.label_requirement},
+        )
+
+    def test_version_four_defaults_to_non_hessian_loss(self) -> None:
+        """Older energy-loss payloads retain their non-Hessian behavior."""
+        data = EnergyStdLoss(starter_learning_rate=1.0).serialize()
+
+        loss = EnergyStdLoss.deserialize(data)
+
+        self.assertEqual(data["@version"], 4)
+        self.assertFalse(loss.has_h)
+        self.assertEqual(loss.start_pref_h, 0.0)
+        self.assertEqual(loss.limit_pref_h, 0.0)
+
     def test_model_hessian_mode_follows_loss_data_requirement(self) -> None:
         from deepmd.pt.train.training import (
             prepare_model_for_data_requirements,

From 074e277fbaaa298fa8740416f9f6b27ff62310b7 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz-bot@users.noreply.github.com>
Date: Fri, 31 Jul 2026 13:42:27 +0800
Subject: [PATCH 10/12] fix(pd): mask padded Hessian loss pairs

Reject unsupported Huber Hessian supervision and normalize Hessian MSE/MAE over real-real Cartesian pairs in mixed-system batches.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 deepmd/pd/loss/ener.py       | 55 +++++++++++++++++++++++++++++++-----
 source/tests/pd/test_loss.py | 37 ++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 7 deletions(-)

diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py
index 1601238370..1a9b2c38d0 100644
--- a/deepmd/pd/loss/ener.py
+++ b/deepmd/pd/loss/ener.py
@@ -37,6 +37,33 @@ def custom_huber_loss(
     return paddle.mean(loss)
 
 
+def masked_pair_mean(
+    elem: paddle.Tensor, maskf: paddle.Tensor, ncomp: int
+) -> paddle.Tensor:
+    """Average pair contributions over real atoms on both pair axes.
+
+    Hessian padding is square: a component is valid only when both its row atom
+    and column atom are real. The reduction is normalized per frame so mixed
+    systems with different atom counts receive the same weighting.
+    """
+    nf, nloc = maskf.shape
+    component_mask = paddle.broadcast_to(
+        maskf.astype(elem.dtype).reshape([nf, nloc, 1]),
+        [nf, nloc, ncomp],
+    ).reshape([nf, nloc * ncomp])
+    masked = elem * component_mask[:, :, None] * component_mask[:, None, :]
+    per_frame_sum = paddle.sum(masked.reshape([nf, -1]), axis=-1)
+    per_frame_dof = paddle.square(paddle.sum(component_mask, axis=-1))
+    has_dof = per_frame_dof > 0
+    safe_dof = paddle.where(has_dof, per_frame_dof, paddle.ones_like(per_frame_dof))
+    per_frame = paddle.where(
+        has_dof,
+        per_frame_sum / safe_dof,
+        paddle.zeros_like(per_frame_sum),
+    )
+    return paddle.mean(per_frame)
+
+
 class EnergyStdLoss(TaskLoss):
     def __init__(
         self,
@@ -652,6 +679,8 @@ def __init__(
         # A scheduled term is active when either endpoint is nonzero. Requiring
         # both endpoints silently disabled valid ramp-up and ramp-down inputs.
         self.has_h = (start_pref_h != 0.0 or limit_pref_h != 0.0) or self.inference
+        if self.use_huber and self.has_h:
+            raise RuntimeError("Huber loss is not implemented for hessian.")
 
         self.start_pref_h = start_pref_h
         self.limit_pref_h = limit_pref_h
@@ -674,12 +703,21 @@ def forward(
         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))
+            maskf = model_pred.get("mask")
+            if maskf is not None:
+                nf, nloc = maskf.shape
+                hessian_shape = [nf, nloc * 3, nloc * 3]
+                diff_h = label["hessian"].reshape(hessian_shape) - model_pred[
+                    "hessian"
+                ].reshape(hessian_shape)
+                l2_hessian_loss = masked_pair_mean(
+                    paddle.square(diff_h), maskf, ncomp=3
+                )
+            else:
+                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
@@ -688,7 +726,10 @@ def forward(
             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))
+                if maskf is not None:
+                    mae_h = masked_pair_mean(paddle.abs(diff_h), maskf, ncomp=3)
+                else:
+                    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:
diff --git a/source/tests/pd/test_loss.py b/source/tests/pd/test_loss.py
index 3c4e21cfb6..b5b0f503ed 100644
--- a/source/tests/pd/test_loss.py
+++ b/source/tests/pd/test_loss.py
@@ -621,6 +621,43 @@ def test_zero_schedule_keeps_standard_energy_loss(self) -> None:
         loss = get_loss(params, start_lr=1.0, _ntypes=0, _model=None)
         self.assertIs(type(loss), EnergyStdLoss)
 
+    def test_huber_is_rejected_for_hessian_supervision(self) -> None:
+        """Paddle must not silently optimize Hessians with raw MSE under Huber."""
+        with self.assertRaisesRegex(
+            RuntimeError, "Huber loss is not implemented for hessian"
+        ):
+            EnergyHessianStdLoss(start_pref_h=1.0, use_huber=True)
+
+    def test_hessian_padding_is_masked_on_both_atom_axes(self) -> None:
+        """Only the real-real Cartesian block contributes to Hessian loss."""
+        residual = np.full((1, 6, 6), 10.0)
+        residual[:, :3, :3] = 2.0
+        model_pred = {
+            "hessian": paddle.zeros([1, 6, 6], dtype="float64"),
+            "mask": paddle.to_tensor([[1.0, 0.0]], dtype="float64"),
+        }
+        label = {
+            "hessian": paddle.to_tensor(residual, dtype="float64"),
+            "find_hessian": 1.0,
+        }
+        loss_module = EnergyHessianStdLoss(
+            starter_learning_rate=1.0,
+            start_pref_h=1.0,
+            limit_pref_h=1.0,
+        )
+
+        def fake_model(**kwargs):
+            return model_pred
+
+        _, loss, more_loss = loss_module(
+            {}, fake_model, label, natoms=2, learning_rate=1.0, mae=True
+        )
+
+        self.assertTrue(np.allclose(loss.numpy(), 4.0))
+        self.assertTrue(np.allclose(more_loss["l2_hessian_loss"].numpy(), 4.0))
+        self.assertTrue(np.allclose(more_loss["rmse_h"].numpy(), 2.0))
+        self.assertTrue(np.allclose(more_loss["mae_h"].numpy(), 2.0))
+
 
 if __name__ == "__main__":
     unittest.main()

From d8094886b198a1e2bc6fa6a155aef37369fa7b88 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Fri, 31 Jul 2026 22:18:45 +0800
Subject: [PATCH 11/12] test(pt): use backend device for hessian loss

Construct PyTorch Hessian test inputs on the configured backend device so CUDA assertions compare colocated tensors.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 source/tests/common/dpmodel/test_loss_padding.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py
index cc571013de..2b0e71fcb5 100644
--- a/source/tests/common/dpmodel/test_loss_padding.py
+++ b/source/tests/common/dpmodel/test_loss_padding.py
@@ -856,6 +856,7 @@ def test_pt_placeholder_pairs_are_excluded_from_metrics(self):
 
         torch = pytest.importorskip("torch")
         from deepmd.pt.loss.ener import EnergyStdLoss as PTEnergyStdLoss
+        from deepmd.pt.utils import env
 
         pred, label = _full_ener_dicts(
             1,
@@ -869,13 +870,13 @@ def test_pt_placeholder_pairs_are_excluded_from_metrics(self):
         label_h[:, : 3 * NA, : 3 * NA] = 2.0
         pred, label = self._set_hessian(pred, label, pred_h, label_h)
         pt_pred = {
-            key: torch.as_tensor(value, device="cpu")
+            key: torch.as_tensor(value, device=env.DEVICE)
             if isinstance(value, np.ndarray)
             else value
             for key, value in pred.items()
         }
         pt_label = {
-            key: torch.as_tensor(value, device="cpu")
+            key: torch.as_tensor(value, device=env.DEVICE)
             if isinstance(value, np.ndarray)
             else value
             for key, value in label.items()
@@ -997,17 +998,18 @@ def test_pt_mae_loss_func_uses_l1_for_the_hessian(self):
 
         torch = pytest.importorskip("torch")
         from deepmd.pt.loss.ener import EnergyStdLoss as PTEnergyStdLoss
+        from deepmd.pt.utils import env
 
         residual = 10.0
         pred, label = self._constant_residual_dicts(residual)
         pt_pred = {
-            key: torch.as_tensor(value, device="cpu")
+            key: torch.as_tensor(value, device=env.DEVICE)
             if isinstance(value, np.ndarray)
             else value
             for key, value in pred.items()
         }
         pt_label = {
-            key: torch.as_tensor(value, device="cpu")
+            key: torch.as_tensor(value, device=env.DEVICE)
             if isinstance(value, np.ndarray)
             else value
             for key, value in label.items()

From ac9583a0491ab2ae5318159cb71b25b8075ce511 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 31 Jul 2026 14:19:54 +0000
Subject: [PATCH 12/12] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/common/dpmodel/test_loss_padding.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py
index 2b0e71fcb5..337c325fdc 100644
--- a/source/tests/common/dpmodel/test_loss_padding.py
+++ b/source/tests/common/dpmodel/test_loss_padding.py
@@ -856,7 +856,9 @@ def test_pt_placeholder_pairs_are_excluded_from_metrics(self):
 
         torch = pytest.importorskip("torch")
         from deepmd.pt.loss.ener import EnergyStdLoss as PTEnergyStdLoss
-        from deepmd.pt.utils import env
+        from deepmd.pt.utils import (
+            env,
+        )
 
         pred, label = _full_ener_dicts(
             1,
@@ -998,7 +1000,9 @@ def test_pt_mae_loss_func_uses_l1_for_the_hessian(self):
 
         torch = pytest.importorskip("torch")
         from deepmd.pt.loss.ener import EnergyStdLoss as PTEnergyStdLoss
-        from deepmd.pt.utils import env
+        from deepmd.pt.utils import (
+            env,
+        )
 
         residual = 10.0
         pred, label = self._constant_residual_dicts(residual)