feat(pt_expt): graph-native Hessian + fix fparam graph export/train (OutisLi P1s)

Two P1s from OutisLi's review of the graph-eligible DPA2 default-flip:

1. fparam graph export/train crash: dp_atomic_model.forward_atomic_graph called
   frame_id_from_n_node(graph.n_node) without n_total, so it fell back to
   int(sum(n_node)) -- an int() on a traced tensor that make_fx/torch.export
   cannot evaluate (GuardOnDataDependentSymNode), breaking both the graph .pt2
   export and compiled training for any numb_fparam>0 model. Pass the static
   flat node count n_total=atype.shape[0]. Regression:
   test_graph_lower_fparam_symbolic_trace_and_compile (red before, green after).

2. Hessian on the graph route: rather than fall back to dense, the graph route
   now computes the Hessian NATIVELY. The Hessian is just
   autograd.functional.hessian of the reduced energy w.r.t. coords -- route-
   agnostic. Added _WrapperForwardEnergyGraph + _cal_hessian_ext_graph and a
   Hessian loop in _call_common_graph (parallel to the dense forward_common_
   atomic loop), differentiating w.r.t. LOCAL coords by rebuilding the carry-all
   graph inside the wrapper (no extended nall->nloc fold needed -- the graph
   reduces over owned nodes). Output matches the dense route bit-tight in the
   same (nf, nloc*3, nloc*3) layout. Graph-builder dispatch factored into the
   shared _build_graph_for_method helper (one owner for _call_common_graph and
   the Hessian wrapper). Regression: test_graph_native_hessian_matches_dense
   (graph == dense Hessian at 1e-9, plus symmetry). Eager-only, like the dense
   Hessian. Benefits all graph-eligible descriptors, not just dpa2.

Author: Han Wang

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -308,7 +308,12 @@ def forward_atomic_graph(
         )
         fparam_node = None
         if fparam is not None:
