diff --git a/deepmd/backend/pt_expt.py b/deepmd/backend/pt_expt.py index 38b66f0104..5fcb96b148 100644 --- a/deepmd/backend/pt_expt.py +++ b/deepmd/backend/pt_expt.py @@ -13,6 +13,9 @@ from deepmd.backend.backend import ( Backend, ) +from deepmd.utils.pt_checkpoint import ( + detect_pt_checkpoint_backend, +) if TYPE_CHECKING: from argparse import ( @@ -51,11 +54,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 +69,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 999e7f69f1..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, @@ -65,6 +66,9 @@ VesinNeighborList, is_vesin_torch_available, ) +from deepmd.utils.pt_checkpoint import ( + detect_pt_checkpoint_backend, +) if TYPE_CHECKING: import ase.neighborlist @@ -76,12 +80,14 @@ NeighborGraph, ) +log = logging.getLogger(__name__) + -# 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", @@ -132,18 +138,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. @@ -185,18 +179,26 @@ 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 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`` 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 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 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 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. @@ -212,14 +214,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") @@ -241,20 +243,20 @@ 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). - if neighbor_graph_method != "dense" and getattr(self, "metadata", {}).get( + # ``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 != "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-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: @@ -268,32 +270,79 @@ 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": + 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: + """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 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"): 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 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": + 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 + ) + 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 - # ``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": @@ -505,9 +554,25 @@ 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`." + ) + 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( @@ -547,13 +612,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 @@ -589,6 +647,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,68 +711,173 @@ 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.model.graph_lower 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. - # 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: + # 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``. + # + # 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: + 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, + ) - 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, + self.exported_module = _eager_runner_graph_spin + else: + from deepmd.pt_expt.model.ener_model import ( + _translate_energy_keys, ) - self.exported_module = _eager_runner_spin + 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 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.""" @@ -1977,18 +2145,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 +2312,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/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 6a9b7e2c59..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,7 +4,8 @@ Covers two pieces: 1. ``Backend.detect_backend_by_model`` sniffs ``.pt`` content - (``.w``/``.b`` -> pt_expt, ``.matrix``/``.bias`` -> pt) so that + (``.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 @@ -17,6 +18,9 @@ import shutil import tempfile import unittest +from unittest import ( + mock, +) import numpy as np import pytest @@ -40,6 +44,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 +55,9 @@ from deepmd.pt_expt.utils.env import ( DEVICE, ) +from deepmd.utils.pt_checkpoint import ( + detect_pt_checkpoint_backend, +) from ...seed import ( GLOBAL_SEED, @@ -110,6 +121,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, @@ -162,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) @@ -189,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 @@ -196,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) @@ -214,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.""" @@ -241,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) @@ -296,6 +419,137 @@ 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() + 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, + ) + + 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). @@ -585,6 +839,81 @@ 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") + self.assertEqual(dp.deep_eval._neighbor_graph_method, "dense") + + 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.""" 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)