diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 1b97ef7f50..5881cc132c 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,23 @@ 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_graph_method : str, default: "dense" - Carry-all graph builder for GRAPH-FORM ``.pt2`` artifacts ONLY - (``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 + 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: "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 + ``"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 backend-selection API is deferred to the dense-nlist deprecation. @@ -160,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 - # 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 is resolved once after model metadata + # identifies the lower ABI. self._neighbor_graph_method = neighbor_graph_method self._is_pt2 = model_file.endswith(".pt2") @@ -184,20 +190,20 @@ 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). - if neighbor_graph_method != "dense" and getattr(self, "metadata", {}).get( - "lower_input_kind" - ) not in ("graph", "dpa1_canonical"): + # ``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"): 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." ) - self._setup_nlist_backend(nlist_backend) + self._setup_neighbor_backend(nlist_backend) if isinstance(auto_batch_size, bool): if auto_batch_size: @@ -211,20 +217,66 @@ 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 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._neighbor_graph_method = self._resolve_neighbor_graph_method( + self._neighbor_graph_method + ) + 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 +571,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 +635,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 +1876,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 +2043,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..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, ) @@ -40,6 +43,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, @@ -47,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, @@ -110,6 +120,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 +339,88 @@ 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() + 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, + } + _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") + self.assertEqual( + dp.deep_eval._neighbor_graph_method, + self.expected_graph_method, + ) + 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).