-            frame_id = frame_id_from_n_node(graph.n_node)
+            # Pass the STATIC flat node count (``atype.shape[0] == N``) so the
+            # helper does not fall back to ``int(sum(n_node))``: that int() on a
+            # traced tensor breaks make_fx / torch.export
+            # (``GuardOnDataDependentSymNode``) for the graph .pt2 export and
+            # compiled-training paths when ``numb_fparam > 0``.
+            frame_id = frame_id_from_n_node(graph.n_node, n_total=atype.shape[0])
             fparam_node = xp.take(fparam, frame_id, axis=0)  # (N, ndf)
         return self.fitting_net.call_graph(
             gg, atype, gr=rot_mat, g2=None, h2=None, fparam=fparam_node, aparam=aparam

deepmd/pt_expt/model/make_model.py

@@ -191,6 +191,154 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor:
         return energy_redu
 
 
+def _build_graph_for_method(
+    method: str,
+    coord: torch.Tensor,
+    atype: torch.Tensor,
+    box: torch.Tensor | None,
+    rcut: float,
+    pair_excl: Any,
+) -> Any:
+    """Build a carry-all ``NeighborGraph`` for the named pt_expt builder.
+
+    Single owning site for the graph-builder dispatch shared by
+    :meth:`_call_common_graph` and the graph Hessian wrapper
+    (:class:`_WrapperForwardEnergyGraph`), so both build the graph identically.
+    """
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        build_neighbor_graph,
+        build_neighbor_graph_ase,
+    )
+
+    if method == "dense":
+        return build_neighbor_graph(coord, atype, box, rcut, pair_excl=pair_excl)
+    if method == "ase":
+        return build_neighbor_graph_ase(coord, atype, box, rcut, pair_excl=pair_excl)
+    if method == "vesin":
+        from deepmd.pt_expt.utils.vesin_graph_builder import (
+            build_neighbor_graph_vesin,
+        )
+
+        return build_neighbor_graph_vesin(coord, atype, box, rcut, pair_excl=pair_excl)
+    if method == "nv":
+        from deepmd.pt_expt.utils.nv_graph_builder import (
+            build_neighbor_graph_nv,
+        )
+
+        return build_neighbor_graph_nv(coord, atype, box, rcut, pair_excl=pair_excl)
+    raise ValueError(
+        f"unknown neighbor_graph_method {method!r}; use 'dense', 'ase', "
+        "'vesin', or 'nv'"
+    )
+
+
+class _WrapperForwardEnergyGraph:
+    """Graph twin of :class:`_WrapperForwardEnergy` for the Hessian.
+
+    Given flattened LOCAL coordinates for one frame, rebuilds the carry-all
+    ``NeighborGraph`` (so ``edge_vec`` tracks the coordinates through
+    ``autograd.functional.hessian``'s double-backward) and returns the scalar
+    reduced-output component. Unlike the dense wrapper this differentiates
+    w.r.t. the ``(nloc, 3)`` LOCAL coordinates directly -- the carry-all graph
+    has no ghost nodes (PBC images enter only as edges, rebuilt from the same
+    coords each call), so there is no extended-region ``nall -> nloc`` fold.
+    """
+
+    def __init__(
+        self,
+        model: Any,
+        kk: str,
+        ci: int,
+        nloc: int,
+        atype: torch.Tensor,  # (1, nloc)
+        box: torch.Tensor | None,  # (1, ...) or None
+        method: str,
+        pair_excl: Any,
+        rcut: float,
+        fparam: torch.Tensor | None,  # (1, ndf) or None
+        aparam: torch.Tensor | None,  # (1, nloc, nda) or None
+    ) -> None:
+        self.model = model
+        self.kk = kk
+        self.ci = ci
+        self.nloc = nloc
+        self.atype = atype
+        self.box = box
+        self.method = method
+        self.pair_excl = pair_excl
+        self.rcut = rcut
+        self.fparam = fparam
+        self.aparam = aparam
+
+    def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor:
+        cc = coord_flat.reshape(1, self.nloc, 3)
+        ng = _build_graph_for_method(
+            self.method, cc, self.atype, self.box, self.rcut, self.pair_excl
+        )
+        atomic_ret = self.model.atomic_model.forward_common_atomic_graph(
+            ng,
+            self.atype.reshape(-1),
+            fparam=self.fparam,
+            aparam=self.aparam,
+        )
+        # atomic_ret[kk]: flat (N, *def), N == nloc for a single-frame carry-all
+        # graph (all nodes owned); reduced output = sum over the node axis.
+        atom_out = atomic_ret[self.kk]
+        return atom_out.sum(dim=0).reshape(-1)[self.ci]
+
+
+def _cal_hessian_ext_graph(
+    model: Any,
+    kk: str,
+    vdef: OutputVariableDef,
+    coord: torch.Tensor,
+    atype: torch.Tensor,
+    box: torch.Tensor | None,
+    fparam: torch.Tensor | None,
+    aparam: torch.Tensor | None,
+    method: str,
+    pair_excl: Any,
+    rcut: float,
+    create_graph: bool = False,
+) -> torch.Tensor:
+    """Graph twin of :func:`_cal_hessian_ext`.
+
+    Computes the Hessian of the reduced output w.r.t. the LOCAL coordinates on
+    the carry-all graph route. Returns shape ``[nf, *vdef.shape, nloc*3,
+    nloc*3]`` -- the local-only counterpart of the dense extended Hessian,
+    already in the same final layout the dense route reaches after
+    ``communicate_extended_output`` folds ``nall -> nloc`` (the graph route
+    reduces over owned nodes, so no fold is needed). Node axis is
+    atom-major/xyz-minor, matching the dense final reshape.
+    """
+    nf, nloc, _ = coord.shape
+    vsize = math.prod(vdef.shape)
+    coord_flat = coord.reshape(nf, nloc * 3)
+    hessians = []
+    for ii in range(nf):
+        for ci in range(vsize):
+            wrapper = _WrapperForwardEnergyGraph(
+                model,
+                kk,
+                ci,
+                nloc,
+                atype[ii : ii + 1],
+                box[ii : ii + 1] if box is not None else None,
+                method,
+                pair_excl,
+                rcut,
+                fparam[ii : ii + 1] if fparam is not None else None,
+                aparam[ii : ii + 1] if aparam is not None else None,
+            )
+            hess = torch.autograd.functional.hessian(
+                wrapper,
+                coord_flat[ii],
+                create_graph=create_graph,
+            )  # (nloc*3, nloc*3)
+            hessians.append(hess)
+    return torch.stack(hessians).reshape(nf, *vdef.shape, nloc * 3, nloc * 3)
+
+
 def make_model(
     T_AtomicModel: type[BaseAtomicModel],
     T_Bases: tuple[type, ...] = (),
@@ -476,11 +624,6 @@ def _call_common_graph(
                 ``energy_redu``, ``energy_derv_r``, ``energy_derv_c_redu``, and
                 ``energy_derv_c`` when ``do_atomic_virial``).
             """
-            from deepmd.dpmodel.utils.neighbor_graph import (
-                build_neighbor_graph,
-                build_neighbor_graph_ase,
-            )
-
             # mirror the dpmodel guard: _resolve_graph_method's eligibility
             # check only protects the default (None) path; an EXPLICIT
             # neighbor_graph_method would otherwise reach the builders for
@@ -498,29 +641,7 @@ def _call_common_graph(
             # exported lower (forward_common_atomic_graph, which no longer
             # re-applies it) consumes a pre-excluded edge_mask.
             pair_excl = getattr(self.atomic_model, "pair_excl", None)
-            if method == "dense":
-                ng = build_neighbor_graph(cc, atype, bb, rcut, pair_excl=pair_excl)
-            elif method == "ase":
-                ng = build_neighbor_graph_ase(cc, atype, bb, rcut, pair_excl=pair_excl)
-            elif method == "vesin":
-                from deepmd.pt_expt.utils.vesin_graph_builder import (
-                    build_neighbor_graph_vesin,
-                )
-
-                ng = build_neighbor_graph_vesin(
-                    cc, atype, bb, rcut, pair_excl=pair_excl
-                )
-            elif method == "nv":
-                from deepmd.pt_expt.utils.nv_graph_builder import (
-                    build_neighbor_graph_nv,
-                )
-
-                ng = build_neighbor_graph_nv(cc, atype, bb, rcut, pair_excl=pair_excl)
-            else:
-                raise ValueError(
-                    f"unknown neighbor_graph_method {method!r}; "
-                    "use 'dense', 'ase', 'vesin', or 'nv'"
-                )
+            ng = _build_graph_for_method(method, cc, atype, bb, rcut, pair_excl)
             nf, nloc = atype.shape[:2]
             atype_flat = atype.reshape(nf * nloc)
             model_predict = self.forward_common_lower_graph(
@@ -548,6 +669,30 @@ def _call_common_graph(
                     and v.shape[:1] == torch.Size([N])
                 ):
                     model_predict[k] = v.reshape(nf, nloc, *v.shape[1:])
+            # Graph-native Hessian (parallel to the dense ``forward_common_atomic``
+            # loop): differentiate the reduced output w.r.t. the LOCAL coords by
+            # rebuilding the graph inside the wrapper. Added AFTER the unravel so
+            # its ``(nf, *def, nloc, 3, nloc, 3)`` shape is returned as-is.
+            # Eager-only, like the dense Hessian (autograd.functional.hessian
+            # does not export/compile).
+            aod = self.atomic_output_def()
+            for kk in aod.keys():
+                vdef = aod[kk]
+                if vdef.reducible and vdef.r_hessian:
+                    model_predict[get_hessian_name(kk)] = _cal_hessian_ext_graph(
+                        self,
+                        kk,
+                        vdef,
+                        cc,
+                        atype,
+                        bb,
+                        fp,
+                        ap,
+                        method,
+                        pair_excl,
+                        rcut,
+                        create_graph=self.training,
+                    )
             return model_predict
 
         def forward_common_atomic(

source/tests/pt_expt/model/test_dpa2_graph_lower.py

@@ -176,6 +176,7 @@ def _make_model(
         repformer_nsel: int = 150,
         repinit_nsel: int = 200,
         repformer_attn: bool = False,
+        numb_fparam: int = 0,
     ) -> EnergyModel:
         ds = _make_dpa2_descriptor(
             ntypes=self.nt,
@@ -189,6 +190,7 @@ def _make_model(
             ds.get_dim_out(),
             1,
             mixed_types=ds.mixed_types(),
+            numb_fparam=numb_fparam,
             precision="float64",
             seed=GLOBAL_SEED,
         ).to(self.device)
@@ -224,6 +226,61 @@ def test_force_virial_parity_vs_legacy(self) -> None:
             graph["energy_derv_c_redu"], legacy["energy_derv_c_redu"], **tol
         )
 
+    def test_graph_native_hessian_matches_dense(self) -> None:
+        """The graph route computes the Hessian natively and it matches the
+        dense route bit-tight.
+
+        A Hessian-enabled graph-eligible DPA2 model default-flips to the graph
+        route; ``_call_common_graph`` now runs its own ``_cal_hessian_ext_graph``
+        loop (differentiating the reduced energy w.r.t. the LOCAL coords by
+        rebuilding the carry-all graph inside the autograd wrapper), so
+        ``energy_derv_r_derv_r`` is produced without falling back to dense.
+        The result must equal the dense route (forced via the
+        ``disable_graph_lower`` escape hatch on the same weights) at fp64
+        parity, in the same ``(nf, 1, nloc*3, nloc*3)`` layout.
+
+        A non-binding repformer sel is used so graph and dense agree on the
+        first derivatives too (attention off -- the fixture default); the
+        Hessian parity would otherwise inherit the binding-sel divergence.
+        """
+        from deepmd.pt_expt.train.training import (
+            _model_uses_graph_lower,
+        )
+
+        model = self._make_model()  # non-binding sel, attention off
+        model.eval()
+        assert model.atomic_model.descriptor.uses_graph_lower() is True
+        model.enable_hessian()
+        # a Hessian model stays graph-routed: the graph now produces the Hessian.
+        assert _model_uses_graph_lower(model) is True
+
+        box = self.cell.reshape(1, 9)
+        # ``EnergyModel.forward`` exposes the Hessian under the ``"hessian"``
+        # key (translated from ``energy_derv_r_derv_r``); it would KeyError if
+        # the graph route failed to produce it.
+        graph_out = model.forward(
+            self.coord.clone().requires_grad_(True), self.atype, box=box
+        )
+        assert "hessian" in graph_out
+        assert graph_out["hessian"].shape == (1, self.natoms * 3, self.natoms * 3)
+
+        # dense reference on the SAME weights via the escape hatch.
+        ref_model = self._make_model()
+        ref_model.eval()
+        ref_model.load_state_dict(model.state_dict(), strict=False)
+        ref_model.atomic_model.descriptor.disable_graph_lower()
+        ref_model.enable_hessian()
+        assert _model_uses_graph_lower(ref_model) is False
+        dense_out = ref_model.forward(
+            self.coord.clone().requires_grad_(True), self.atype, box=box
+        )
+        torch.testing.assert_close(
+            graph_out["hessian"], dense_out["hessian"], rtol=1e-9, atol=1e-9
+        )
+        # Hessian symmetry (a genuine second derivative, not a shape artifact).
+        h = graph_out["hessian"][0]
+        torch.testing.assert_close(h, h.transpose(-1, -2), rtol=1e-9, atol=1e-9)
+
     def test_disable_graph_lower_escape_hatch(self) -> None:
         """``descriptor.disable_graph_lower()`` is the documented legacy-dense
         escape hatch: it flips ``uses_graph_lower()`` to ``False`` so the
@@ -495,6 +552,63 @@ def test_compiled_training_graph_smoke(self) -> None:
             **tol,
         )
 
