From 378110e8c5e96dc12da09db3b1eeb6b3208401fc Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 21 Jul 2026 12:19:49 +0800 Subject: [PATCH 1/2] fix(pt_expt): preserve graph routing for raw checkpoints Select the eager lower ABI from the restored model capability instead of forcing every raw checkpoint through the padded dense neighbor-list path. Graph-eligible energy models now reuse the graph DeepEval contract, keeping energy, force, virial, and atomic outputs aligned with public forward semantics when descriptor statistics are nonzero. Retain the existing dense and spin paths, expose graph builder selection for raw graph checkpoints, and cover both plain and compiled checkpoint layouts with a deterministic DPA1 regression. --- deepmd/pt_expt/infer/deep_eval.py | 154 ++++++++++++++---- .../infer/test_deep_eval_pt_checkpoint.py | 109 +++++++++++++ 2 files changed, 229 insertions(+), 34 deletions(-) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 1b97ef7f50..cebd645957 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +import warnings from collections.abc import ( Callable, ) @@ -76,11 +77,11 @@ ) -# Public output keys emitted by the graph-form AOTI forward -# (``forward_lower_graph_exportable``) keyed by the output-variable category that -# ``request_defs`` carries. The graph path is LOCAL-only (``N == sum(n_node)`` -# nodes, no ghosts), so its outputs are already at local-atom resolution -- no -# ``communicate_extended_output`` fold-back is needed. +# Public output keys emitted by graph-lower forwards, keyed by the +# output-variable category that ``request_defs`` carries. The graph path is +# local-only (``N == sum(n_node)`` nodes, no ghosts), so its outputs are already +# at local-atom resolution and require no ``communicate_extended_output`` +# fold-back. _GRAPH_CATEGORY_TO_KEY = { OutputVariableCategory.OUT: "atom_energy", OutputVariableCategory.REDU: "energy", @@ -133,18 +134,21 @@ class DeepEval(DeepEvalBackend): If True, automatic batch size will be used. If int, it will be used as the initial batch size. neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional - The ASE neighbor list class to produce the neighbor list. If None, the - neighbor list will be built natively in the model. + The ASE neighbor list class for nlist-routed artifacts. If None, the + neighbor list will be built natively. Explicit neighbor lists are + rejected for graph-routed artifacts. nlist_backend : str, default: "auto" - Neighbor-list builder for the NLIST/extended lower path (``.pte`` and - nlist-form ``.pt2``): ``"auto"`` / ``"vesin"`` / ``"native"``. Not - used by graph-form ``.pt2`` artifacts. + Neighbor-list builder for the NLIST/extended lower path (``.pte``, + nlist-form ``.pt2``, and dense-routed ``.pt`` checkpoints): + ``"auto"`` / ``"vesin"`` / ``"native"``. Explicit non-default values + are ignored with a warning for graph-routed artifacts. neighbor_graph_method : str, default: "dense" - Carry-all graph builder for GRAPH-FORM ``.pt2`` artifacts ONLY + Carry-all graph builder for graph-form ``.pt2`` artifacts and + graph-routed ``.pt`` checkpoints (``metadata["lower_input_kind"] == "graph"``): ``"dense"`` / ``"ase"`` (backend-agnostic) or ``"vesin"`` / ``"nv"`` (on-device O(N)). A - non-default value on any other artifact raises at construction — the - knob would silently do nothing there; use ``nlist_backend`` for the + non-default value on any other artifact raises at construction because + the knob would silently do nothing there; use ``nlist_backend`` for the nlist path instead. All builders emit the same neighbor set, so the choice is performance-only. Consolidating the two knobs into a single backend-selection API is deferred to the dense-nlist deprecation. @@ -166,8 +170,8 @@ def __init__( self.output_def = output_def self.model_path = model_file self.neighbor_list = neighbor_list - # World-2 graph-form ``.pt2`` (lower_input_kind == "graph") builder select: - # "dense"/"ase" (backend-agnostic) or "vesin"/"nv" (on-device O(N)). + # Graph-lower builder selection: "dense"/"ase" (backend-agnostic) or + # "vesin"/"nv" (on-device O(N)). self._neighbor_graph_method = neighbor_graph_method self._is_pt2 = model_file.endswith(".pt2") @@ -184,15 +188,15 @@ def __init__( "`.pt` (training checkpoint)." ) - # neighbor_graph_method is consumed ONLY by graph-form .pt2 eval - # (_eval_model_graph); fail fast instead of silently ignoring it on - # nlist-form artifacts (there, the builder knob is nlist_backend). + # ``neighbor_graph_method`` is consumed only by graph-lower evaluation. + # Fail fast instead of silently ignoring it on nlist-form artifacts, + # where the corresponding builder knob is ``nlist_backend``. if neighbor_graph_method != "dense" and getattr(self, "metadata", {}).get( "lower_input_kind" ) not in ("graph", "dpa1_canonical"): raise ValueError( f"neighbor_graph_method={neighbor_graph_method!r} only applies to " - "graph-form .pt2 artifacts (lower_input_kind == 'graph'); this " + "graph-routed artifacts (lower_input_kind == 'graph'); this " f"model is not graph-form. Use nlist_backend to select the " "neighbor-list builder for the nlist path." ) @@ -217,14 +221,33 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None: ``"native"`` uses the dense all-pairs builder; ``"vesin"`` forces the O(N) ``vesin.torch`` cell list (raising if it is unavailable or the model/inputs are unsupported); ``"auto"`` uses vesin when applicable and - silently falls back to the native builder otherwise. Results are - unchanged either way -- only the neighbor-search cost differs. + silently falls back to the native builder otherwise. Graph-routed + artifacts use ``neighbor_graph_method`` instead, reject an explicit ASE + neighbor list, and warn when an explicit nlist backend is ignored. + Results are unchanged either way; only the neighbor-search cost differs. """ if nlist_backend not in ("auto", "vesin", "native"): raise ValueError( f"Unknown nlist_backend '{nlist_backend}'; " "expected 'auto', 'vesin', or 'native'." ) + if self.metadata.get("lower_input_kind") in ("graph", "dpa1_canonical"): + if self.neighbor_list is not None: + raise ValueError( + "neighbor_list only applies to nlist-routed artifacts; use " + "neighbor_graph_method for this graph-routed model." + ) + if nlist_backend != "auto": + warnings.warn( + f"nlist_backend={nlist_backend!r} is ignored for graph-routed " + "artifacts; use neighbor_graph_method to select the graph " + "builder.", + UserWarning, + stacklevel=2, + ) + self._use_vesin = False + self._nlist_builder = None + return is_spin = bool(getattr(self, "_is_spin", False)) ase_provided = self.neighbor_list is not None # reason vesin cannot be used (None means it can) @@ -519,6 +542,11 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: wrapper.load_state_dict(state_dict) model = wrapper.model["Default"].eval() + from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower, + ) + + use_graph_lower = model_uses_graph_lower(model) self._dpmodel = model self._is_spin = ( model_params.get("type") == "spin_ener" or "spin" in model_params @@ -578,19 +606,77 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: else None ), "is_spin": self._is_spin, - "lower_input_kind": "nlist", + "lower_input_kind": "graph" if use_graph_lower else "nlist", } + if use_graph_lower: + from deepmd.pt_expt.utils.serialization import ( + _graph_edge_dtype, + ) + + self.metadata["graph_edge_dtype"] = _graph_edge_dtype(model, "graph") if self._is_spin: self.metadata["ntypes_spin"] = model.spin.get_ntypes_spin() self.metadata["use_spin"] = [bool(v) for v in model.spin.use_spin] - # Eager runner with the same signature as the .pt2/.pte exported module. - # Use forward_common_lower (not forward_lower) to match the export-time - # output keys ("energy", "energy_redu", "energy_derv_r", ...) that - # communicate_extended_output downstream consumes. + # Eager runners use the same ABI as the corresponding exported lower. + # Graph-eligible checkpoints preserve the source model's default + # graph-forward semantics; all other checkpoints use the dense lower. + # The dense runner emits internal keys consumed by + # ``communicate_extended_output``, while the graph runner emits local + # public keys consumed directly by ``_eval_model_graph``. + # # Non-spin: (ext_coord, ext_atype, nlist, mapping, fparam, aparam) # Spin: (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) - if self._is_spin: + if use_graph_lower: + from deepmd.pt_expt.model.ener_model import ( + _translate_energy_keys, + ) + + do_grad_r = model.do_grad_r("energy") + do_grad_c = model.do_grad_c("energy") + + def _eager_runner_graph( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + model_ret = model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=True, + do_atomic_virial=True, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + return _translate_energy_keys( + model_ret, + do_grad_r=do_grad_r, + do_grad_c=do_grad_c, + do_atomic_virial=True, + local=True, + ) + + self.exported_module = _eager_runner_graph + elif self._is_spin: def _eager_runner_spin( ext_coord: torch.Tensor, @@ -1761,18 +1847,18 @@ def _eval_model_graph( request_defs: list[OutputVariableDef], charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: - """Evaluate a graph-form ``.pt2`` (``lower_input_kind == "graph"``). + """Evaluate a graph-lower model (``lower_input_kind == "graph"``). Builds a carry-all :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph` from the eval system at its exact (tight) edge count and feeds the positional schema ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, - fparam, aparam, charge_spin)`` to the exported forward. The AOTI - artifact's edge axis is dynamic, so no ``edge_capacity`` padding is needed. The - ``graph_edge_dtype`` metadata selects float32 geometry for compressed - DPA1 and float64 for generic graph descriptors. The forward returns the - LOCAL public keys directly, so results are reshaped without + fparam, aparam, charge_spin)`` to the graph forward. Exported AOTI + artifacts use a dynamic edge axis, so no ``edge_capacity`` padding is + needed. The ``graph_edge_dtype`` metadata selects float32 geometry for + compressed DPA1 and float64 for generic graph descriptors. The forward + returns local public keys directly, so results are reshaped without ``communicate_extended_output``. """ from deepmd.pt_expt.utils.env import ( @@ -1928,7 +2014,7 @@ def _build_eval_graph( box_input: np.ndarray | None, device: "torch.device", ) -> "NeighborGraph": - """Build the carry-all NeighborGraph for graph-form ``.pt2`` inference. + """Build the carry-all NeighborGraph for graph-lower inference. Dispatches on ``self._neighbor_graph_method``: ``dense``/``ase`` run backend-agnostic (numpy); ``vesin``/``nv`` run on-device (torch, O(N)). diff --git a/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py b/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py index 6a9b7e2c59..e1d9106135 100644 --- a/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py +++ b/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py @@ -40,6 +40,10 @@ from deepmd.pt_expt.infer.deep_eval import DeepEval as PtExptDeepEval from deepmd.pt_expt.model import ( EnergyModel, + get_model, +) +from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower, ) from deepmd.pt_expt.train.wrapper import ( ModelWrapper, @@ -110,6 +114,39 @@ def _build_model_and_params( return model, model_params +def _build_graph_dpa1_model_and_params() -> tuple[EnergyModel, dict]: + """Build a graph-routed DPA1 model with nonzero descriptor statistics.""" + model_params = { + "type_map": ["H", "O"], + "descriptor": { + "type": "dpa1", + "sel": 20, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [3, 6], + "axis_neuron": 2, + "attn": 4, + "attn_layer": 0, + "smooth_type_embedding": True, + "set_davg_zero": False, + "type_one_side": True, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "ener", + "neuron": [8, 8], + "precision": "float64", + "seed": 1, + }, + } + model = get_model(copy.deepcopy(model_params)).to(torch.float64).to(DEVICE).eval() + with torch.no_grad(): + model.atomic_model.descriptor.se_atten.mean.fill_(0.01) + model.atomic_model.descriptor.se_atten.stddev.fill_(0.1) + return model, model_params + + def _save_pt_checkpoint( model: EnergyModel, model_params: dict, @@ -296,6 +333,78 @@ def test_unsupported_extension_raises(self) -> None: os.unlink(bogus) +class TestPtExptLoadPtGraphDPA1(unittest.TestCase): + """Raw DPA1 checkpoints retain the source model's graph-forward semantics.""" + + @classmethod + def setUpClass(cls) -> None: + cls.model, cls.model_params = _build_graph_dpa1_model_and_params() + cls.pt_paths = { + "plain": tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name, + "compiled": tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name, + } + _save_pt_checkpoint(cls.model, cls.model_params, cls.pt_paths["plain"]) + _save_pt_checkpoint_compiled( + cls.model, + cls.model_params, + cls.pt_paths["compiled"], + ) + + @classmethod + def tearDownClass(cls) -> None: + for path in cls.pt_paths.values(): + if os.path.exists(path): + os.unlink(path) + + def test_eval_matches_public_forward(self) -> None: + self.assertTrue(model_uses_graph_lower(self.model)) + + coords = np.array([[[1.0, 1.0, 1.0], [2.0, 1.0, 1.0], [1.0, 2.0, 1.0]]]) + cells = np.eye(3).reshape(1, 9) * 10.0 + atom_types = np.array([0, 1, 0], dtype=np.int32) + + output_names = ( + "energy", + "force", + "virial", + "atom_energy", + "atom_virial", + ) + coord_t = torch.tensor( + coords, dtype=torch.float64, device=DEVICE + ).requires_grad_(True) + atype_t = torch.tensor( + atom_types.reshape(1, -1), dtype=torch.int64, device=DEVICE + ) + cell_t = torch.tensor(cells, dtype=torch.float64, device=DEVICE) + expected = self.model.forward( + coord_t, + atype_t, + cell_t, + do_atomic_virial=True, + ) + + for layout, path in self.pt_paths.items(): + with self.subTest(layout=layout): + dp = DeepPot(path, auto_batch_size=False) + self.assertEqual(dp.deep_eval.metadata["lower_input_kind"], "graph") + actual = dict( + zip( + output_names, + dp.eval(coords, cells, atom_types, atomic=True), + strict=True, + ) + ) + for name in output_names: + np.testing.assert_allclose( + actual[name], + expected[name].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + + class TestPtExptLoadPtCompiledLayout(unittest.TestCase): """`.pt` saved after pt_expt training compilation (`_CompiledModel` wrap). From 7f73a80b544567a9bff96aa9de4974a89dfd385c Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 21 Jul 2026 13:41:11 +0800 Subject: [PATCH 2/2] perf(pt_expt): auto-select scalable graph builders MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolve the graph neighbor backend once when DeepEval loads the model. Prefer the batched nvalchemiops builder on CUDA, fall back to Vesin when available, and retain the dense all-pairs implementation as the dependency-free final fallback. This removes the single-core NumPy O(N²) graph construction bottleneck from the default dp test path while preserving every explicit builder selection. Cover the resolved backend in the raw DPA1 checkpoint regression. --- deepmd/pt_expt/infer/deep_eval.py | 55 ++++++++++++++----- .../infer/test_deep_eval_pt_checkpoint.py | 16 ++++++ 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index cebd645957..5881cc132c 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -142,12 +142,14 @@ class DeepEval(DeepEvalBackend): nlist-form ``.pt2``, and dense-routed ``.pt`` checkpoints): ``"auto"`` / ``"vesin"`` / ``"native"``. Explicit non-default values are ignored with a warning for graph-routed artifacts. - neighbor_graph_method : str, default: "dense" + neighbor_graph_method : str, default: "auto" Carry-all graph builder for graph-form ``.pt2`` artifacts and graph-routed ``.pt`` checkpoints - (``metadata["lower_input_kind"] == "graph"``): ``"dense"`` / ``"ase"`` - (backend-agnostic) or ``"vesin"`` / ``"nv"`` (on-device O(N)). A - non-default value on any other artifact raises at construction because + (``metadata["lower_input_kind"] == "graph"``): ``"auto"`` selects + ``"nv"`` on CUDA when nvalchemiops is available, then ``"vesin"`` when + available, and finally falls back to ``"dense"``. Explicit + ``"dense"`` / ``"ase"`` / ``"vesin"`` / ``"nv"`` choices are preserved. + A non-default value on any other artifact raises at construction because the knob would silently do nothing there; use ``nlist_backend`` for the nlist path instead. All builders emit the same neighbor set, so the choice is performance-only. Consolidating the two knobs into a single @@ -164,14 +166,14 @@ def __init__( auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, nlist_backend: str = "auto", - neighbor_graph_method: str = "dense", + neighbor_graph_method: str = "auto", **kwargs: Any, ) -> None: self.output_def = output_def self.model_path = model_file self.neighbor_list = neighbor_list - # Graph-lower builder selection: "dense"/"ase" (backend-agnostic) or - # "vesin"/"nv" (on-device O(N)). + # Graph-lower builder selection is resolved once after model metadata + # identifies the lower ABI. self._neighbor_graph_method = neighbor_graph_method self._is_pt2 = model_file.endswith(".pt2") @@ -191,9 +193,9 @@ def __init__( # ``neighbor_graph_method`` is consumed only by graph-lower evaluation. # Fail fast instead of silently ignoring it on nlist-form artifacts, # where the corresponding builder knob is ``nlist_backend``. - if neighbor_graph_method != "dense" and getattr(self, "metadata", {}).get( - "lower_input_kind" - ) not in ("graph", "dpa1_canonical"): + if neighbor_graph_method not in ("auto", "dense") and getattr( + self, "metadata", {} + ).get("lower_input_kind") not in ("graph", "dpa1_canonical"): raise ValueError( f"neighbor_graph_method={neighbor_graph_method!r} only applies to " "graph-routed artifacts (lower_input_kind == 'graph'); this " @@ -201,7 +203,7 @@ def __init__( "neighbor-list builder for the nlist path." ) - self._setup_nlist_backend(nlist_backend) + self._setup_neighbor_backend(nlist_backend) if isinstance(auto_batch_size, bool): if auto_batch_size: @@ -215,8 +217,32 @@ def __init__( else: raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize") - def _setup_nlist_backend(self, nlist_backend: str) -> None: - """Resolve the neighbor-list construction strategy from a user choice. + @staticmethod + def _resolve_neighbor_graph_method(method: str) -> str: + """Resolve the graph builder once for the active device.""" + if method not in ("auto", "dense", "ase", "vesin", "nv"): + raise ValueError( + f"Unknown neighbor_graph_method {method!r}; " + "expected 'auto', 'dense', 'ase', 'vesin', or 'nv'." + ) + if method != "auto": + return method + + from deepmd.pt.utils.nv_nlist import ( + is_nv_available, + ) + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + if DEVICE.type == "cuda" and is_nv_available(): + return "nv" + if is_vesin_torch_available(): + return "vesin" + return "dense" + + def _setup_neighbor_backend(self, nlist_backend: str) -> None: + """Resolve the graph or neighbor-list construction strategy. ``"native"`` uses the dense all-pairs builder; ``"vesin"`` forces the O(N) ``vesin.torch`` cell list (raising if it is unavailable or the @@ -245,6 +271,9 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None: UserWarning, stacklevel=2, ) + self._neighbor_graph_method = self._resolve_neighbor_graph_method( + self._neighbor_graph_method + ) self._use_vesin = False self._nlist_builder = None return diff --git a/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py b/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py index e1d9106135..de93d8f926 100644 --- a/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py +++ b/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py @@ -31,6 +31,9 @@ from deepmd.infer import ( DeepPot, ) +from deepmd.pt.utils.nv_nlist import ( + is_nv_available, +) from deepmd.pt_expt.descriptor.se_e2_a import ( DescrptSeA, ) @@ -51,6 +54,9 @@ from deepmd.pt_expt.utils.env import ( DEVICE, ) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + is_vesin_torch_available, +) from ...seed import ( GLOBAL_SEED, @@ -339,6 +345,12 @@ class TestPtExptLoadPtGraphDPA1(unittest.TestCase): @classmethod def setUpClass(cls) -> None: cls.model, cls.model_params = _build_graph_dpa1_model_and_params() + if DEVICE.type == "cuda" and is_nv_available(): + cls.expected_graph_method = "nv" + elif is_vesin_torch_available(): + cls.expected_graph_method = "vesin" + else: + cls.expected_graph_method = "dense" cls.pt_paths = { "plain": tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name, "compiled": tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name, @@ -388,6 +400,10 @@ def test_eval_matches_public_forward(self) -> None: with self.subTest(layout=layout): dp = DeepPot(path, auto_batch_size=False) self.assertEqual(dp.deep_eval.metadata["lower_input_kind"], "graph") + self.assertEqual( + dp.deep_eval._neighbor_graph_method, + self.expected_graph_method, + ) actual = dict( zip( output_names,