From 368e940396fe22bba12870b6da8256ed380ad688 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 21 Jul 2026 12:19:49 +0800 Subject: [PATCH 1/5] 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 | 153 ++++++++++++++---- .../infer/test_deep_eval_pt_checkpoint.py | 109 +++++++++++++ 2 files changed, 228 insertions(+), 34 deletions(-) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 999e7f69f1..9cc198afe5 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -77,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", @@ -185,18 +185,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. @@ -218,8 +221,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") @@ -241,15 +244,15 @@ def __init__( # place that sees every model load regardless of archive kind. _warn_legacy_edge_vec(self.metadata) - # 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." ) @@ -274,14 +277,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)) # Native-spin (NeighborGraph route) graph-form artifacts never touch # this NLIST builder at all -- graph-form eval uses @@ -589,6 +611,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 @@ -648,19 +675,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, @@ -1977,18 +2062,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 ( @@ -2144,7 +2229,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 09e9442299c7cf59511099ac7b022909b12a4625 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 21 Jul 2026 13:41:11 +0800 Subject: [PATCH 2/5] 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 9cc198afe5..348fb0d5d5 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -193,12 +193,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 @@ -215,14 +217,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") @@ -247,9 +249,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 " @@ -257,7 +259,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: @@ -271,8 +273,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 @@ -301,6 +327,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, From 7c48a6c43c543f7a083ed4df81337547f881e585 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 28 Jul 2026 14:04:51 +0800 Subject: [PATCH 3/5] fix(pt_expt): preserve native-spin graph checkpoint ABI --- deepmd/pt_expt/infer/deep_eval.py | 51 +++++++++++++++++++++++++++++-- 1 file changed, 48 insertions(+), 3 deletions(-) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 348fb0d5d5..f79efed295 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -723,9 +723,54 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: # ``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 use_graph_lower: + # Graph: (..., source_row_ptr, fparam, aparam, charge_spin) + # Graph spin: (..., source_row_ptr, spin, fparam, aparam, charge_spin) + # Nlist: (ext_coord, ext_atype, nlist, mapping, fparam, aparam) + # Nlist spin: (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) + if use_graph_lower and self._is_spin: + + def _eager_runner_graph_spin( + 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, + spin: 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, + spin=spin, + ) + return model._translate_eager_call( + model_ret, + atype, + do_atomic_virial=True, + ) + + self.exported_module = _eager_runner_graph_spin + elif use_graph_lower: from deepmd.pt_expt.model.ener_model import ( _translate_energy_keys, ) From 8ef72a39199304eabce292bd7358342909388df8 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 28 Jul 2026 20:12:16 +0800 Subject: [PATCH 4/5] fix(pt_expt): distinguish raw checkpoint backends by weights --- deepmd/backend/pt_expt.py | 60 ++++++++++---- deepmd/pt_expt/infer/deep_eval.py | 32 +++----- .../infer/test_deep_eval_pt_checkpoint.py | 79 ++++++++++++++++++- 3 files changed, 135 insertions(+), 36 deletions(-) diff --git a/deepmd/backend/pt_expt.py b/deepmd/backend/pt_expt.py index 38b66f0104..7f24cfbec9 100644 --- a/deepmd/backend/pt_expt.py +++ b/deepmd/backend/pt_expt.py @@ -1,12 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from collections.abc import ( Callable, + Mapping, ) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, + Any, ClassVar, ) @@ -27,6 +29,46 @@ ) +def detect_pt_checkpoint_backend(checkpoint: Any) -> str | None: + """Detect the backend dialect of a raw PyTorch checkpoint. + + Parameters + ---------- + checkpoint : Any + A checkpoint payload or its unwrapped model state dictionary. + + Returns + ------- + str or None + ``"pt-expt"`` or ``"pt"`` when the parameter names identify one + backend unambiguously, otherwise ``None``. + """ + state_dict = checkpoint + if isinstance(state_dict, Mapping) and "model" in state_dict: + state_dict = state_dict["model"] + if not isinstance(state_dict, Mapping): + return None + + keys = tuple(key for key in state_dict if isinstance(key, str)) + + # Weight names are decisive. pt_expt DPA4 also contains ordinary + # torch-native ``.bias`` parameters, so bias names cannot override a + # clear ``.w`` versus ``.matrix`` distinction. + has_pt_expt_weight = any(key.endswith(".w") for key in keys) + has_pt_weight = any(key.endswith(".matrix") for key in keys) + if has_pt_expt_weight or has_pt_weight: + if has_pt_expt_weight == has_pt_weight: + return None + return "pt-expt" if has_pt_expt_weight else "pt" + + # Bias-only state dictionaries retain the original dialect fallback. + has_pt_expt_bias = any(key.endswith(".b") for key in keys) + has_pt_bias = any(key.endswith(".bias") for key in keys) + if has_pt_expt_bias == has_pt_bias: + return None + return "pt-expt" if has_pt_expt_bias else "pt" + + @Backend.register("pt-expt") @Backend.register("pytorch-exportable") class PyTorchExportableBackend(Backend): @@ -51,11 +93,8 @@ def match_filename(cls, filename: str) -> int: Returns ------- - 1 for the regular `.pte` / `.pt2` suffixes (default behaviour). - - 2 for `.pt` files whose state-dict uses pt_expt's dpmodel - parameter naming (`.w`/`.b`); this outranks the legacy pt - backend's default suffix score (1) so pt_expt-trained `.pt` - checkpoints route here, while genuine pt-trained `.pt` files - (which use `.matrix`/`.bias`) keep going to the pt backend. + - 2 for `.pt` files whose state dictionary uses the pt_expt parameter + dialect. This outranks the pt backend's default suffix score (1). - 0 otherwise. """ score = super().match_filename(filename) @@ -69,21 +108,14 @@ def match_filename(cls, filename: str) -> int: # weights_only=True avoids unpickling arbitrary code from an # untrusted .pt — sniffing only needs the dict keys. - sd = torch.load(filename, map_location="cpu", weights_only=True) + checkpoint = torch.load(filename, map_location="cpu", weights_only=True) except Exception: # Not a valid torch archive (corrupt file, wrong format, or a # weights_only=True restriction trip). Surrender the claim so # the dispatcher falls back to the default suffix match — pt's # default score (1) will pick up the file under `dp --pt`. return 0 - if isinstance(sd, dict) and "model" in sd: - sd = sd["model"] - keys = list(sd.keys()) if hasattr(sd, "keys") else [] - has_pt_expt = any(k.endswith(".w") or k.endswith(".b") for k in keys) - has_pt = any(k.endswith(".matrix") or k.endswith(".bias") for k in keys) - if has_pt_expt and not has_pt: - return 2 - return 0 + return 2 if detect_pt_checkpoint_backend(checkpoint) == "pt-expt" else 0 def is_available(self) -> bool: """Check if the backend is available. diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index f79efed295..57d40fe9f3 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -13,6 +13,9 @@ import numpy as np import torch +from deepmd.backend.pt_expt import ( + detect_pt_checkpoint_backend, +) from deepmd.dpmodel.model.transform_output import ( communicate_extended_output, ) @@ -132,18 +135,6 @@ def _reshape_charge_spin( ) from err -def _is_pt_backend_dpa4_params(model_params: dict[str, Any]) -> bool: - """Return whether a training checkpoint should be loaded by the pt backend.""" - model_type = str(model_params.get("type", "")).lower() - if model_type in {"sezm", "dpa4", "sezm_spin"}: - return True - descriptor = model_params.get("descriptor") - if isinstance(descriptor, dict): - descriptor_type = str(descriptor.get("type", "")).lower() - return descriptor_type in {"sezm", "dpa4"} - return False - - def _warn_legacy_edge_vec(metadata: dict) -> None: """Warn once per model load when an edge_vec-schema artifact is opened. @@ -556,9 +547,17 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: # Match the training resume path (training.py:712) — weights_only=True # avoids unpickling arbitrary code from untrusted checkpoints. - state_dict = torch.load(model_file, map_location=DEVICE, weights_only=True) + checkpoint = torch.load(model_file, map_location=DEVICE, weights_only=True) + checkpoint_backend = detect_pt_checkpoint_backend(checkpoint) + state_dict = checkpoint if isinstance(state_dict, dict) and "model" in state_dict: state_dict = state_dict["model"] + if checkpoint_backend == "pt": + raise ValueError( + f"Checkpoint '{model_file}' uses the regular `pt` parameter " + "dialect. Load it with `dp --pt`, or export it to `.pt2` / " + "`.pte` before loading it with `pt_expt`." + ) extra = state_dict.get("_extra_state") if isinstance(state_dict, dict) else None if not (isinstance(extra, dict) and "model_params" in extra): raise ValueError( @@ -598,13 +597,6 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: state_dict = head_state model_params = head_params - if _is_pt_backend_dpa4_params(model_params): - raise ValueError( - "DPA4/SeZM `.pt` checkpoints belong to the regular `pt` backend. " - "Use the `pt` backend for eager checkpoint inference, or export " - "the checkpoint to `.pt2` / `.pte` before loading it with `pt_expt`." - ) - model = get_model(deepcopy(model_params)).to(DEVICE) # Strip the `_CompiledModel` wrapper that pt_expt training applies 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 de93d8f926..6ed5d3e343 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 @@ -4,8 +4,9 @@ Covers two pieces: 1. ``Backend.detect_backend_by_model`` sniffs ``.pt`` content - (``.w``/``.b`` -> pt_expt, ``.matrix``/``.bias`` -> pt) so that - ``dp test -m foo.pt`` routes to the right backend. + (``.w`` weights -> pt_expt, ``.matrix`` weights -> pt, with bias names + used only as a fallback) so that ``dp test -m foo.pt`` routes to the + right backend. 2. ``pt_expt.DeepEval._load_pt`` reconstructs the model from ``_extra_state["model_params"]``, loads ``state_dict``, and runs inference in eager mode, producing outputs that match a direct @@ -710,6 +711,80 @@ def _spin_eager_reference(model, COORD, ATYPE, SPIN, BOX): return {k: v.detach().cpu().numpy() for k, v in ref.items()} +class TestPtExptLoadPtNativeSpinDPA4(unittest.TestCase): + """A native-spin DPA4 checkpoint reaches the graph-spin eager runner.""" + + @classmethod + def setUpClass(cls) -> None: + from ..model.test_dpa4_native_spin import ( + NATIVE_SPIN_CONFIG, + _build_native_spin_model_cpu, + ) + from .test_deep_eval_spin import ( + ATYPE, + BOX, + COORD, + SPIN, + ) + + cls.ATYPE = ATYPE + cls.BOX = BOX + cls.COORD = COORD + cls.SPIN = SPIN + cls.model = _build_native_spin_model_cpu().to(DEVICE).eval() + cls.ref = _spin_eager_reference( + cls.model, cls.COORD, cls.ATYPE, cls.SPIN, cls.BOX + ) + cls.pt_path = tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name + _save_pt_checkpoint( + cls.model, + copy.deepcopy(NATIVE_SPIN_CONFIG), + cls.pt_path, + ) + + @classmethod + def tearDownClass(cls) -> None: + if os.path.exists(cls.pt_path): + os.unlink(cls.pt_path) + + def test_auto_dispatch_and_eval_match_eager(self) -> None: + backend = Backend.detect_backend_by_model(self.pt_path) + self.assertIs(backend, Backend.get_backend("pt-expt")) + + dp = DeepPot( + self.pt_path, + auto_batch_size=False, + neighbor_graph_method="dense", + ) + self.assertTrue(dp.has_spin) + self.assertEqual(dp.deep_eval.metadata["lower_input_kind"], "graph") + + energy, force, virial, force_mag, mask_mag = dp.eval( + self.COORD, + self.BOX, + self.ATYPE, + atomic=False, + spin=self.SPIN, + ) + for name, actual in ( + ("energy", energy), + ("force", force), + ("virial", virial), + ("force_mag", force_mag), + ): + np.testing.assert_allclose( + actual.reshape(-1), + self.ref[name].reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + np.testing.assert_array_equal( + mask_mag.reshape(-1), + self.ref["mask_mag"].reshape(-1), + ) + + class _SpinFilesMixin: """Build .pt + .pte for the chosen ``spin_config`` once per class.""" From 5f37d97b495cc8516e4ea8a411cf9419c37167d6 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 30 Jul 2026 12:38:12 +0800 Subject: [PATCH 5/5] fix(pt_expt): enforce raw checkpoint routing contracts --- deepmd/backend/pt_expt.py | 45 +-- deepmd/pt_expt/infer/deep_eval.py | 359 +++++++++--------- deepmd/pt_expt/model/graph_lower.py | 32 ++ deepmd/pt_expt/utils/serialization.py | 27 +- deepmd/utils/pt_checkpoint.py | 51 +++ .../infer/test_deep_eval_pt_checkpoint.py | 167 +++++++- .../pt_expt/utils/test_graph_pt2_metadata.py | 8 +- 7 files changed, 431 insertions(+), 258 deletions(-) create mode 100644 deepmd/utils/pt_checkpoint.py diff --git a/deepmd/backend/pt_expt.py b/deepmd/backend/pt_expt.py index 7f24cfbec9..5fcb96b148 100644 --- a/deepmd/backend/pt_expt.py +++ b/deepmd/backend/pt_expt.py @@ -1,20 +1,21 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from collections.abc import ( Callable, - Mapping, ) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, - Any, ClassVar, ) from deepmd.backend.backend import ( Backend, ) +from deepmd.utils.pt_checkpoint import ( + detect_pt_checkpoint_backend, +) if TYPE_CHECKING: from argparse import ( @@ -29,46 +30,6 @@ ) -def detect_pt_checkpoint_backend(checkpoint: Any) -> str | None: - """Detect the backend dialect of a raw PyTorch checkpoint. - - Parameters - ---------- - checkpoint : Any - A checkpoint payload or its unwrapped model state dictionary. - - Returns - ------- - str or None - ``"pt-expt"`` or ``"pt"`` when the parameter names identify one - backend unambiguously, otherwise ``None``. - """ - state_dict = checkpoint - if isinstance(state_dict, Mapping) and "model" in state_dict: - state_dict = state_dict["model"] - if not isinstance(state_dict, Mapping): - return None - - keys = tuple(key for key in state_dict if isinstance(key, str)) - - # Weight names are decisive. pt_expt DPA4 also contains ordinary - # torch-native ``.bias`` parameters, so bias names cannot override a - # clear ``.w`` versus ``.matrix`` distinction. - has_pt_expt_weight = any(key.endswith(".w") for key in keys) - has_pt_weight = any(key.endswith(".matrix") for key in keys) - if has_pt_expt_weight or has_pt_weight: - if has_pt_expt_weight == has_pt_weight: - return None - return "pt-expt" if has_pt_expt_weight else "pt" - - # Bias-only state dictionaries retain the original dialect fallback. - has_pt_expt_bias = any(key.endswith(".b") for key in keys) - has_pt_bias = any(key.endswith(".bias") for key in keys) - if has_pt_expt_bias == has_pt_bias: - return None - return "pt-expt" if has_pt_expt_bias else "pt" - - @Backend.register("pt-expt") @Backend.register("pytorch-exportable") class PyTorchExportableBackend(Backend): diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 57d40fe9f3..b8b85d2a4b 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 logging import warnings from collections.abc import ( Callable, @@ -13,9 +14,6 @@ import numpy as np import torch -from deepmd.backend.pt_expt import ( - detect_pt_checkpoint_backend, -) from deepmd.dpmodel.model.transform_output import ( communicate_extended_output, ) @@ -68,6 +66,9 @@ VesinNeighborList, is_vesin_torch_available, ) +from deepmd.utils.pt_checkpoint import ( + detect_pt_checkpoint_backend, +) if TYPE_CHECKING: import ase.neighborlist @@ -79,6 +80,8 @@ NeighborGraph, ) +log = logging.getLogger(__name__) + # Public output keys emitted by graph-lower forwards, keyed by the # output-variable category that ``request_defs`` carries. The graph path is @@ -178,18 +181,21 @@ class DeepEval(DeepEvalBackend): neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional 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. + rejected for graph-routed artifacts because switching those models to + the nlist lower can change their predictions. Use + ``neighbor_graph_method="ase"`` for ASE-based graph construction. nlist_backend : str, default: "auto" 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. + are rejected for graph-routed artifacts. 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"``): ``"auto"`` selects - ``"nv"`` on CUDA when nvalchemiops is available, then ``"vesin"`` when - available, and finally falls back to ``"dense"``. Explicit + ``"nv"`` on CUDA when nvalchemiops is available and otherwise falls + back to ``"dense"``. ``"vesin"`` remains explicit opt-in because it + loops over frames in Python. 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 @@ -240,9 +246,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 not in ("auto", "dense") and getattr( - self, "metadata", {} - ).get("lower_input_kind") not in ("graph", "dpa1_canonical"): + if neighbor_graph_method != "auto" 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 " @@ -282,10 +288,15 @@ def _resolve_neighbor_graph_method(method: str) -> str: DEVICE, ) - if DEVICE.type == "cuda" and is_nv_available(): - return "nv" - if is_vesin_torch_available(): - return "vesin" + if DEVICE.type == "cuda": + if is_nv_available(): + return "nv" + log.warning( + "nvalchemi-toolkit-ops is unavailable; falling back from " + "neighbor_graph_method='auto' to the dense graph builder. " + "Install it with `pip install nvalchemi-toolkit-ops` to enable " + "the NV graph builder." + ) return "dense" def _setup_neighbor_backend(self, nlist_backend: str) -> None: @@ -295,8 +306,8 @@ def _setup_neighbor_backend(self, nlist_backend: str) -> None: 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. Graph-routed - artifacts use ``neighbor_graph_method`` instead, reject an explicit ASE - neighbor list, and warn when an explicit nlist backend is ignored. + artifacts use ``neighbor_graph_method`` instead and reject nlist-specific + controls rather than silently changing or ignoring the requested lower. Results are unchanged either way; only the neighbor-search cost differs. """ if nlist_backend not in ("auto", "vesin", "native"): @@ -307,16 +318,18 @@ def _setup_neighbor_backend(self, nlist_backend: str) -> None: 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." + "neighbor_list cannot be used with this graph-routed model: " + "switching to the nlist lower would change its inference " + "semantics. Use neighbor_graph_method='ase' for ASE-based " + "graph construction, or load an explicitly exported " + "nlist-form artifact." ) 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, + raise ValueError( + f"nlist_backend={nlist_backend!r} only applies to " + "nlist-routed artifacts and cannot be used with this " + "graph-routed model. Use neighbor_graph_method to select " + "the graph builder." ) self._neighbor_graph_method = self._resolve_neighbor_graph_method( self._neighbor_graph_method @@ -325,17 +338,11 @@ def _setup_neighbor_backend(self, nlist_backend: str) -> None: self._nlist_builder = None return is_spin = bool(getattr(self, "_is_spin", False)) - # Native-spin (NeighborGraph route) graph-form artifacts never touch - # this NLIST builder at all -- graph-form eval uses - # ``neighbor_graph_method`` instead (see ``_eval_model_graph_spin``). - # Only the virtual-atom (dense/nlist) spin scheme actually needs the - # vesin restriction below. - is_native_spin_graph = is_spin and ( - getattr(self, "metadata", {}).get("lower_input_kind") == "graph" - ) + # Graph-routed artifacts returned above. Any remaining spin model uses + # the virtual-atom nlist scheme, which vesin does not support. ase_provided = self.neighbor_list is not None # reason vesin cannot be used (None means it can) - unsupported = "spin models" if is_spin and not is_native_spin_graph else None + unsupported = "spin models" if is_spin else None if nlist_backend == "native": self._use_vesin = False elif nlist_backend == "vesin": @@ -558,6 +565,14 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "dialect. Load it with `dp --pt`, or export it to `.pt2` / " "`.pte` before loading it with `pt_expt`." ) + if checkpoint_backend is None: + raise ValueError( + f"Cannot determine the parameter dialect of checkpoint " + f"'{model_file}': its state dictionary is ambiguous or does " + "not contain backend-specific parameter names. Load it with " + "the backend that produced it, or use an unambiguous " + "pt_expt training checkpoint." + ) extra = state_dict.get("_extra_state") if isinstance(state_dict, dict) else None if not (isinstance(extra, dict) and "model_params" in extra): raise ValueError( @@ -699,11 +714,11 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "lower_input_kind": "graph" if use_graph_lower else "nlist", } if use_graph_lower: - from deepmd.pt_expt.utils.serialization import ( - _graph_edge_dtype, + from deepmd.pt_expt.model.graph_lower import ( + graph_edge_dtype, ) - self.metadata["graph_edge_dtype"] = _graph_edge_dtype(model, "graph") + 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] @@ -719,148 +734,150 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: # Graph spin: (..., source_row_ptr, spin, fparam, aparam, charge_spin) # Nlist: (ext_coord, ext_atype, nlist, mapping, fparam, aparam) # Nlist spin: (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) - if use_graph_lower and self._is_spin: - - def _eager_runner_graph_spin( - 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, - spin: 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, - spin=spin, - ) - return model._translate_eager_call( - model_ret, - atype, - do_atomic_virial=True, - ) + if use_graph_lower: + if self._is_spin: + + def _eager_runner_graph_spin( + 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, + spin: 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, + spin=spin, + ) + return model._translate_eager_call( + model_ret, + atype, + do_atomic_virial=True, + ) - self.exported_module = _eager_runner_graph_spin - elif 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_spin + else: + from deepmd.pt_expt.model.ener_model import ( + _translate_energy_keys, ) - self.exported_module = _eager_runner_graph - elif self._is_spin: - - def _eager_runner_spin( - ext_coord: torch.Tensor, - ext_atype: torch.Tensor, - ext_spin: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None = None, - ) -> dict[str, torch.Tensor]: - ext_coord = ext_coord.detach().requires_grad_(True) - return model.forward_common_lower( - ext_coord, - ext_atype, - ext_spin, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=True, - ) + 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_spin + self.exported_module = _eager_runner_graph else: + if self._is_spin: + + def _eager_runner_spin( + ext_coord: torch.Tensor, + ext_atype: torch.Tensor, + ext_spin: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + ext_coord = ext_coord.detach().requires_grad_(True) + return model.forward_common_lower( + ext_coord, + ext_atype, + ext_spin, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=True, + ) - def _eager_runner( - ext_coord: torch.Tensor, - ext_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None = None, - ) -> dict[str, torch.Tensor]: - ext_coord = ext_coord.detach().requires_grad_(True) - return model.forward_common_lower( - ext_coord, - ext_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=True, - ) + self.exported_module = _eager_runner_spin + else: + + def _eager_runner( + ext_coord: torch.Tensor, + ext_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + ext_coord = ext_coord.detach().requires_grad_(True) + return model.forward_common_lower( + ext_coord, + ext_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=True, + ) - self.exported_module = _eager_runner + self.exported_module = _eager_runner def get_rcut(self) -> float: """Get the cutoff radius of this model.""" diff --git a/deepmd/pt_expt/model/graph_lower.py b/deepmd/pt_expt/model/graph_lower.py index f196bd4fe5..e2f32dcfcf 100644 --- a/deepmd/pt_expt/model/graph_lower.py +++ b/deepmd/pt_expt/model/graph_lower.py @@ -5,6 +5,38 @@ Any, ) +import torch + + +def graph_edge_dtype(model: Any, lower_kind: str) -> str: + """Return the graph edge-vector dtype encoded by a deployment artifact. + + Parameters + ---------- + model : Any + Model exposing an atomic-model descriptor. + lower_kind : str + Concrete lower-forward schema. + + Returns + ------- + str + ``"float32"`` for eligible geometrically compressed DPA1 graph + lowers, otherwise ``"float64"``. + """ + atomic_model = getattr(model, "atomic_model", None) + descriptor = getattr(atomic_model, "descriptor", None) + descriptor_block = getattr(descriptor, "se_atten", None) + statistics = getattr(descriptor_block, "mean", None) + if ( + lower_kind in ("graph", "dpa1_canonical") + and bool(getattr(descriptor, "geo_compress", False)) + and isinstance(statistics, torch.Tensor) + and statistics.dtype == torch.float32 + ): + return "float32" + return "float64" + def model_uses_graph_lower(model: Any) -> bool: """Return whether a model's default lower uses ``NeighborGraph``. diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 897d91cec3..ff297cb4d2 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -26,6 +26,9 @@ from deepmd.dpmodel.utils.serialization import ( traverse_model_dict, ) +from deepmd.pt_expt.model.graph_lower import ( + graph_edge_dtype, +) # --------------------------------------------------------------------------- # AOTInductor ``.pt2`` archive layout. @@ -1015,28 +1018,6 @@ def _build_dynamic_shapes( return (*base, None, None, None, None, None, None, None, None) -def _graph_edge_dtype(model: torch.nn.Module, lower_kind: str) -> str: - """Return the graph edge-vector dtype encoded by the deployment artifact. - - Geometrically compressed DPA1 with float32 descriptor statistics evaluates - both descriptor directions in float32 and therefore accepts float32 - geometry directly. Other graph descriptors retain the model-agnostic - float64 geometry ABI. - """ - atomic_model = getattr(model, "atomic_model", None) - descriptor = getattr(atomic_model, "descriptor", None) - descriptor_block = getattr(descriptor, "se_atten", None) - statistics = getattr(descriptor_block, "mean", None) - if ( - lower_kind in ("graph", "dpa1_canonical") - and bool(getattr(descriptor, "geo_compress", False)) - and isinstance(statistics, torch.Tensor) - and statistics.dtype == torch.float32 - ): - return "float32" - return "float64" - - def _supports_graph_export(model: torch.nn.Module) -> bool: """Whether the model has an exportable graph-lower implementation. @@ -1173,7 +1154,7 @@ def _probe_has_message_passing(obj: object) -> bool | None: # "graph" → NeighborGraph (atype, n_node, edge_index, edge_vec, edge_mask) # The C++ loader branches on this to build the matching inputs. meta["lower_input_kind"] = lower_kind - meta["graph_edge_dtype"] = _graph_edge_dtype(model, lower_kind) + meta["graph_edge_dtype"] = graph_edge_dtype(model, lower_kind) # Model-level pair-type exclusion (``pair_exclude_types``): a list of # ``[ti, tj]`` type pairs whose interaction is dropped. Exclusion is a diff --git a/deepmd/utils/pt_checkpoint.py b/deepmd/utils/pt_checkpoint.py new file mode 100644 index 0000000000..de71c50a09 --- /dev/null +++ b/deepmd/utils/pt_checkpoint.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Utilities shared by PyTorch checkpoint backends.""" + +from collections.abc import ( + Mapping, +) +from typing import ( + Any, +) + + +def detect_pt_checkpoint_backend(checkpoint: Any) -> str | None: + """Detect the parameter dialect of a raw PyTorch checkpoint. + + Parameters + ---------- + checkpoint : Any + A checkpoint payload or its unwrapped model state dictionary. + + Returns + ------- + str or None + ``"pt-expt"`` or ``"pt"`` when the parameter names identify one + backend unambiguously, otherwise ``None``. + """ + state_dict = checkpoint + if isinstance(state_dict, Mapping) and "model" in state_dict: + state_dict = state_dict["model"] + if not isinstance(state_dict, Mapping): + return None + + keys = tuple(key for key in state_dict if isinstance(key, str)) + + # Weight names are decisive. pt_expt DPA4 also contains ordinary + # torch-native ``.bias`` parameters, so bias names cannot override a + # clear ``.w`` versus ``.matrix`` distinction. + has_pt_expt_weight = any(key.endswith(".w") for key in keys) + has_pt_weight = any(key.endswith(".matrix") for key in keys) + if has_pt_expt_weight or has_pt_weight: + if has_pt_expt_weight == has_pt_weight: + return None + return "pt-expt" if has_pt_expt_weight else "pt" + + # A lone ``.b`` is specific to pt_expt's NativeLayer. A lone ``.bias`` + # is not specific to pt because pt_expt models can contain torch-native + # modules with that suffix, so it remains deliberately unclassified. + has_pt_expt_bias = any(key.endswith(".b") for key in keys) + has_pt_bias = any(key.endswith(".bias") for key in keys) + if has_pt_expt_bias and not has_pt_bias: + return "pt-expt" + return None 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 6ed5d3e343..fb1f86afa3 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 @@ -4,9 +4,9 @@ Covers two pieces: 1. ``Backend.detect_backend_by_model`` sniffs ``.pt`` content - (``.w`` weights -> pt_expt, ``.matrix`` weights -> pt, with bias names - used only as a fallback) so that ``dp test -m foo.pt`` routes to the - right backend. + (``.w`` weights -> pt_expt and ``.matrix`` weights -> pt, with only the + pt_expt-specific ``.b`` suffix used as a fallback) so that + ``dp test -m foo.pt`` routes to the right backend. 2. ``pt_expt.DeepEval._load_pt`` reconstructs the model from ``_extra_state["model_params"]``, loads ``state_dict``, and runs inference in eager mode, producing outputs that match a direct @@ -18,6 +18,9 @@ import shutil import tempfile import unittest +from unittest import ( + mock, +) import numpy as np import pytest @@ -32,9 +35,6 @@ 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, ) @@ -55,8 +55,8 @@ from deepmd.pt_expt.utils.env import ( DEVICE, ) -from deepmd.pt_expt.utils.vesin_neighbor_list import ( - is_vesin_torch_available, +from deepmd.utils.pt_checkpoint import ( + detect_pt_checkpoint_backend, ) from ...seed import ( @@ -206,12 +206,45 @@ def _save_pt_checkpoint_compiled( torch.save({"model": cooked}, path) +class TestPtCheckpointBackendDetection(unittest.TestCase): + """Checkpoint dialect detection must be conservative and deterministic.""" + + def test_parameter_name_matrix(self) -> None: + cases = { + "wrapped pt_expt weight with torch bias": ( + {"model": {"layer.w": object(), "module.bias": object()}}, + "pt-expt", + ), + "unwrapped pt weight": ({"layer.matrix": object()}, "pt"), + "pt_expt bias only": ({"layer.b": object()}, "pt-expt"), + "torch bias only": ({"layer.bias": object()}, None), + "mixed weights": ( + {"left.w": object(), "right.matrix": object()}, + None, + ), + "mixed biases": ( + {"left.b": object(), "right.bias": object()}, + None, + ), + "no string keys": ({1: object()}, None), + "non-mapping payload": (object(), None), + "wrapped non-mapping payload": ({"model": object()}, None), + } + for name, (checkpoint, expected) in cases.items(): + with self.subTest(name=name): + self.assertEqual( + detect_pt_checkpoint_backend(checkpoint), + expected, + ) + + class TestBackendDispatchPt(unittest.TestCase): """``Backend.detect_backend_by_model`` must sniff `.pt` content.""" def setUp(self) -> None: # Real pt_expt-trained checkpoint (uses `.w`/`.b` keys). model, model_params = _build_model_and_params() + self.output_def = ModelOutputDef(model.atomic_output_def()) self.pt_expt_pt = tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name _save_pt_checkpoint(model, model_params, self.pt_expt_pt) @@ -233,6 +266,18 @@ def setUp(self) -> None: self.pt_pt, ) + # Mixed weight dialects are deliberately unclassified. + self.ambiguous_pt = tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name + torch.save( + { + "model": { + "left.w": torch.zeros(1), + "right.matrix": torch.zeros(1), + } + }, + self.ambiguous_pt, + ) + # File that exists but is not a valid torch checkpoint — sniffing # must fail gracefully and fall back to suffix dispatch. self.bogus_pt = tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name @@ -240,7 +285,12 @@ def setUp(self) -> None: f.write(b"not a real torch file") def tearDown(self) -> None: - for p in (self.pt_expt_pt, self.pt_pt, self.bogus_pt): + for p in ( + self.pt_expt_pt, + self.pt_pt, + self.ambiguous_pt, + self.bogus_pt, + ): if os.path.exists(p): os.unlink(p) @@ -258,6 +308,16 @@ def test_bogus_pt_falls_back_to_suffix(self) -> None: backend = Backend.detect_backend_by_model(self.bogus_pt) self.assertIs(backend, Backend.get_backend("pt")) + def test_forced_pt_expt_rejects_other_dialects(self) -> None: + cases = ( + (self.pt_pt, "regular `pt` parameter dialect"), + (self.ambiguous_pt, "Cannot determine the parameter dialect"), + ) + for checkpoint, message in cases: + with self.subTest(checkpoint=checkpoint): + with self.assertRaisesRegex(ValueError, message): + PtExptDeepEval(checkpoint, self.output_def) + class TestPtExptLoadPt(unittest.TestCase): """``pt_expt.DeepEval._load_pt`` produces outputs matching the source model.""" @@ -285,6 +345,25 @@ def test_metadata_accessors(self) -> None: self.assertEqual(de.get_dim_aparam(), 0) self.assertFalse(de._is_spin) + def test_nlist_controls_remain_available(self) -> None: + output_def = ModelOutputDef(self.model.atomic_output_def()) + PtExptDeepEval( + self.pt_path, + output_def, + neighbor_list=mock.sentinel.neighbor_list, + ) + PtExptDeepEval( + self.pt_path, + output_def, + nlist_backend="native", + ) + with self.assertRaisesRegex(ValueError, "only applies to graph-routed"): + PtExptDeepEval( + self.pt_path, + output_def, + neighbor_graph_method="dense", + ) + def test_eval_matches_source_model(self) -> None: """Run inference via DeepPot(.pt) and compare to direct forward.""" dp = DeepPot(self.pt_path) @@ -340,18 +419,49 @@ def test_unsupported_extension_raises(self) -> None: os.unlink(bogus) +class TestNeighborGraphMethodResolution(unittest.TestCase): + """Auto graph-builder selection must cover each host policy explicitly.""" + + def test_auto_resolution(self) -> None: + cases = ( + ("cpu", False, "dense", False), + ("cuda", True, "nv", False), + ("cuda", False, "dense", True), + ) + for device_type, nv_available, expected, warns in cases: + with self.subTest( + device_type=device_type, + nv_available=nv_available, + ): + with ( + mock.patch( + "deepmd.pt_expt.utils.env.DEVICE", + torch.device(device_type), + ), + mock.patch( + "deepmd.pt.utils.nv_nlist.is_nv_available", + return_value=nv_available, + ), + ): + if warns: + with self.assertLogs( + "deepmd.pt_expt.infer.deep_eval", + level="WARNING", + ): + actual = PtExptDeepEval._resolve_neighbor_graph_method( + "auto" + ) + else: + actual = PtExptDeepEval._resolve_neighbor_graph_method("auto") + self.assertEqual(actual, expected) + + 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() - 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, @@ -401,10 +511,6 @@ 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, @@ -421,6 +527,28 @@ def test_eval_matches_public_forward(self) -> None: err_msg=name, ) + def test_nlist_controls_fail_without_changing_graph_semantics(self) -> None: + output_def = ModelOutputDef(self.model.atomic_output_def()) + path = self.pt_paths["plain"] + with self.assertRaisesRegex( + ValueError, + "switching to the nlist lower would change", + ): + PtExptDeepEval( + path, + output_def, + neighbor_list=mock.sentinel.neighbor_list, + ) + with self.assertRaisesRegex( + ValueError, + "only applies to nlist-routed artifacts", + ): + PtExptDeepEval( + path, + output_def, + nlist_backend="native", + ) + class TestPtExptLoadPtCompiledLayout(unittest.TestCase): """`.pt` saved after pt_expt training compilation (`_CompiledModel` wrap). @@ -758,6 +886,7 @@ def test_auto_dispatch_and_eval_match_eager(self) -> None: ) self.assertTrue(dp.has_spin) self.assertEqual(dp.deep_eval.metadata["lower_input_kind"], "graph") + self.assertEqual(dp.deep_eval._neighbor_graph_method, "dense") energy, force, virial, force_mag, mask_mag = dp.eval( self.COORD, diff --git a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py index bcd0b3b0c8..5ee3ed91df 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -17,8 +17,10 @@ import pytest import torch +from deepmd.pt_expt.model.graph_lower import ( + graph_edge_dtype, +) from deepmd.pt_expt.utils.serialization import ( - _graph_edge_dtype, _needs_with_comm_artifact, _supports_graph_export, deserialize_to_file, @@ -156,8 +158,8 @@ class _AtomicModel: class _Model: atomic_model = _AtomicModel() - assert _graph_edge_dtype(_Model(), "graph") == expected - assert _graph_edge_dtype(_Model(), "nlist") == "float64" + assert graph_edge_dtype(_Model(), "graph") == expected + assert graph_edge_dtype(_Model(), "nlist") == "float64" assert _supports_graph_export(_Model()) is (statistics_dtype == torch.float32)