+    def test_graph_lower_fparam_symbolic_trace_and_compile(self) -> None:
+        """A graph-eligible DPA2 model with ``numb_fparam > 0`` must export
+        (``make_fx`` symbolic) AND inductor-compile.
+
+        Regression for the ``frame_id_from_n_node(graph.n_node)`` call in
+        ``dp_atomic_model.forward_atomic_graph``: without a STATIC ``n_total``
+        it fell back to ``int(sum(n_node))``, which make_fx / torch.export
+        cannot evaluate on a traced tensor (``GuardOnDataDependentSymNode``),
+        so both the graph ``.pt2`` export and compiled training failed for any
+        fparam model. Both must now trace/compile and match eager bit-tight.
+        """
+        from deepmd.pt_expt.train.training import (
+            _trace_and_compile_graph,
+        )
+        from deepmd.pt_expt.utils.serialization import (
+            build_synthetic_graph_inputs,
+        )
+
+        model = self._make_model(numb_fparam=2).to("cpu")
+        model.eval()
+        assert model.atomic_model.descriptor.uses_graph_lower() is True
+
+        sample = build_synthetic_graph_inputs(
+            model,
+            e_max=175,
+            nframes=2,
+            nloc=7,
+            dtype=torch.float64,
+            device=torch.device("cpu"),
+        )
+        atype, n_node, ei, ev, em, fp, ap, cs = sample
+        assert fp is not None and fp.shape[-1] == 2, "fparam must be present"
+
+        # (a) symbolic-trace export path
+        traced = model.forward_lower_graph_exportable(
+            atype, n_node, ei, ev, em,
+            fparam=fp, aparam=ap, do_atomic_virial=True,
+            charge_spin=cs, tracing_mode="symbolic", _allow_non_fake_inputs=True,
+        )
+        out = traced(atype, n_node, ei, ev, em, fp, ap, cs)
+        ref = model.forward_common_lower_graph(
+            atype, n_node, ei, ev, em, fparam=fp, aparam=ap, do_atomic_virial=True
+        )
+        tol = {"rtol": 1e-12, "atol": 1e-12}
+        torch.testing.assert_close(out["energy"], ref["energy_redu"], **tol)
+        torch.testing.assert_close(
+            out["force"], ref["energy_derv_r"].reshape(out["force"].shape), **tol
+        )
+
+        # (b) compiled-training path (fparam threaded through the compile)
+        compiled_lower, _ = _trace_and_compile_graph(model, fp, None, None)
+        compiled_out = compiled_lower(atype, n_node, ei, ev, em, fp, ap, cs)
+        ctol = {"rtol": 1e-10, "atol": 1e-10}
+        torch.testing.assert_close(
+            compiled_out["energy"], ref["energy_redu"], **ctol
+        )
+
     # ------------------------------------------------------------------
     # 2. the one piece of real code: the border_op graph exchange override.
     # ------------------------------------------------------------------