diff --git a/backend/find_pytorch.py b/backend/find_pytorch.py index f969c5ebce..1e7117b54a 100644 --- a/backend/find_pytorch.py +++ b/backend/find_pytorch.py @@ -126,6 +126,19 @@ def get_pt_requirement(pt_version: str = "") -> dict: else: mpi_requirement = [] + # Vesin's torch bindings require PyTorch 2.3. Older PyTorch targets and + # macOS x86_64, whose last PyTorch wheels are 2.2, use native neighbors. + vesin_requirement = [] + if pt_version == "" or Version(pt_version) >= Version("2.3"): + supported_platform = ( + "platform_system != 'Darwin' or platform_machine != 'x86_64'" + ) + vesin_requirement = [ + f"vesin[torch]>=0.6; {supported_platform}", + # The extra does not constrain the independently versioned bindings. + f"vesin-torch>=0.6; {supported_platform}", + ] + return { "torch": [ # uv has different local version behaviors, i.e. `==2.3.1` cannot match `==2.3.1+cpu` @@ -137,11 +150,7 @@ def get_pt_requirement(pt_version: str = "") -> dict: # https://github.com/pytorch/pytorch/commit/7e0c26d4d80d6602aed95cb680dfc09c9ce533bc else "torch>=2.1.0", "e3nn>=0.5.9", - # O(N) cell-list neighbor list for fast Python/ASE inference; the - # torch bindings (vesin-torch) ship only as a PyPI extra, so keep it - # under the torch extra rather than the core deps (conda-forge has - # vesin but not vesin-torch). - "vesin[torch]", + *vesin_requirement, # GPU O(N) cell-list neighbor list for large systems. Restricted to # Python >= 3.11 (the package requires it while deepmd-kit still # supports 3.10) and to Linux: it is a CUDA package, and its diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 0eca4776ed..402e09ba77 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -1,9 +1,24 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """ASE calculator interface module.""" +import logging +import math +import weakref +from contextlib import ( + ExitStack, +) +from functools import ( + partial, +) from pathlib import ( Path, ) +from shutil import ( + copyfile, +) +from tempfile import ( + TemporaryDirectory, +) from typing import ( TYPE_CHECKING, Any, @@ -12,6 +27,7 @@ Union, ) +import numpy as np from ase.calculators.calculator import ( Calculator, PropertyNotImplementedError, @@ -23,6 +39,10 @@ ) if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + from ase import ( Atoms, ) @@ -30,8 +50,84 @@ NeighborList, ) + from deepmd.infer.deep_eval import ( + DeepEvalBackend, + ) + __all__ = ["DP"] +log = logging.getLogger(__name__) + + +def _compress_dpa4c_for_ase(model: Any) -> bool: + """Select supported DPA4C compression without changing other model families.""" + from deepmd.pt_expt.descriptor.dpa4c import ( + DescrptDPA4C, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + compression_ineligibility_reason, + ) + + parts = model.atomic_model.fused_decomposition() + if ( + parts is None + or parts[0] is None + or not isinstance(parts[0].descriptor, DescrptDPA4C) + ): + return False + reason = compression_ineligibility_reason(parts[0].descriptor) + if reason is None: + return True + log.warning( + "DPA4C compression is unavailable: %s Freezing without compression.", reason + ) + return False + + +def _checkpoint_freezer( + evaluator: "DeepEvalBackend", head: str | None, *, freeze_requested: bool +) -> "Callable[..., None] | None": + """Select an exporter for the loaded checkpoint's model and head.""" + from deepmd.pt.infer.deep_eval import DeepEval as PTDeepEval + from deepmd.pt.model.model.sezm_model import ( + SeZMModel, + ) + from deepmd.pt_expt.infer.deep_eval import DeepEval as PTExptDeepEval + + if isinstance(evaluator, PTDeepEval) and isinstance( + evaluator.get_model(), SeZMModel + ): + from deepmd.pt.entrypoints.freeze_pt2 import ( + freeze_sezm_to_pt2, + ) + + return partial( + freeze_sezm_to_pt2, + head=head if evaluator.multi_task else None, + with_comm=False, + ) + if ( + isinstance(evaluator, PTExptDeepEval) + and evaluator.metadata["lower_input_kind"] == "graph" + ): + from deepmd.pt_expt.entrypoints.main import ( + freeze, + ) + + # The checkpoint loader selects Default when no head is supplied. + return partial( + freeze, + head="Default" if head is None else head, + lower_kind="graph", + with_comm=False, + compress=( + _compress_dpa4c_for_ase(evaluator._dpmodel) + if freeze_requested + else False + ), + ) + return None + class DP(Calculator): """Implementation of ASE deepmd calculator. @@ -51,6 +147,34 @@ class DP(Calculator): The neighbor list object. If None, then build the native neighbor list. head : Union[str, None], optional a specific model branch choosing from pretrained model, by default None + nlist_backend : str, optional + Neighbor search for models using the extended-coordinate input schema. + The default ``"auto"`` selects an available accelerated search. + neighbor_graph_method : str, optional + Graph search for PT-expt graph models or PT SeZM/DPA4, by default + ``"auto"``. Single-frame CUDA inference prefers Vesin. PT-expt graph + models also accept ``"dense"``, ``"ase"``, ``"cell"``, ``"vesin"``, + and ``"nv"``; SeZM edge-vector inputs accept ``"vesin"`` and ``"nv"``. + neighbor_skin : float, optional + Vesin candidate-search skin in angstrom, by default 1.0. Set to zero + to rebuild candidates every step. All searches still use the model's + exact cutoff for the returned neighbors. Applies to both graph and + extended-coordinate inference, independently of checkpoint freezing. + freeze : bool, optional + Export an eligible ``.pt`` checkpoint to ``.pt2`` for inference, + by default False. + Supports PT DPA4/SeZM and graph-eligible PT-expt models. Each call + retains a private artifact for its evaluator's lifetime and atomically + publishes a copy as ``-ase_single.pt2`` beside the + checkpoint. Concurrent processes retain their own selected model heads; + the last publication replaces the shared copy. Export uses the + backend's freeze-time kernel policy and omits the + multi-rank communication artifact. DPA4C models are also compressed + when their configuration supports it; otherwise a warning explains + why the export remains uncompressed. Other models ignore this option + and are loaded directly. Direct checkpoint inference for supported + models can be substantially slower; False is intended for debugging + without the compilation overhead. Examples -------- @@ -92,21 +216,74 @@ def __init__( neighbor_list: Optional["NeighborList"] = None, head: str | None = None, nlist_backend: str = "auto", + freeze: bool = False, + neighbor_graph_method: str = "auto", + neighbor_skin: float = 1.0, **kwargs: Any, ) -> None: Calculator.__init__(self, label=label, **kwargs) + if not math.isfinite(neighbor_skin) or neighbor_skin < 0: + raise ValueError( + f"neighbor_skin must be finite and non-negative, got {neighbor_skin}" + ) + model_path = Path(model).resolve() self.dp = DeepPot( - str(Path(model).resolve()), + str(model_path), neighbor_list=neighbor_list, head=head, nlist_backend=nlist_backend, + neighbor_graph_method=neighbor_graph_method, + neighbor_skin=neighbor_skin, ) + freezer = ( + _checkpoint_freezer(self.dp.deep_eval, head, freeze_requested=freeze) + if model_path.suffix == ".pt" + else None + ) + if freezer is not None and freeze: + output = model_path.with_name(f"{model_path.stem}-ase_single.pt2") + # Release the eager model before the exporter constructs its model. + self.dp.close() + del self.dp + log.info("Freezing ASE checkpoint %s to %s", model_path, output) + with ExitStack() as cleanup: + temporary = TemporaryDirectory( + prefix=f".{model_path.stem}-ase-", dir=model_path.parent + ) + cleanup.callback(temporary.cleanup) + private = Path(temporary.name) / output.name + freezer(str(model_path), str(private)) + self.dp = DeepPot( + str(private), + neighbor_list=neighbor_list, + nlist_backend=nlist_backend, + neighbor_graph_method=neighbor_graph_method, + neighbor_skin=neighbor_skin, + ) + # A separate inode keeps later writes to the public artifact + # from changing the file reopened by evaluator serialization. + publication = Path(temporary.name) / "publish.pt2" + copyfile(private, publication) + publication.replace(output) + weakref.finalize(self.dp.deep_eval, temporary.cleanup) + cleanup.pop_all() + elif freezer is not None: + log.warning( + "ASE inference directly from a .pt checkpoint can be SUBSTANTIALLY " + "slower than a frozen .pt2 model. Use DP(..., freeze=True), or " + "freeze the checkpoint manually and load the .pt2 file. " + "freeze=False is intended for debugging and similar short runs " + "that avoid the upfront compilation overhead." + ) if type_dict: self.type_dict = type_dict else: self.type_dict = dict( zip(self.dp.get_type_map(), range(self.dp.get_ntypes()), strict=True) ) + self._cached_numbers: np.ndarray | None = None + self._cached_types: np.ndarray | None = None + self._cached_type_dict: dict[str, int] = {} def calculate( self, @@ -134,8 +311,22 @@ def calculate( cell = self.atoms.get_cell().reshape([1, -1]) else: cell = None - symbols = self.atoms.get_chemical_symbols() - atype = [self.type_dict[k] for k in symbols] + numbers = self.atoms.numbers + if ( + not np.array_equal(numbers, self._cached_numbers) + or self.type_dict != self._cached_type_dict + ): + self._cached_types = np.fromiter( + ( + self.type_dict[symbol] + for symbol in self.atoms.get_chemical_symbols() + ), + dtype=np.int32, + count=len(numbers), + ) + self._cached_numbers = numbers.copy() + self._cached_type_dict = self.type_dict.copy() + atype = self._cached_types fparam = self.atoms.info.get("fparam", None) aparam = self.atoms.info.get("aparam", None) diff --git a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py index d7e1a5274e..270bbadf3d 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py +++ b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py @@ -110,16 +110,19 @@ def neighbor_graph_from_ijs( j = xp.astype(xp.asarray(j, device=dev), xp.int64) nframe_id = xp.astype(xp.asarray(nframe_id, device=dev), xp.int64) # Within-frame indices become node indices through the frame offsets. - offset = xp.take(xp.cumulative_sum(n_node) - n_node, nframe_id, axis=0) - i_flat = i + offset - j_flat = j + offset + if nf == 1: + i_flat, j_flat = i, j + else: + offset = xp.take(xp.cumulative_sum(n_node) - n_node, nframe_id, axis=0) + i_flat = i + offset + j_flat = j + offset r_i = xp.take(coord_flat, i_flat, axis=0) r_j = xp.take(coord_flat, j_flat, axis=0) edge_vec = r_j - r_i if box is not None: box = xp.asarray(box, device=dev) box = xp.reshape(box, (nf, 3, 3)) - box_per_edge = xp.take(box, nframe_id, axis=0) # (E, 3, 3) + box_per_edge = box[0] if nf == 1 else xp.take(box, nframe_id, axis=0) S = xp.astype(xp.asarray(S, device=dev), box.dtype) # S @ box per edge via broadcast sum (NEVER np.einsum, which breaks on torch): # shift[e, b] = sum_a S[e, a] * box[e, a, b] diff --git a/deepmd/dpmodel/utils/neighbor_list.py b/deepmd/dpmodel/utils/neighbor_list.py index c9b6f5923b..bc8d379fba 100644 --- a/deepmd/dpmodel/utils/neighbor_list.py +++ b/deepmd/dpmodel/utils/neighbor_list.py @@ -63,8 +63,9 @@ class NeighborList: Implementations turn local coordinates into the extended representation: the coordinates and atom types of local-plus-ghost (periodic-image) atoms, a candidate neighbor list indexing the extended atoms, and a mapping from each - extended atom to its local owner. Implementations are stateless -- - ``rcut``/``sel`` are supplied by the model at call time. + extended atom to its local owner. Implementations may retain search storage + and neighbor candidates, while ``rcut``/``sel`` are supplied by the model at + call time. Returned geometry always describes the current coordinates. """ def build( diff --git a/deepmd/main.py b/deepmd/main.py index 1206b2fa0e..6d967553fc 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -369,6 +369,13 @@ def main_parser() -> argparse.ArgumentParser: "always freezes to 'graph', since the dense lower is deprecated in this " "backend -- requesting 'nlist' for one logs a warning and is overridden.", ) + parser_frz.add_argument( + "--compress", + action="store_true", + help="(Supported backend: PyTorch Exportable) Compress the model before " + "export using the default compression settings. Eligible .pt2 models " + "use the accelerated compressed graph without an intermediate export.", + ) # * test script ******************************************************************** parser_tst = subparsers.add_parser( @@ -1098,6 +1105,12 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: """ parser = main_parser() parsed_args = parser.parse_args(args=args) + if ( + parsed_args.command == "freeze" + and parsed_args.compress + and BACKEND_TABLE[parsed_args.backend] != "pytorch-exportable" + ): + parser.error("freeze --compress requires --pt-expt (PyTorch Exportable)") if parsed_args.command is None: parser.print_help() else: diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index b3cd33798a..81d47e45ea 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -531,6 +531,7 @@ def _make_sample_inputs( mapping_t, ) if has_spin: + # Local spins have a frame stride independent of ghost storage. return ( edge_schema.coord, edge_schema.atype, @@ -538,7 +539,7 @@ def _make_sample_inputs( edge_schema.edge_vec, edge_schema.edge_scatter_index, edge_schema.edge_mask, - ext_spin[:, :nloc], + ext_spin[:, :nloc].contiguous(), fparam, aparam, charge_spin, @@ -917,6 +918,7 @@ def freeze_sezm_to_pt2( device: torch.device | None = None, head: str | None = None, atomic_virial: bool = True, + with_comm: bool = True, ) -> None: """Freeze a SeZM checkpoint into an AOTInductor ``.pt2`` archive. @@ -938,6 +940,9 @@ def freeze_sezm_to_pt2( default: the edge-force scatter assembles the per-atom virial as a free by-product of the single backward, so exporting it carries no compute cost. + with_comm + Whether to include a communication artifact when the model supports + cross-rank message passing. Disable for single-process inference. Notes ----- @@ -955,6 +960,7 @@ def freeze_sezm_to_pt2( target_device=target_device, head=head, atomic_virial=atomic_virial, + with_comm=with_comm, ) @@ -965,6 +971,7 @@ def _freeze_sezm_to_pt2( target_device: torch.device, head: str | None, atomic_virial: bool, + with_comm: bool, ) -> None: """Build one AOTInductor archive under an established kernel policy.""" from torch._inductor import ( @@ -1086,11 +1093,13 @@ def _freeze_sezm_to_pt2( # report supports_edge_parallel()=True since the Source Freeze Propagation # gate's per-node partials are completed across ranks by # DescrptSeZM._gate_partial_exchange (issue #5906). - with_comm = ( - model.export_lower_input_kind() == "edge_vec" and model.supports_edge_parallel() + has_comm_artifact = ( + with_comm + and model.export_lower_input_kind() == "edge_vec" + and model.supports_edge_parallel() ) with_comm_bytes: bytes | None = None - if with_comm: + if has_comm_artifact: log.info( "Compiling the parallel with-comm artifact (second AOTInductor " "compilation)..." @@ -1106,7 +1115,7 @@ def _freeze_sezm_to_pt2( output_keys=output_keys, is_spin=is_spin, do_atomic_virial=atomic_virial, - has_comm_artifact=with_comm, + has_comm_artifact=has_comm_artifact, ) with zipfile.ZipFile(out_path_str, "a") as zf: zf.writestr("model/extra/metadata.json", json.dumps(metadata)) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index f468379565..6cf55f588f 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -69,17 +69,12 @@ GLOBAL_PT_FLOAT_PRECISION, RESERVED_PRECISION_DICT, ) -from deepmd.pt.utils.nv_nlist import ( - NvNeighborList, - is_nv_available, -) from deepmd.pt.utils.utils import ( to_numpy_array, to_torch_tensor, ) -from deepmd.pt_expt.utils.vesin_neighbor_list import ( - VesinNeighborList, - is_vesin_torch_available, +from deepmd.pt_expt.utils.auto_neighbor_list import ( + make_neighbor_list, ) from deepmd.utils.batch_size import ( RetrySignal, @@ -143,6 +138,15 @@ class DeepEval(DeepEvalBackend): 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. + nlist_backend : str, default: "auto" + Extended-coordinate neighbor search: ``"auto"``, ``"native"``, + ``"vesin"``, or ``"nv"``. Auto prefers Vesin for a single frame and + Toolkit-Ops for CUDA batches when available. + neighbor_graph_method : str, default: "auto" + SeZM/DPA4 inference search: ``"auto"``, ``"vesin"``, or ``"nv"``. + The strategy runs inside the model's native forward geometry stage. + neighbor_skin : float, default: 0.0 + Vesin candidate-search skin in angstrom. Zero disables Verlet caching. **kwargs : dict Keyword arguments. """ @@ -157,11 +161,15 @@ def __init__( head: str | int | None = None, no_jit: bool = False, nlist_backend: str = "auto", + neighbor_graph_method: str = "auto", + neighbor_skin: float = 0.0, **kwargs: Any, ) -> None: self.output_def = output_def self.model_path = model_file self.neighbor_list = neighbor_list + self._neighbor_graph_method = neighbor_graph_method + self._neighbor_skin = neighbor_skin # data modifier, populated only for frozen .pth models that carry one self.modifier = None if str(self.model_path).endswith(".pt"): @@ -282,22 +290,42 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None: ``"native"`` uses the dense all-pairs builder; ``"vesin"`` / ``"nv"`` force the O(N) ``vesin.torch`` / ``nvalchemiops`` cell list (raising if unavailable or the model/inputs are unsupported); ``"auto"`` picks the - first available O(N) builder (vesin, then nv) and otherwise falls back to - the native builder. Results are unchanged either way -- only the + available O(N) builder by frame count and otherwise falls back to the + native builder. Results are unchanged either way -- only the neighbor-search cost differs. """ inner = self.dp.model["Default"] self_built = getattr(inner, "use_self_built_nlist", None) if callable(self_built) and self_built(): - # The model builds its own neighbor list and runs the native path; - # an external strategy would bypass it, so always use native. - log.debug( - "Ignoring nlist_backend=%r: %s uses its own built-in neighbor list.", - nlist_backend, - type(inner).__name__, - ) + # The strategy belongs to the model's geometry stage. Keeping the + # native forward entry preserves its compiled core computation. + configure = getattr(inner, "set_inference_neighbor_list", None) self._nlist_builder = None + if not callable(configure): + return + method = self._neighbor_graph_method + if method not in ("auto", "vesin", "nv"): + raise ValueError( + "PT SeZM/DPA4 supports neighbor_graph_method 'auto', " + f"'vesin', or 'nv'; got {method!r}." + ) + if method != "auto" and nlist_backend != "auto": + raise ValueError( + "Select only one of neighbor_graph_method and nlist_backend." + ) + if method == "auto": + method = nlist_backend + configure( + make_neighbor_list( + method, DEVICE, self._neighbor_skin, return_mode="edges" + ) + ) return + if self._neighbor_graph_method != "auto": + raise ValueError( + "neighbor_graph_method applies to PT SeZM/DPA4 models; " + "use nlist_backend for this model." + ) if nlist_backend not in ("auto", "vesin", "nv", "native"): raise ValueError( f"Unknown nlist_backend '{nlist_backend}'; " @@ -325,7 +353,6 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None: unsupported = "non-energy models" ase_provided = self.neighbor_list is not None - builder = None if nlist_backend in ("vesin", "nv"): if unsupported is not None: raise ValueError( @@ -337,36 +364,11 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None: f"nlist_backend='{nlist_backend}' conflicts with an " "explicitly supplied ASE neighbor_list; pass only one." ) - if nlist_backend == "vesin": - if not is_vesin_torch_available(): - raise ImportError( - "nlist_backend='vesin' was requested but 'vesin.torch' " - "is not installed. Install it (`pip install " - "vesin[torch]`) or use nlist_backend='native' (or 'auto')." - ) - builder = VesinNeighborList() - elif DEVICE.type != "cuda": - raise ValueError( - "nlist_backend='nv' requires CUDA inference tensors; " - f"current DEVICE is {DEVICE!s}. Use nlist_backend='native' " - "(or 'auto') for CPU inference." - ) - elif not is_nv_available(): - raise ImportError( - "nlist_backend='nv' was requested but 'nvalchemi-toolkit-ops'" - " is not installed. Install it (`pip install " - "nvalchemi-toolkit-ops`) or use nlist_backend='native' " - "(or 'auto')." - ) - else: - builder = NvNeighborList() - elif nlist_backend == "auto" and unsupported is None and not ase_provided: - # Pick the first available O(N) builder; nv is GPU-only. - if is_vesin_torch_available(): - builder = VesinNeighborList() - elif is_nv_available() and DEVICE.type == "cuda": - builder = NvNeighborList() - self._nlist_builder = builder + self._nlist_builder = ( + make_neighbor_list(nlist_backend, DEVICE, self._neighbor_skin) + if unsupported is None and not ase_provided + else None + ) def get_rcut(self) -> float: """Get the cutoff radius of this model.""" diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index b1fdf148db..35a87e3294 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -511,17 +511,12 @@ strip_saved_tensor_detach, trace_pad_dim, ) -from deepmd.pt.utils.nv_nlist import ( - NvNeighborList, - is_nv_available, +from deepmd.pt_expt.utils.auto_neighbor_list import ( + AdaptiveNeighborList, ) from deepmd.pt_expt.utils.edge_schema import ( edge_schema_from_extended, ) -from deepmd.pt_expt.utils.vesin_neighbor_list import ( - VesinNeighborList, - is_vesin_torch_available, -) from deepmd.utils.version import ( check_version_compatibility, ) @@ -531,49 +526,6 @@ SeZMModel_ = make_model(SeZMAtomicModel) -def _select_neighbor_builder(nf: int, device: torch.device) -> NeighborList: - """Select the O(N) neighbor builder for the given batch shape and device. - - A single CPU frame uses ``vesin``: its lightweight cell list has no - per-frame batching to amortize and avoids the ``nvalchemiops`` (Warp) - launch overhead. Every other case -- any CUDA input or any multi-frame - batch -- uses ``nvalchemiops``, whose batched kernel amortizes the launch - cost across frames. Both builders keep every neighbor within ``rcut``, so - the resulting edge set never depends on ``sel``. ``vesin`` also serves as - the fallback when ``nvalchemiops`` is absent. - - Parameters - ---------- - nf : int - Number of frames in the batch. - device : torch.device - Device of the input coordinates. - - Returns - ------- - NeighborList - A builder instance whose ``build`` method produces the neighbor list. - - Raises - ------ - RuntimeError - If neither ``nvalchemiops`` nor ``vesin`` is importable. - """ - if device.type == "cpu" and nf == 1 and is_vesin_torch_available(): - return VesinNeighborList() - if is_nv_available(): - return NvNeighborList() - if is_vesin_torch_available(): - return VesinNeighborList() - raise RuntimeError( - "SeZM neighbor-list construction requires a neighbor-list backend, but " - "neither 'vesin' nor 'nvalchemiops' is importable. Install one of:\n" - " pip install 'vesin[torch]' # portable CPU/CUDA cell list\n" - " pip install nvalchemi-toolkit-ops # batched CUDA kernels\n" - "'vesin' also ships with the PyTorch extra: pip install 'deepmd-kit[torch]'." - ) - - # Apply the process-global PyTorch workarounds the compile pipeline relies on # (autotune log suppression, DDP optimiser, and the 2.12 divisibility repair) # once, before the first compilation in this run. @@ -719,6 +671,7 @@ def __init__( self.redu_prec = env.GLOBAL_PT_ENER_FLOAT_PRECISION self.use_compile = bool(use_compile) self.enable_tf32 = bool(enable_tf32) + self._inference_neighbor_list: NeighborList | None = None # LoRA injection happens in Trainer.__init__ after pre-trained state is loaded. self.lora_config: dict[str, Any] | None = None if lora is None else dict(lora) self._dens_compiled = False @@ -2764,6 +2717,19 @@ def use_self_built_nlist(self) -> bool: """ return True + def set_inference_neighbor_list(self, neighbor_list: NeighborList | None) -> None: + """Select inference geometry construction while retaining native forward. + + The strategy is runtime state, independent of model parameters and + serialization. Training retains its own batch-aware builder policy. + """ + self._inference_neighbor_list = neighbor_list + + def _get_neighbor_builder(self) -> NeighborList: + if not self.training and self._inference_neighbor_list is not None: + return self._inference_neighbor_list + return AdaptiveNeighborList(method="native") + def build_neighbor_list( self, coord: Float[Tensor, "nf nloc 3"] | Float[Tensor, "nf nloc_x3"], @@ -2772,7 +2738,7 @@ def build_neighbor_list( ) -> EdgeNeighborList: """Build the unified edge-vector schema for the ``forward`` entry.""" nf, nloc = atype.shape[:2] - return _select_neighbor_builder(nf, coord.device).build( + return self._get_neighbor_builder().build( coord.view(nf, nloc, 3), atype, box, @@ -2799,7 +2765,7 @@ def build_extended_neighbor_list( contract order ``(extended_coord, extended_atype, nlist, mapping)``. """ nf, nloc = atype.shape[:2] - return _select_neighbor_builder(nf, coord.device).build( + return self._get_neighbor_builder().build( coord.view(nf, nloc, 3), atype, box, diff --git a/deepmd/pt/utils/compile_compat.py b/deepmd/pt/utils/compile_compat.py index f82307b488..fd93f57032 100644 --- a/deepmd/pt/utils/compile_compat.py +++ b/deepmd/pt/utils/compile_compat.py @@ -634,6 +634,11 @@ def build_inductor_compile_options(*, inference: bool = False) -> dict[str, Any] if tune_train >= 2: compile_options["max_autotune_gemm"] = True if inference: + # PyTorch 2.11 CPU horizontal fusion can reuse a scatter's index + # array in an earlier gather, before the array is declared. Keep + # those loops separate while retaining SIMD and vertical fusion. + if _torch_release() == (2, 11): + compile_options["cpp.max_horizontal_fusion_size"] = 1 # The peak-memory reordering pass sizes buffers through # ``sizevars.size_hint(numel, fallback=0)``. The inference graph is # lowered from ``make_fx`` fake placeholders whose edge-count symbols diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 9aeae29ca3..00a5cc219c 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -1,13 +1,21 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Compress a pt_expt model (.pte) by tabulating embedding nets.""" +"""Build and export compressed pt_expt inference models.""" import logging +from typing import ( + TYPE_CHECKING, +) from deepmd.pt_expt.utils.serialization import ( deserialize_to_file, serialize_from_file, ) +if TYPE_CHECKING: + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + log = logging.getLogger(__name__) @@ -19,14 +27,14 @@ def enable_compression( check_frequency: int = -1, training_script: str | None = None, ) -> None: - """Compress a .pte model by tabulating embedding nets. + """Compress a .pte or .pt2 model by tabulating embedding nets. Parameters ---------- input_file : str - Path to the input .pte model file. + Path to the input .pte or .pt2 model file. output : str - Path to the output compressed .pte model file. + Path to the output compressed .pte or .pt2 model file. stride : float The uniform stride of the first table. extrapolate : int @@ -41,16 +49,60 @@ def enable_compression( BaseModel, ) - # 1. Load the .pte model model_dict = serialize_from_file(input_file) model = BaseModel.deserialize(model_dict["model"]) + compress_model( + model, + output, + stride=stride, + extrapolate=extrapolate, + check_frequency=check_frequency, + training_script=training_script, + min_nbor_dist=model_dict.get("min_nbor_dist"), + model_def_script=model_dict.get("model_def_script"), + ) - # 2. Get or compute min_nbor_dist. Measuring it is a dense all-pairs pass - # over the training data, so it is only run for models that tabulate - # from the shortest observed distance. - min_nbor_dist = model.get_min_nbor_dist() - if min_nbor_dist is None: - min_nbor_dist = model_dict.get("min_nbor_dist") + +def compress_model( + model: "BaseModel", + output: str, + *, + stride: float = 0.01, + extrapolate: int = 5, + check_frequency: int = -1, + training_script: str | None = None, + min_nbor_dist: float | None = None, + model_def_script: dict | None = None, + with_comm: bool = True, +) -> None: + """Compress a loaded model and export its inference artifact. + + Parameters + ---------- + model + Model whose parameters and output statistics are loaded. + output + Destination ``.pte`` or ``.pt2`` file. + stride + Uniform spacing of the first compression table. + extrapolate + Scale of the extrapolation region for models that require it. + check_frequency + Frequency of table overflow checks. + training_script + Training configuration used when neighbor statistics are required + but absent from the model and ``min_nbor_dist``. + min_nbor_dist + Fallback minimum neighbor distance from artifact metadata. + model_def_script + Training model configuration to preserve in the exported artifact. + with_comm + Whether to include a communication artifact for cross-rank models. + """ + # === Step 1. Resolve the domain required by the compression tables. === + model_min_nbor_dist = model.get_min_nbor_dist() + if model_min_nbor_dist is not None: + min_nbor_dist = model_min_nbor_dist if min_nbor_dist is None and not model.compression_needs_min_nbor_dist(): log.info( "The model tabulates over an analytically bounded domain; " @@ -64,9 +116,10 @@ def enable_compression( ) if training_script is None: raise ValueError( - "The model does not have a minimum neighbor distance, " - "so the training script and data must be provided " - "(via -t,--training-script)." + "The model does not have a minimum neighbor distance. " + "Provide training_script to compression, or freeze without " + "--compress and then run 'dp --pt-expt compress -i model.pte " + "-o compressed.pte -t INPUT.json' with the training script and data." ) from deepmd.common import ( j_loader, @@ -95,8 +148,12 @@ def enable_compression( model.min_nbor_dist = min_nbor_dist - # 3. Enable compression (also ensures fake ops are registered now that - # the C++ custom op library is loaded via enable_compression imports) + # === Step 2. Fold deployment constants and build the compressed snapshot. === + # The neutral vacuum reference must be resolved before a charge-conditioned + # descriptor specializes its tables to the deployment charge state. + model.eval() + model.fold_vacuum_reference() + uncompressed_model_dict = model.serialize() from deepmd.pt_expt.utils.tabulate_ops import ( ensure_fake_registered, ) @@ -111,16 +168,14 @@ def enable_compression( check_frequency, ) - # 4. Serialize the compressed model dict (includes tabulated data) compressed_model_dict = model.serialize() - # 5. Re-export the compressed model. + # === Step 3. Export the graph supported by the compressed model. === # - # A geometrically compressed graph-lower descriptor (DPA1 / se_atten strip, - # attn_layer == 0) evaluates its tabulated embedding through fused CUDA - # operators that make_fx can trace, so its compressed graph exports directly - # to a graph-lower ``.pt2``: ``deserialize_to_file`` bakes the end-to-end - # fused table operator and the mandatory per-atom virial (see its docstring). + # Compressed graph descriptors use fused operators with traceable fake + # implementations. Their graphs export directly to ``.pt2``; the exporter + # selects a compact canonical schema when the target supports it and + # preserves the per-atom virial alongside the total energy and forces. # Every other compressed descriptor keeps tabulated operators make_fx cannot # trace: those export the UNCOMPRESSED graph and carry the compressed dict in # ``model.json`` so ``deserialize()`` restores the compression state for the @@ -129,23 +184,24 @@ def enable_compression( model_uses_graph_lower, ) - model_def_script = model_dict.get("model_def_script") if output.endswith(".pt2") and model_uses_graph_lower(model): log.info("Re-exporting compressed graph...") deserialize_to_file( output, {"model": compressed_model_dict, "model_def_script": model_def_script}, lower_kind="auto", + with_comm=with_comm, ) else: log.info("Re-exporting compressed model...") deserialize_to_file( output, - {"model": model_dict["model"], "model_def_script": model_def_script}, + {"model": uncompressed_model_dict, "model_def_script": model_def_script}, model_json_override={ "model": compressed_model_dict, "model_def_script": model_def_script, "min_nbor_dist": float(min_nbor_dist), }, + with_comm=with_comm, ) log.info("Compressed model saved to %s", output) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index e4ff2f813d..f6736b0221 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -528,6 +528,9 @@ def freeze( output: str = "frozen_model", head: str | None = None, lower_kind: str = "nlist", + *, + with_comm: bool = True, + compress: bool = False, ) -> None: """Freeze a pt_expt checkpoint into a .pte/.pt2 exported model. @@ -558,6 +561,14 @@ def freeze( graph lower, so ``lower_kind`` is resolved to ``"graph"`` for it regardless of the requested value -- before the export ABI and the default output suffix are selected. + with_comm : bool, optional + Include a communication artifact for models with cross-rank message + passing, by default True. Disable for single-process inference. + compress : bool, optional + Compress the loaded model before export, by default False. Uses the + default settings of ``dp compress`` and selects the accelerated graph + schema for eligible ``.pt2`` models without compiling an intermediate + artifact. """ import torch @@ -668,12 +679,26 @@ def freeze( ) do_atomic_virial = True + if compress: + from deepmd.pt_expt.entrypoints.compress import ( + compress_model, + ) + + compress_model( + m, + output, + model_def_script=single_model_params, + with_comm=with_comm, + ) + return + model_dict_serialized = m.serialize() deserialize_to_file( output, {"model": model_dict_serialized, "model_def_script": single_model_params}, do_atomic_virial=do_atomic_virial, lower_kind=lower_kind, + with_comm=with_comm, ) log.info("Saved frozen model to %s (lower_kind=%s)", output, lower_kind) @@ -942,6 +967,7 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None: output=FLAGS.output, head=FLAGS.head, lower_kind=getattr(FLAGS, "lower_kind", "nlist"), + compress=getattr(FLAGS, "compress", False), ) elif FLAGS.command == "change-bias": change_bias( diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 94c1e1f635..951e5a5176 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -63,12 +63,14 @@ charge_states_per_frame, single_charge_state, ) +from deepmd.pt_expt.utils.auto_neighbor_list import ( + make_neighbor_list, +) from deepmd.pt_expt.utils.edge_schema import ( edge_schema_from_extended, ) -from deepmd.pt_expt.utils.vesin_neighbor_list import ( - VesinNeighborList, - is_vesin_torch_available, +from deepmd.pt_expt.utils.vesin_search import ( + VesinNeighborSearch, ) from deepmd.utils.charge_state import ( CHARGE_STATE_TABLE_RANGES, @@ -86,6 +88,9 @@ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) + from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + ) # Public output keys emitted by graph-lower forwards, keyed by the @@ -179,25 +184,28 @@ class DeepEval(DeepEvalBackend): 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 rejected for graph-routed artifacts. + ``"auto"`` / ``"vesin"`` / ``"nv"`` / ``"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 via :func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder` - at each eval call (CUDA: ``nv`` if importable; else ``vesin`` only when - ``nf == 1`` and importable; else ``dense``). Explicit + at each eval call (CUDA: ``vesin`` for one frame if importable, then + ``nv`` if importable, otherwise ``dense``). Explicit ``"dense"`` / ``"ase"`` / ``"cell"`` / ``"vesin"`` / ``"nv"`` choices are preserved. The CPU-only ``cell`` search runs on the host; its graph tensors are transferred to the model device before inference. - 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. Training keeps a separate auto policy + Edge-vector artifacts accept ``"vesin"`` and ``"nv"`` through the same + option. A non-default value on a nlist artifact raises at construction; + use ``nlist_backend`` for that path. All builders emit the same neighbor + set, so the choice is performance-only. Training keeps a separate auto policy (:func:`~deepmd.pt_expt.utils.graph_builder.resolve_neighbor_graph_method`) that never selects ``vesin``. Consolidating the two knobs into a single backend-selection API is deferred to the dense-nlist deprecation. + neighbor_skin : float, default: 0.0 + Vesin candidate-search skin in angstrom, shared by graph, edge-vector, + and extended-coordinate inputs. Zero disables Verlet caching. **kwargs : dict Keyword arguments. """ @@ -211,14 +219,24 @@ def __init__( neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, nlist_backend: str = "auto", neighbor_graph_method: str = "auto", + neighbor_skin: float = 0.0, **kwargs: Any, ) -> None: + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + self.output_def = output_def self.model_path = model_file self.neighbor_list = neighbor_list # Graph-lower builder selection is resolved once after model metadata # identifies the lower ABI. self._neighbor_graph_method = neighbor_graph_method + self._neighbor_skin = neighbor_skin + # CUDA destination grouping belongs to the graph CSR construction. + self._vesin_search = VesinNeighborSearch( + skin=neighbor_skin, sorted=DEVICE.type != "cuda" + ) self._is_pt2 = model_file.endswith(".pt2") # Only a compressed ``.pt2`` folds its charge state into constants; a # model that reads the condition as an ordinary input needs no rebuild. @@ -242,21 +260,22 @@ 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-lower evaluation. - # Fail fast instead of silently ignoring it on nlist-form artifacts, - # where the corresponding builder knob is ``nlist_backend``. + # ``neighbor_graph_method`` selects searches for graph and edge-vector + # inputs. Extended-coordinate inputs use ``nlist_backend`` instead. if neighbor_graph_method != "auto" and getattr(self, "metadata", {}).get( "lower_input_kind" ) not in ( "graph", "dpa1_canonical", "dpa4c_canonical", + "edge_vec", ): raise ValueError( f"neighbor_graph_method={neighbor_graph_method!r} only applies to " - "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." + "graph or edge-vector artifacts (lower_input_kind: 'graph', " + "'dpa1_canonical', 'dpa4c_canonical', or 'edge_vec'). " + f"This model uses {self.metadata.get('lower_input_kind')!r}; " + "use nlist_backend to select its neighbor-list builder." ) self._setup_neighbor_backend(nlist_backend) @@ -303,18 +322,22 @@ def _resolve_neighbor_graph_method(method: str, nf: int | None = None) -> str: 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. Graph-routed + ``"native"`` uses the dense all-pairs builder; ``"vesin"`` and ``"nv"`` + select accelerated searches, raising for unavailable or unsupported + inputs. ``"auto"`` selects by frame count and available libraries, + falling back to the native builder when none applies. 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"): + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + if nlist_backend not in ("auto", "vesin", "nv", "native"): raise ValueError( f"Unknown nlist_backend '{nlist_backend}'; " - "expected 'auto', 'vesin', or 'native'." + "expected 'auto', 'vesin', 'nv', or 'native'." ) if self.metadata.get("lower_input_kind") in ( "graph", @@ -339,40 +362,57 @@ def _setup_neighbor_backend(self, nlist_backend: str) -> None: self._neighbor_graph_method = self._resolve_neighbor_graph_method( self._neighbor_graph_method ) - self._use_vesin = False self._nlist_builder = None return + if ( + self.metadata.get("lower_input_kind") == "edge_vec" + and self._neighbor_graph_method != "auto" + ): + if self._neighbor_graph_method not in ("vesin", "nv"): + raise ValueError( + "Edge-vector artifacts support neighbor_graph_method " + "'auto', 'vesin', or 'nv'; " + f"got {self._neighbor_graph_method!r}." + ) + if nlist_backend != "auto": + raise ValueError( + "Select only one of neighbor_graph_method and nlist_backend." + ) + nlist_backend = self._neighbor_graph_method is_spin = bool(getattr(self, "_is_spin", False)) - # Graph-routed artifacts returned above. Any remaining spin model uses - # the virtual-atom nlist scheme, which vesin does not support. + # Native-spin edge artifacts share the energy neighbor strategy. + # Virtual-atom spin expansion retains the native nlist path. ase_provided = self.neighbor_list is not None - # reason vesin cannot be used (None means it can) - unsupported = "spin models" if is_spin else None - if nlist_backend == "native": - self._use_vesin = False - elif nlist_backend == "vesin": - if not is_vesin_torch_available(): - raise ImportError( - "nlist_backend='vesin' was requested but 'vesin.torch' is " - "not installed. Install it (`pip install vesin[torch]`) or " - "use nlist_backend='native' (or 'auto')." - ) + unsupported = ( + "spin models" + if is_spin and self.metadata.get("lower_input_kind") != "edge_vec" + else None + ) + if nlist_backend in ("vesin", "nv"): if unsupported is not None: raise ValueError( - f"nlist_backend='vesin' is not supported for {unsupported}; " + f"nlist_backend={nlist_backend!r} is not supported for {unsupported}; " "use nlist_backend='native' (or 'auto')." ) if ase_provided: raise ValueError( - "nlist_backend='vesin' conflicts with an explicitly " + f"nlist_backend={nlist_backend!r} conflicts with an explicitly " "supplied ASE neighbor_list; pass only one." ) - self._use_vesin = True - else: # auto: use vesin when possible, otherwise fall back silently - self._use_vesin = ( - is_vesin_torch_available() and unsupported is None and not ase_provided + self._nlist_builder = ( + make_neighbor_list( + nlist_backend, + DEVICE, + self._neighbor_skin, + return_mode=( + "edges" + if self.metadata.get("lower_input_kind") == "edge_vec" + else "extended" + ), ) - self._nlist_builder = VesinNeighborList() if self._use_vesin else None + if unsupported is None and not ase_provided + else None + ) def _init_from_model_json(self, model_json_str: str) -> None: """Deserialize model.json and derive model API from the dpmodel instance.""" @@ -473,7 +513,10 @@ def _load_pte(self, model_file: str) -> None: "model_def_script.json": "", "metadata.json": "", } - exported = torch.export.load(model_file, extra_files=extra_files) + # Archive loading stages raw storages on CPU before restoring their + # recorded devices, including when the caller selects another default. + with torch.device("cpu"): + exported = torch.export.load(model_file, extra_files=extra_files) self.exported_module = exported.module() mds = extra_files["model_def_script.json"] self._model_def_script = json.loads(mds) if mds else {} @@ -1639,28 +1682,6 @@ def _build_nlist_ase_single( return extended_coord, extended_atype, nlist, mapping - @staticmethod - def _build_edge_inputs_from_nlist( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Convert a padded neighbor list into the compact-edge schema.""" - nloc = nlist.shape[1] - schema = edge_schema_from_extended( - extended_coord, - extended_atype[:, :nloc], - nlist, - mapping, - ) - return ( - schema.edge_index, - schema.edge_vec, - schema.edge_scatter_index, - schema.edge_mask, - ) - def _prepare_nlist_inputs( self, coords: np.ndarray, @@ -1719,47 +1740,9 @@ def _prepare_nlist_inputs( atype_t, ) - if fparam is not None: - fparam_t = torch.tensor( - fparam.reshape(nframes, self.get_dim_fparam()), - dtype=torch.float64, - device=DEVICE, - ) - elif self.get_dim_fparam() > 0: - # Exported models (.pt2/.pte) are compiled with fparam as a - # required input. Fill with default values from metadata. - default_fp = self.metadata.get("default_fparam") - if default_fp is not None: - fparam_t = ( - torch.tensor(default_fp, dtype=torch.float64, device=DEVICE) - .unsqueeze(0) - .expand(nframes, -1) - .contiguous() - ) - else: - raise ValueError( - f"fparam is required for this model (dim_fparam={self.get_dim_fparam()}) " - "but was not provided, and no default_fparam is stored in the model." - ) - else: - fparam_t = None - - if aparam is not None: - aparam_t = torch.tensor( - aparam.reshape(nframes, natoms, self.get_dim_aparam()), - dtype=torch.float64, - device=DEVICE, - ) - elif self.get_dim_aparam() > 0: - # Exported models (.pt2/.pte) are compiled with aparam as a - # required positional input. Unlike fparam, there is no default. - raise ValueError( - f"aparam is required for this model (dim_aparam={self.get_dim_aparam()}) " - "but was not provided." - ) - else: - aparam_t = None - + fparam_t, aparam_t = self._prepare_optional_lower_inputs( + fparam, aparam, nframes, natoms, DEVICE + ) charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) return ( @@ -1857,21 +1840,19 @@ def _prepare_inputs( coords, cells, atom_types, fparam, aparam, charge_spin ) if self.metadata.get("lower_input_kind") == "edge_vec": - edge_index_t, edge_vec_t, edge_scatter_t, edge_mask_t = ( - self._build_edge_inputs_from_nlist( - ext_coord_t, - ext_atype_t, - nlist_t, - mapping_t, - ) - ) - model_inputs = ( + edge_schema = edge_schema_from_extended( ext_coord_t, ext_atype_t[:, :natoms], - edge_index_t, - edge_vec_t, - edge_scatter_t, - edge_mask_t, + nlist_t, + mapping_t, + ) + model_inputs = ( + edge_schema.coord, + edge_schema.atype, + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_scatter_index, + edge_schema.edge_mask, fparam_t, aparam_t, charge_spin_t, @@ -1904,6 +1885,8 @@ def _prepare_optional_lower_inputs( device=device, ) elif self.get_dim_fparam() > 0: + # Compiled lower interfaces require a tensor for frame parameters; + # omitted values use the defaults stored in the model metadata. default_fp = self.metadata.get("default_fparam") if default_fp is None: raise ValueError( @@ -1926,6 +1909,7 @@ def _prepare_optional_lower_inputs( device=device, ) elif self.get_dim_aparam() > 0: + # Atomic parameters have no model-level default. raise ValueError( f"aparam is required for this model (dim_aparam={self.get_dim_aparam()}) " "but was not provided." @@ -2022,130 +2006,29 @@ def _eval_model_spin( request_defs, charge_spin=charge_spin, ) - nframes = coords.shape[0] - if len(atom_types.shape) == 1: - natoms = len(atom_types) - atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) - else: - natoms = len(atom_types[0]) - from deepmd.pt_expt.utils.env import ( DEVICE, ) - coord_input = coords.reshape(nframes, natoms, 3) - if self.neighbor_list is not None: - extended_coord, extended_atype, nlist, mapping = self._build_nlist_ase( - coord_input, - cells, - atom_types, - ) - ext_coord_t = torch.tensor( - extended_coord, dtype=torch.float64, device=DEVICE - ) - ext_atype_t = torch.tensor(extended_atype, dtype=torch.int64, device=DEVICE) - nlist_t = torch.tensor(nlist, dtype=torch.int64, device=DEVICE) - mapping_t = torch.tensor(mapping, dtype=torch.int64, device=DEVICE) - else: - coord_t = torch.tensor(coord_input, dtype=torch.float64, device=DEVICE) - atype_t = torch.tensor(atom_types, dtype=torch.int64, device=DEVICE) - cells_t = ( - torch.tensor(cells, dtype=torch.float64, device=DEVICE) - if cells is not None - else None - ) - ext_coord_t, ext_atype_t, nlist_t, mapping_t = self._build_nlist_native( - coord_t, - cells_t, - atype_t, - ) - - # Extend spin to ghost atoms using mapping + model_inputs, mapping_t, nframes, natoms = self._prepare_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin + ) spin_t = torch.tensor( spins.reshape(nframes, natoms, 3), dtype=torch.float64, device=DEVICE ) - batch_idx = ( - torch.arange(nframes, dtype=torch.long, device=DEVICE) - .unsqueeze(1) - .expand_as(mapping_t) - ) - ext_spin_t = spin_t[batch_idx, mapping_t] - - if fparam is not None: - fparam_t = torch.tensor( - fparam.reshape(nframes, self.get_dim_fparam()), - dtype=torch.float64, - device=DEVICE, - ) - elif self.get_dim_fparam() > 0: - # Exported models (.pt2/.pte) are compiled with fparam as a - # required input. Fill with default values from metadata. - default_fp = self.metadata.get("default_fparam") - if default_fp is not None: - fparam_t = ( - torch.tensor(default_fp, dtype=torch.float64, device=DEVICE) - .unsqueeze(0) - .expand(nframes, -1) - .contiguous() - ) - else: - raise ValueError( - f"fparam is required for this model (dim_fparam={self.get_dim_fparam()}) " - "but was not provided, and no default_fparam is stored in the model." - ) - else: - fparam_t = None - - if aparam is not None: - aparam_t = torch.tensor( - aparam.reshape(nframes, natoms, self.get_dim_aparam()), - dtype=torch.float64, - device=DEVICE, - ) - elif self.get_dim_aparam() > 0: - raise ValueError( - f"aparam is required for this model (dim_aparam={self.get_dim_aparam()}) " - "but was not provided." - ) - else: - aparam_t = None - - charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) - - # Build the lower inputs for the model's spin ABI. The native scheme - # shares the energy edge contract and feeds the owned-atom spins (the - # descriptor only needs local spins; ghost neighbours resolve to their - # local owners). The deepspin scheme keeps the extended nlist contract. + # Native spin belongs to the local owners of the edge schema. The + # virtual-atom scheme extends it to ghost sites through the same mapping + # used to fold the model's coordinate derivatives back to local atoms. if self.metadata.get("lower_input_kind") == "edge_vec": - edge_schema = edge_schema_from_extended( - ext_coord_t, - ext_atype_t[:, :natoms], - nlist_t, - mapping_t, - ) - model_inputs = ( - edge_schema.coord, - edge_schema.atype, - edge_schema.edge_index, - edge_schema.edge_vec, - edge_schema.edge_scatter_index, - edge_schema.edge_mask, - spin_t, - fparam_t, - aparam_t, - charge_spin_t, - ) + model_inputs = (*model_inputs[:6], spin_t, *model_inputs[6:]) else: - model_inputs = ( - ext_coord_t, - ext_atype_t, - ext_spin_t, - nlist_t, - mapping_t, - fparam_t, - aparam_t, - charge_spin_t, + batch_idx = ( + torch.arange(nframes, dtype=torch.long, device=DEVICE) + .unsqueeze(1) + .expand_as(mapping_t) ) + ext_spin_t = spin_t[batch_idx, mapping_t] + model_inputs = (*model_inputs[:2], ext_spin_t, *model_inputs[2:]) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: @@ -2200,158 +2083,16 @@ def _eval_model_graph_spin( request_defs: list[OutputVariableDef], charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: - """Evaluate a graph-form native-spin ``.pt2`` (``lower_input_kind == - "graph"`` and ``is_spin``). - - Mirrors :meth:`_eval_model_graph`'s carry-all - :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph` - construction (SAME builder, SAME positional ABI up through - ``source_row_ptr``), then inserts the owned-atom ``spin`` tensor - ``(N, 3)`` at positional index 10 of - ``NativeSpinEnergyModel.forward_lower_graph_exportable`` -- the node - axis IS the owned-local-atom axis for single-rank eval (no ghost - nodes), so ``spin`` needs no extension/mapping, unlike the dense - spin path's ``ext_spin_t``. ``charge_spin`` rides the conditional - slot-13 tail (combined native-spin + charge-spin FiLM models; the - slot is dropped from the exported signature otherwise). The forward - returns LOCAL public keys directly (``atom_energy``, - ``energy``, ``force``, ``force_mag``, ``virial``, ``atom_virial``), - so results are reshaped without ``communicate_extended_output``, - same as the non-spin graph path. - """ - from deepmd.pt_expt.utils.env import ( - DEVICE, - ) + """Evaluate native-spin graph or compact canonical inputs. - nframes = coords.shape[0] - if len(atom_types.shape) == 1: - natoms = len(atom_types) - atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) - else: - natoms = len(atom_types[0]) - - coord_input = coords.reshape(nframes, natoms, 3) - box_input = cells.reshape(nframes, 9) if cells is not None else None - graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE) - - atype_t = torch.tensor( - np.asarray(atom_types).reshape(-1), dtype=torch.int64, device=DEVICE - ) - n_node_t = torch.as_tensor(graph.n_node, dtype=torch.int64, device=DEVICE) - edge_index_t = torch.as_tensor( - graph.edge_index, dtype=torch.int64, device=DEVICE - ) - edge_dtype = ( - torch.float32 - if self.metadata.get("graph_edge_dtype") == "float32" - else torch.float64 - ) - edge_vec_t = torch.as_tensor( - graph.edge_vec, - dtype=edge_dtype, - device=DEVICE, - ) - edge_mask_t = torch.as_tensor(graph.edge_mask, dtype=torch.bool, device=DEVICE) - destination_order_t = torch.as_tensor( - graph.destination_order, - device=DEVICE, - ) - destination_row_ptr_t = torch.as_tensor( - graph.destination_row_ptr, - dtype=torch.int64, - device=DEVICE, - ) - source_order_t = torch.as_tensor( - graph.source_order, - device=DEVICE, - ) - source_row_ptr_t = torch.as_tensor( - graph.source_row_ptr, - dtype=torch.int64, - device=DEVICE, - ) - - spin_t = torch.tensor( - np.asarray(spins).reshape(nframes * natoms, 3), - dtype=torch.float64, - device=DEVICE, - ) - - fparam_t, aparam_t = self._prepare_optional_lower_inputs( - fparam, aparam, nframes, natoms, DEVICE + Spin belongs to the local node axis, so no ghost extension is needed. + The shared input preparation places spin before the generic graph's + optional conditioning inputs, or last in the compact canonical ABI. + Both forwards return local public outputs directly. + """ + model_inputs, nframes, natoms = self._prepare_graph_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin, spins=spins ) - if aparam_t is not None: - # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) -- - # the same axis as ``atype``/``spin`` (mirrors _eval_model_graph). - aparam_t = aparam_t.reshape(nframes * natoms, -1) - - if self.metadata.get("lower_input_kind") == "dpa4c_canonical": - # The compact canonical ABI has no conditional tail, so the moment - # is the last slot rather than the eleventh, and any conditioning - # input the artifact cannot receive is an error rather than a - # silently dropped argument. - if ( - self.get_dim_fparam() > 0 - or self.get_dim_aparam() > 0 - or int(self.metadata.get("dim_chg_spin", 0) or 0) > 0 - ): - raise NotImplementedError( - "compact canonical artifacts carry no fparam/aparam/" - "charge_spin inputs; a model requiring them must not be " - "frozen with a canonical lower kind." - ) - # A compressed descriptor still serves a condition, through the - # constants its fold rewrites rather than through an argument. - self._apply_charge_state(charge_spin) - from deepmd.dpmodel.utils.neighbor_graph import ( - NeighborGraph, - ) - from deepmd.pt_expt.utils.canonical_graph import ( - canonical_graph_from_neighbor_graph, - ) - - compact = canonical_graph_from_neighbor_graph( - NeighborGraph( - n_node=n_node_t, - edge_index=edge_index_t, - edge_vec=edge_vec_t, - edge_mask=edge_mask_t, - n_local=n_node_t, - destination_order=destination_order_t, - destination_row_ptr=destination_row_ptr_t, - source_order=source_order_t, - source_row_ptr=source_row_ptr_t, - destination_sorted=bool(graph.destination_sorted), - ) - ) - model_inputs = ( - atype_t, - compact.n_node, - compact.n_local, - compact.source, - compact.edge_vec, - compact.destination_row_ptr, - compact.source_row_ptr, - compact.source_order, - spin_t.to(torch.float32), - ) - else: - model_inputs = ( - atype_t, - n_node_t, - n_node_t, - edge_index_t, - edge_vec_t, - edge_mask_t, - destination_order_t, - destination_row_ptr_t, - source_order_t, - source_row_ptr_t, - spin_t, - fparam_t, - aparam_t, - self._make_charge_spin_input(nframes, charge_spin), - ) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: @@ -2380,85 +2121,113 @@ def _eval_model_graph( request_defs: list[OutputVariableDef], charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: - """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 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``. + """Evaluate generic graph or compact canonical inputs. + + Exported artifacts accept a dynamic edge axis. Input preparation uses + the artifact's geometry dtype and layout; both forwards return local + public outputs without extended-atom communication. """ + model_inputs, nframes, natoms = self._prepare_graph_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin + ) + if self._is_pt2: + model_ret = self._pt2_runner(*model_inputs) + else: + model_ret = self.exported_module(*model_inputs) + + results = [] + for odef in request_defs: + shape = self._get_output_shape(odef, nframes, natoms) + gkey = _GRAPH_CATEGORY_TO_KEY.get(odef.category) + val = model_ret.get(gkey) if gkey is not None else None + if val is not None: + results.append(val.detach().cpu().numpy().reshape(shape)) + else: + results.append( + np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION) + ) + return tuple(results) + + def _prepare_graph_inputs( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None, + aparam: np.ndarray | None, + charge_spin: np.ndarray | None, + *, + spins: np.ndarray | None = None, + ) -> tuple[tuple[torch.Tensor | None, ...], int, int]: + """Pack local graph inputs according to the exported artifact's ABI.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + canonical_graph_from_neighbor_graph, + ) from deepmd.pt_expt.utils.env import ( DEVICE, ) nframes = coords.shape[0] - if len(atom_types.shape) == 1: + if atom_types.ndim == 1: natoms = len(atom_types) - atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) + atom_types = np.tile(atom_types, (nframes, 1)) else: - natoms = len(atom_types[0]) - - coord_input = coords.reshape(nframes, natoms, 3) - box_input = cells.reshape(nframes, 9) if cells is not None else None - # Build the carry-all graph at its exact edge count; the exported edge - # axis is dynamic. - graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE) - - atype_t = torch.tensor( - np.asarray(atom_types).reshape(-1), dtype=torch.int64, device=DEVICE + natoms = atom_types.shape[1] + graph = self._build_eval_graph( + coords.reshape(nframes, natoms, 3), + atom_types, + cells.reshape(nframes, 9) if cells is not None else None, + DEVICE, + ) + atype_t = torch.as_tensor( + atom_types.reshape(-1), dtype=torch.int64, device=DEVICE ) - # graph fields may be numpy (dense/ase) or torch, possibly on CUDA - # (vesin/nv) -- torch.as_tensor handles both and moves to DEVICE. - n_node_t = torch.as_tensor(graph.n_node, dtype=torch.int64, device=DEVICE) - edge_index_t = torch.as_tensor( - graph.edge_index, dtype=torch.int64, device=DEVICE + canonical = self.metadata.get("lower_input_kind") in ( + "dpa1_canonical", + "dpa4c_canonical", ) edge_dtype = ( torch.float32 if self.metadata.get("graph_edge_dtype") == "float32" else torch.float64 ) - edge_vec_t = torch.as_tensor( - graph.edge_vec, - dtype=edge_dtype, - device=DEVICE, - ) - edge_mask_t = torch.as_tensor(graph.edge_mask, dtype=torch.bool, device=DEVICE) - destination_order_t = torch.as_tensor( - graph.destination_order, - device=DEVICE, - ) - destination_row_ptr_t = torch.as_tensor( - graph.destination_row_ptr, - dtype=torch.int64, - device=DEVICE, - ) - source_order_t = torch.as_tensor( - graph.source_order, - device=DEVICE, - ) - source_row_ptr_t = torch.as_tensor( - graph.source_row_ptr, - dtype=torch.int64, - device=DEVICE, - ) - - if self.metadata.get("lower_input_kind") in ( - "dpa1_canonical", - "dpa4c_canonical", - ): - # The canonical ABI has NO fparam/aparam/charge_spin slots; the - # export gate (fitting_eligible) rejects such models today, so - # this is unreachable -- assert it loudly so a future loosening - # of the eligibility fails here instead of silently dropping - # conditioning inputs at inference. + if not isinstance(graph, CanonicalGraph): + # CPU builders return NumPy arrays; device builders return tensors. + # All graph fields cross the model-device boundary together. + n_node = torch.as_tensor(graph.n_node, dtype=torch.int64, device=DEVICE) + graph = NeighborGraph( + n_node=n_node, + n_local=n_node, + edge_index=torch.as_tensor( + graph.edge_index, dtype=torch.int64, device=DEVICE + ), + edge_vec=torch.as_tensor( + graph.edge_vec, dtype=edge_dtype, device=DEVICE + ), + edge_mask=torch.as_tensor( + graph.edge_mask, dtype=torch.bool, device=DEVICE + ), + destination_order=torch.as_tensor( + graph.destination_order, device=DEVICE + ), + destination_row_ptr=torch.as_tensor( + graph.destination_row_ptr, dtype=torch.int64, device=DEVICE + ), + source_order=torch.as_tensor(graph.source_order, device=DEVICE), + source_row_ptr=torch.as_tensor( + graph.source_row_ptr, dtype=torch.int64, device=DEVICE + ), + destination_sorted=bool(graph.destination_sorted), + ) + if canonical: + graph = canonical_graph_from_neighbor_graph(graph) + if canonical: + # The compact ABI has no conditional input slots. Compressed charge + # conditioning instead updates the constants through its fold. if ( self.get_dim_fparam() > 0 or self.get_dim_aparam() > 0 @@ -2467,84 +2236,53 @@ def _eval_model_graph( raise NotImplementedError( "compact canonical artifacts carry no fparam/aparam/" "charge_spin inputs; a model requiring them must not be " - "frozen with a canonical lower kind (the export eligibility " - "gate should have rejected it)." + "frozen with a canonical lower kind." ) - # A compressed descriptor still serves a condition, through the - # constants its fold rewrites rather than through an argument. self._apply_charge_state(charge_spin) - from deepmd.dpmodel.utils.neighbor_graph import ( - NeighborGraph, - ) - from deepmd.pt_expt.utils.canonical_graph import ( - canonical_graph_from_neighbor_graph, - ) - - generic_graph = NeighborGraph( - n_node=n_node_t, - edge_index=edge_index_t, - edge_vec=edge_vec_t, - edge_mask=edge_mask_t, - n_local=n_node_t, - destination_order=destination_order_t, - destination_row_ptr=destination_row_ptr_t, - source_order=source_order_t, - source_row_ptr=source_row_ptr_t, - destination_sorted=bool(graph.destination_sorted), - ) - compact = canonical_graph_from_neighbor_graph(generic_graph) model_inputs = ( atype_t, - compact.n_node, - compact.n_local, - compact.source, - compact.edge_vec, - compact.destination_row_ptr, - compact.source_row_ptr, - compact.source_order, + graph.n_node, + graph.n_local, + graph.source, + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, ) else: + model_inputs = ( + atype_t, + graph.n_node, + graph.n_local, + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + ) + if spins is not None: + spin_t = torch.as_tensor( + spins.reshape(nframes * natoms, 3), + dtype=torch.float32 if canonical else torch.float64, + device=DEVICE, + ) + model_inputs = (*model_inputs, spin_t) + if not canonical: fparam_t, aparam_t = self._prepare_optional_lower_inputs( fparam, aparam, nframes, natoms, DEVICE ) if aparam_t is not None: - # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) - # -- the same axis as ``atype`` (the shared helper above - # returns the dense rectangular (nf, natoms, nda) layout). + # Graph inputs flatten frame and atom axes into one node axis. aparam_t = aparam_t.reshape(nframes * natoms, -1) - charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) model_inputs = ( - atype_t, - n_node_t, - n_node_t, - edge_index_t, - edge_vec_t, - edge_mask_t, - destination_order_t, - destination_row_ptr_t, - source_order_t, - source_row_ptr_t, + *model_inputs, fparam_t, aparam_t, - charge_spin_t, + self._make_charge_spin_input(nframes, charge_spin), ) - if self._is_pt2: - model_ret = self._pt2_runner(*model_inputs) - else: - model_ret = self.exported_module(*model_inputs) - - results = [] - for odef in request_defs: - shape = self._get_output_shape(odef, nframes, natoms) - gkey = _GRAPH_CATEGORY_TO_KEY.get(odef.category) - val = model_ret.get(gkey) if gkey is not None else None - if val is not None: - results.append(val.detach().cpu().numpy().reshape(shape)) - else: - results.append( - np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION) - ) - return tuple(results) + return model_inputs, nframes, natoms def _build_eval_graph( self, @@ -2552,8 +2290,8 @@ def _build_eval_graph( atom_types: np.ndarray, box_input: np.ndarray | None, device: "torch.device", - ) -> "NeighborGraph": - """Build the carry-all NeighborGraph for graph-lower inference. + ) -> "NeighborGraph | CanonicalGraph": + """Build the local graph for the artifact's input layout. Dispatches on ``self._neighbor_graph_method``: ``auto`` is resolved call-time via @@ -2565,7 +2303,8 @@ def _build_eval_graph( (carry-all, sel-free), so the selection is a pure performance choice and results are unchanged. The result is canonicalized to the destination-major graph-form ``.pt2`` ABI after construction; the - caller transfers its fields to the model device. + caller transfers its fields to the model device. Vesin emits the + compact canonical layout directly when no pair exclusions are needed. """ method = self._neighbor_graph_method if method == "auto": @@ -2672,9 +2411,23 @@ def _build_eval_graph( ) if method == "vesin": from deepmd.pt_expt.utils.vesin_graph_builder import ( + build_canonical_graph_vesin, build_neighbor_graph_vesin, ) + if pair_excl is None and self.metadata.get("lower_input_kind") in ( + "dpa1_canonical", + "dpa4c_canonical", + ): + return build_canonical_graph_vesin( + cc, + bb.reshape(-1, 3, 3) if bb is not None else None, + self._rcut, + self._vesin_search, + node_mask=aa >= 0 + if (np.asarray(atom_types) < 0).any() + else None, + ) return build_neighbor_graph_vesin( cc, aa, @@ -2682,6 +2435,7 @@ def _build_eval_graph( self._rcut, canonicalize=True, pair_excl=pair_excl, + search=self._vesin_search, ) from deepmd.pt_expt.utils.nv_graph_builder import ( build_neighbor_graph_nv, diff --git a/deepmd/pt_expt/kernels/dpa4c/graph_compress.py b/deepmd/pt_expt/kernels/dpa4c/graph_compress.py index a7805d5c1a..b114330f75 100644 --- a/deepmd/pt_expt/kernels/dpa4c/graph_compress.py +++ b/deepmd/pt_expt/kernels/dpa4c/graph_compress.py @@ -92,6 +92,7 @@ "build_charge_state_artifacts", "build_compression_artifacts", "build_radial_table", + "compression_ineligibility_reason", "coupling_records", "descriptor_profile", "dpa4c_graph_compress", @@ -212,6 +213,55 @@ def mega_eligible(descriptor: Any) -> bool: ) +def compression_ineligibility_reason(descriptor: Any) -> str | None: + """Return why a DPA4C descriptor cannot build a compressed snapshot. + + The same device-independent conditions govern explicit compression and + automatic deployment. Kernel availability is checked by the inference + dispatcher rather than by artifact construction. + + Parameters + ---------- + descriptor + DPA4C descriptor with its configuration and parameters loaded. + + Returns + ------- + str or None + Diagnostic for an unsupported configuration or initial charge state, + or ``None`` when these compression prerequisites are satisfied. + """ + if descriptor.exclude_types: + return ( + "DPA4C compressed inference has no type-exclusion branch, so a " + "descriptor with a non-empty `exclude_types` cannot reach the " + "fused path. Drop the exclusions, or deploy the uncompressed " + "graph archive." + ) + if descriptor.stddev.dtype != torch.float32: + return ( + "DPA4C compression requires descriptor precision `float32`, " + f"got {descriptor.stddev.dtype}" + ) + if not mega_eligible(descriptor): + return ( + "DPA4C compression supports " + f"channels {SUPPORTED_CHANNELS}, lmax {SUPPORTED_LMAX} and " + f"radial_modes {SUPPORTED_RADIAL_MODES}, got " + f"channels={descriptor.channels}, lmax={descriptor.lmax}, " + f"radial_modes={descriptor.radial_modes}" + ) + if ( + descriptor.charge_spin_embedding is not None + and descriptor.default_chg_spin is None + ): + return ( + "A charge-conditioned DPA4C snapshot requires `default_chg_spin` " + "to supply its initial frame condition." + ) + return None + + # === Compiled profile === @@ -673,27 +723,10 @@ def build_compression_artifacts( If the descriptor excludes type pairs, or is not an fp32 model with a compiled specialization. """ - if getattr(descriptor, "exclude_types", None): - raise ValueError( - "DPA4C compressed CUDA has no type-exclusion branch, so a " - "descriptor with a non-empty `exclude_types` cannot reach the " - "fused path. Drop the exclusions, or deploy the uncompressed " - "graph archive." - ) + reason = compression_ineligibility_reason(descriptor) + if reason is not None: + raise ValueError(reason) descriptor_state = descriptor.stddev - if descriptor_state.dtype != torch.float32: - raise ValueError( - "DPA4C compressed CUDA requires descriptor precision `float32`, " - f"got {descriptor_state.dtype}" - ) - if not mega_eligible(descriptor): - raise ValueError( - "DPA4C compressed CUDA supports " - f"channels {SUPPORTED_CHANNELS}, lmax {SUPPORTED_LMAX} and " - f"radial_modes {SUPPORTED_RADIAL_MODES}, got " - f"channels={descriptor.channels}, lmax={descriptor.lmax}, " - f"radial_modes={descriptor.radial_modes}" - ) device = descriptor_state.device profile = descriptor_profile( descriptor.channels, diff --git a/deepmd/pt_expt/utils/auto_neighbor_list.py b/deepmd/pt_expt/utils/auto_neighbor_list.py new file mode 100644 index 0000000000..fcc5e17314 --- /dev/null +++ b/deepmd/pt_expt/utils/auto_neighbor_list.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Batch-aware selection of persistent torch neighbor-list strategies.""" + +from typing import ( + Any, + Literal, +) + +import torch + +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) +from deepmd.dpmodel.utils.neighbor_list import ( + EdgeNeighborList, + NeighborList, +) +from deepmd.pt.utils.nv_nlist import ( + NvNeighborList, + is_nv_available, +) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + VesinNeighborList, + is_vesin_torch_available, +) + + +class AdaptiveNeighborList(NeighborList): + """Select a cell-list backend from the device and batch shape. + + ``auto`` favors persistent Vesin searches for single frames and CPU + inference, and Toolkit-Ops for CUDA batches. ``native`` retains the + unbuffered batched policy: Vesin for one CPU frame, Toolkit-Ops otherwise, + with either backend serving as the fallback when the other is unavailable. + The selection is independent of the requested neighbor representation. + """ + + def __init__( + self, skin: float = 0.0, *, method: Literal["auto", "native"] = "auto" + ) -> None: + self._method = method + self._vesin = VesinNeighborList(skin=0.0 if method == "native" else skin) + self._nv = NvNeighborList() + + def build( + self, + coord: Any, + atype: Any, + box: Any, + rcut: float, + sel: list[int], + return_mode: str = "extended", + pair_excl: PairExcludeMask | None = None, + ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: + """Build the requested schema with the available inference strategy.""" + vesin_available = is_vesin_torch_available() + single_vesin = atype.shape[0] == 1 and ( + self._method == "auto" or coord.device.type == "cpu" + ) + batched_nv = self._method == "native" or coord.device.type == "cuda" + if single_vesin and vesin_available: + builder = self._vesin + elif batched_nv and is_nv_available(): + builder = self._nv + elif vesin_available: + builder = self._vesin + elif is_nv_available(): + builder = self._nv + else: + raise RuntimeError( + "Neighbor-list construction requires a neighbor-list backend. " + "Install one of:\n" + " pip install 'vesin[torch]'\n" + " pip install nvalchemi-toolkit-ops\n" + "Vesin also ships with the PyTorch extra: " + "pip install 'deepmd-kit[torch]'." + ) + return builder.build( + coord, atype, box, rcut, sel, return_mode=return_mode, pair_excl=pair_excl + ) + + +def make_neighbor_list( + method: str, + device: torch.device, + skin: float = 0.0, + *, + return_mode: Literal["extended", "edges"] = "extended", +) -> NeighborList | None: + """Select a search strategy compatible with the lower-input contract. + + ``None`` delegates extended-coordinate construction to the caller's native + builder. Edge inputs require an explicit carry-all builder: falling back to + a fixed-capacity neighbor list would discard neighbors before conversion. + """ + if method == "native": + if return_mode == "edges": + return AdaptiveNeighborList(method="native") + return None + if method == "auto": + if ( + return_mode == "edges" + or is_vesin_torch_available() + or (device.type == "cuda" and is_nv_available()) + ): + return AdaptiveNeighborList(skin=skin) + return None + if method == "vesin": + if not is_vesin_torch_available(): + raise ImportError( + "Vesin neighbor search requires vesin[torch]>=0.6; " + "install it or select the native neighbor builder." + ) + return VesinNeighborList(skin=skin) + if method == "nv": + if device.type != "cuda": + raise ValueError("NV neighbor search requires CUDA inference tensors") + if not is_nv_available(): + raise ImportError("NV neighbor search requires nvalchemi-toolkit-ops") + return NvNeighborList() + raise ValueError( + f"Unknown neighbor backend {method!r}; expected 'auto', 'native', 'vesin', or 'nv'." + ) diff --git a/deepmd/pt_expt/utils/canonical_graph.py b/deepmd/pt_expt/utils/canonical_graph.py index e01163a073..6495f97005 100644 --- a/deepmd/pt_expt/utils/canonical_graph.py +++ b/deepmd/pt_expt/utils/canonical_graph.py @@ -202,9 +202,74 @@ def canonical_graph_from_neighbor_graph( return result +def canonical_graph_from_pairs( + center: torch.Tensor, + source: torch.Tensor, + edge_vec: torch.Tensor, + n_node: torch.Tensor, + node_count: int, +) -> CanonicalGraph: + """Build a local canonical graph from active, unpadded neighbor pairs. + + Parameters + ---------- + center, source : torch.Tensor + Valid global node indices with shape (E,), int64, on one device. + edge_vec : torch.Tensor + Neighbor-minus-center vectors in angstrom with shape (E, 3), on the + index device. Geometry is converted to float32 for deployment. + n_node : torch.Tensor + Per-frame local node counts with shape (nf,), int64, on the same device. + node_count : int + Total number of local nodes across frames, equal to ``n_node.sum()``. + + Returns + ------- + CanonicalGraph + Destination-major graph with dual CSR and at least two storage slots. + All nodes are local; no coordinate differentiation is required. + """ + edge_count = source.shape[0] + if max(edge_count, node_count) > UINT32_MAX: + raise ValueError("canonical graph node or edge count exceeds the uint32 range") + destination_order = torch.argsort(center, stable=True) + center = center.index_select(0, destination_order) + source = source.index_select(0, destination_order) + edge_vec = edge_vec.index_select(0, destination_order).to(torch.float32) + boundaries = torch.arange(node_count + 1, dtype=torch.int64, device=source.device) + destination_row_ptr = torch.searchsorted(center, boundaries) + source_order = torch.argsort(source, stable=True) + source_row_ptr = torch.searchsorted( + source.index_select(0, source_order), boundaries + ) + # Guard slots keep the exported edge axis dynamic even for zero or one + # physical edge. CSR offsets exclude these slots from every node row. + if edge_count < 2: + padding = 2 - edge_count + source = torch.cat((source, source.new_zeros(padding))) + edge_vec = torch.cat((edge_vec, edge_vec.new_zeros((padding, 3)))) + source_order = torch.cat( + ( + source_order, + torch.arange(edge_count, 2, dtype=torch.int64, device=source.device), + ) + ) + n_node = n_node.contiguous() + return CanonicalGraph( + n_node=n_node, + n_local=n_node, + source=source.to(torch.uint32), + edge_vec=edge_vec, + destination_row_ptr=destination_row_ptr, + source_row_ptr=source_row_ptr, + source_order=source_order.to(torch.uint32), + ) + + __all__ = [ "UINT32_MAX", "CanonicalGraph", "canonical_graph_from_neighbor_graph", + "canonical_graph_from_pairs", "validate_canonical_graph_shapes", ] diff --git a/deepmd/pt_expt/utils/edge_schema.py b/deepmd/pt_expt/utils/edge_schema.py index b2cbad736d..9e8a3b2ab8 100644 --- a/deepmd/pt_expt/utils/edge_schema.py +++ b/deepmd/pt_expt/utils/edge_schema.py @@ -352,6 +352,8 @@ def merge_frame_edge_schemas(frames: list[EdgeNeighborList]) -> EdgeNeighborList """Merge per-frame local edge schemas into a batched flattened schema.""" if not frames: raise ValueError("at least one frame schema is required") + if len(frames) == 1: + return frames[0] coord = torch.cat([frame.coord for frame in frames], dim=0) atype = torch.cat([frame.atype for frame in frames], dim=0) nloc = coord.shape[1] diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index 3b117732a6..b19b1f6d35 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -34,18 +34,17 @@ def resolve_auto_graph_builder( ) -> str: """Resolve ``neighbor_graph_method="auto"`` to a concrete inference builder. - Single owner of the inference / DeepEval auto ladder. Training uses + Single owner of graph inference / DeepEval auto selection. Training uses :func:`resolve_neighbor_graph_method`, which never selects ``vesin``. - Mirrors :func:`deepmd.pt.model.model.sezm_model._select_neighbor_builder`: - ``vesin`` is eligible only for a single-frame batch (``nf == 1``), because - its API loops frames in Python (~1 ms/frame). Multi-frame batches stay on - ``nv`` (CUDA) or ``dense`` so ``auto_batch_size`` / ``dp test`` do not - regress to the per-frame loop. + ``vesin`` is preferred for single-frame CUDA inference. Its API loops over + frames in Python, so multi-frame CUDA batches prefer ``nv`` and fall back + to ``dense`` when Toolkit-Ops is unavailable. Notes ----- - * CUDA + ``nvalchemiops``: ``nv`` (any ``nf``). + * CUDA + ``nf == 1`` + ``vesin.torch``: ``vesin``. + * CUDA + ``nvalchemiops``: ``nv``. * CPU + the operator library: ``cell``, whose search is threaded. * ``nf == 1`` + ``vesin.torch``: ``vesin``. * otherwise: ``dense``. @@ -92,6 +91,8 @@ def resolve_auto_graph_builder( raise ValueError(f"nf must be >= 1, got {nf}") dev = torch.device(device) + if dev.type == "cuda" and nf == 1 and is_vesin_torch_available(): + return "vesin" nv_available = is_nv_available() if dev.type == "cuda" and nv_available: return "nv" diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py index aacd1383bc..663cff1aaf 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -359,5 +359,12 @@ def build_neighbor_graph_nv( if pair_excl is not None: graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) if with_csr or canonicalize: - graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize) + # Matrix slots are decoded in ascending center order; compact filtering + # preserves that order, whereas an exclusion mask can introduce holes. + graph = attach_edge_csr( + graph, + nf * nloc, + canonicalize=canonicalize, + destination_sorted=pair_excl is None or compact, + ) return graph diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 1bec3c9b31..8104cce43b 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1334,7 +1334,10 @@ def _serialize_from_file_pte(model_file: str) -> dict: "model_def_script.json": "", "metadata.json": "", } - torch.export.load(model_file, extra_files=extra_files) + # Storage reconstruction requires CPU staging independently of the caller's + # default device; tensor placement is restored from the archive metadata. + with torch.device("cpu"): + torch.export.load(model_file, extra_files=extra_files) model_dict = json.loads(extra_files["model.json"]) model_dict = _json_to_numpy(model_dict) metadata = ( @@ -1661,6 +1664,8 @@ def deserialize_to_file( model_json_override: dict | None = None, do_atomic_virial: bool = False, lower_kind: str = "nlist", + *, + with_comm: bool = True, ) -> None: """Deserialize a dictionary to a .pte or .pt2 model file. @@ -1704,6 +1709,10 @@ def deserialize_to_file( explicitly selects other levels. The selected schema is recorded as ``lower_input_kind`` in ``metadata.json``. + with_comm : bool, optional + Include a communication artifact in ``.pt2`` exports for models with + cross-rank message passing, by default True. Disable for + single-process inference. """ lower_kind = _resolve_target_lower_kind(model_file, data, lower_kind) uses_dpa4_defaults = model_file.endswith(".pt2") and _uses_dpa4_kernel_defaults( @@ -1731,6 +1740,7 @@ def deserialize_to_file( model_json_override, do_atomic_virial, lower_kind, + with_comm=with_comm, ) else: _deserialize_to_file_pte( @@ -2668,6 +2678,8 @@ def _deserialize_to_file_pt2( model_json_override: dict | None = None, do_atomic_virial: bool = False, lower_kind: str = "nlist", + *, + with_comm: bool = True, ) -> None: """Deserialize a dictionary to a .pt2 model file (AOTInductor). @@ -2678,11 +2690,12 @@ def _deserialize_to_file_pt2( For models whose descriptor reports ``has_message_passing_across_ranks() == True`` (DPA2, DPA3 with ``use_loc_mapping=False``, or hybrids wrapping such children), - compiles a SECOND ``with-comm`` artifact and packs it alongside the - regular one. The ``with-comm`` variant accepts comm-dict tensors as - additional positional inputs and drives MPI ghost-atom exchange via - ``deepmd_export::border_op``. The C++ ``DeepPotPTExpt`` loader picks - the artifact based on the LAMMPS rank count at runtime. + compiles a SECOND ``with-comm`` artifact when ``with_comm`` is True and + packs it alongside the regular one. The ``with-comm`` variant accepts + comm-dict tensors as additional positional inputs and drives MPI + ghost-atom exchange via ``deepmd_export::border_op``. The C++ + ``DeepPotPTExpt`` loader picks the artifact based on the LAMMPS rank + count at runtime. Layout inside the .pt2 ZIP (PyTorch 2.11 strict layout): regular → artifact at ``model/`` (AOTInductor's own layout) @@ -2711,6 +2724,9 @@ def _deserialize_to_file_pt2( lower_kind=lower_kind, ) metadata["output_keys"] = output_keys + metadata["has_comm_artifact"] = with_comm and bool( + metadata.get("has_comm_artifact") + ) # A charge-state update swaps constants in one compiled artifact. A model # that also needs the with-comm lower would own two independent constant diff --git a/deepmd/pt_expt/utils/tabulate_ops.py b/deepmd/pt_expt/utils/tabulate_ops.py index 3e2f3db13b..5fedf66910 100644 --- a/deepmd/pt_expt/utils/tabulate_ops.py +++ b/deepmd/pt_expt/utils/tabulate_ops.py @@ -50,6 +50,12 @@ def _try_register_fake(op_name: str, fn: Callable[..., Any]) -> None: """Register a fake implementation, silently skipping if already registered.""" if op_name in _registered: return + # CompositeImplicitAutograd operators use their decomposition on fake + # tensors; PyTorch forbids a separate Meta registration for such operators. + if torch._C._dispatch_has_kernel_for_dispatch_key( + op_name, "CompositeImplicitAutograd" + ): + return try: torch.library.register_fake(op_name)(fn) _registered.add(op_name) diff --git a/deepmd/pt_expt/utils/vesin_graph_builder.py b/deepmd/pt_expt/utils/vesin_graph_builder.py index cc999bd1af..9e9b73defb 100644 --- a/deepmd/pt_expt/utils/vesin_graph_builder.py +++ b/deepmd/pt_expt/utils/vesin_graph_builder.py @@ -1,18 +1,16 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Carry-all NeighborGraph builder backed by vesin.torch (O(N) cell list). - -World-2 counterpart of vesin_neighbor_list.py: instead of building the dense -quartet, it returns per-frame local (i, j, S), then delegates to the array-API -``neighbor_graph_from_ijs`` (which recomputes ``edge_vec`` differentiably from -the ORIGINAL grad-carrying coords). torch-only => lives in pt_expt. - -Scope note: ``vesin.torch``'s API is single-system, so this builder LOOPS over -frames in Python (~1 ms/frame call overhead measured on GPU). It is intended -for ``nf == 1`` inference and CPU use. Inference ``neighbor_graph_method="auto"`` -(:func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`) selects -vesin only when ``nf == 1`` and ``vesin.torch`` is importable (and ``nv`` is -unavailable on CUDA); multi-frame batches stay on ``nv``/``dense``. Training -auto never selects vesin. Prefer ``nv`` (:mod:`.nv_graph_builder`) for batched +"""Vesin graph builders for differentiable and compact deployment inputs. + +Generic graphs recompute displacements from local neighbor pairs and periodic +image shifts, preserving coordinate and cell derivatives. Compact canonical +graphs use Vesin's detached displacements for exported models with explicit +forces. Both layouts retain every active neighbor within the exact cutoff. + +Vesin's API is single-system, so this builder loops over frames in Python. +It is intended for ``nf == 1`` inference and CPU use. Inference ``neighbor_graph_method="auto"`` +(:func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`) prefers +vesin for single-frame CUDA inputs; multi-frame batches stay on ``nv``/``dense``. +Training auto never selects vesin. Prefer ``nv`` (:mod:`.nv_graph_builder`) for batched multi-frame GPU work, which batches all frames in one kernel. """ @@ -40,57 +38,16 @@ attach_edge_csr, neighbor_graph_from_ijs, ) +from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + canonical_graph_from_pairs, +) from deepmd.pt_expt.utils.vesin_neighbor_list import ( is_vesin_torch_available, ) - - -def vesin_search_ijs( - positions: torch.Tensor, - cell: torch.Tensor | None, - periodic: bool, - rcut: float, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Raw ``vesin.torch`` neighbor search returning ``(ii, jj, ss)`` as int64. - - The caller is responsible for ensuring ``vesin.torch`` is importable (check - :func:`~deepmd.pt_expt.utils.vesin_neighbor_list.is_vesin_torch_available` - before calling). ``positions`` must be detached (the search is - non-differentiable). - - Parameters - ---------- - positions : (nloc, 3) local-frame coordinates, already detached. - cell : (3, 3) box matrix for periodic systems, or ``None`` for non-periodic. - For non-periodic systems a zero box is constructed internally. - periodic : whether the system is periodic. - rcut : neighbor cutoff radius. - device : device to pin as the ambient default. ``vesin.torch`` allocates - some internal tensors on the ambient default device, which may be a - fake/other device in some test contexts (e.g. a placeholder CUDA - default); pinning it here prevents spurious CUDA initializations. - - Returns - ------- - ii : (E,) int64 center local indices. - jj : (E,) int64 neighbor local indices. - ss : (E, 3) int64 periodic image shifts. - """ - import vesin.torch as _vesin_torch - - box = ( - cell if periodic else torch.zeros((3, 3), dtype=positions.dtype, device=device) - ) - nl = _vesin_torch.NeighborList(cutoff=float(rcut), full_list=True) - with torch.device(device): - ii, jj, ss = nl.compute( - points=positions, - box=box, - periodic=periodic, - quantities="ijS", - ) - return ii.to(torch.int64), jj.to(torch.int64), ss.to(torch.int64).reshape(-1, 3) +from deepmd.pt_expt.utils.vesin_search import ( + VesinNeighborSearch, +) def build_neighbor_graph_vesin( @@ -104,6 +61,7 @@ def build_neighbor_graph_vesin( canonicalize: bool = False, pair_excl: PairExcludeMask | None = None, compact: bool = False, + search: VesinNeighborSearch | None = None, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using vesin.torch's O(N) cell list. @@ -135,6 +93,9 @@ def build_neighbor_graph_vesin( compact Passed to :func:`apply_pair_exclusion`; see that function for details. Ignored when ``pair_excl`` is ``None``. + search + Optional persistent search shared across evaluations. Its skin only + controls candidate reuse; the returned graph retains the exact cutoff. Returns ------- @@ -173,45 +134,44 @@ def build_neighbor_graph_vesin( canonicalize=canonicalize, ) + if search is None: + search = VesinNeighborSearch( + sorted=(with_csr or canonicalize) and torch.device(dev).type != "cuda" + ) i_parts, j_parts, S_parts, nf_parts = [], [], [], [] for f in range(nf): - pts = coord[f].detach() - cell_f = box[f].detach() if periodic else None - ii, jj, ss = vesin_search_ijs(pts, cell_f, periodic, rcut, dev) + cell_f = box[f] if periodic else None + ii, jj, ss = search.compute(coord[f], cell_f, rcut, frame=f) i_parts.append(ii) j_parts.append(jj) S_parts.append(ss) nf_parts.append(torch.full((ii.shape[0],), f, dtype=torch.int64, device=dev)) - # guard torch.cat against empty part lists (nf == 0), mirroring ase_builder - i_all = ( - torch.cat(i_parts) - if i_parts - else torch.zeros((0,), dtype=torch.int64, device=dev) - ) - j_all = ( - torch.cat(j_parts) - if j_parts - else torch.zeros((0,), dtype=torch.int64, device=dev) - ) - S_all = ( - torch.cat(S_parts) - if S_parts - else torch.zeros((0, 3), dtype=torch.int64, device=dev) - ) - nf_all = ( - torch.cat(nf_parts) - if nf_parts - else torch.zeros((0,), dtype=torch.int64, device=dev) - ) + # A single frame already owns contiguous edge arrays. Empty batches retain + # the same empty-array contract as the backend-independent ASE builder. + if nf == 1: + i_all, j_all, S_all, nf_all = i_parts[0], j_parts[0], S_parts[0], nf_parts[0] + elif nf > 1: + i_all = torch.cat(i_parts) + j_all = torch.cat(j_parts) + S_all = torch.cat(S_parts) + nf_all = torch.cat(nf_parts) + else: + i_all = torch.zeros((0,), dtype=torch.int64, device=dev) + j_all = torch.zeros((0,), dtype=torch.int64, device=dev) + S_all = torch.zeros((0, 3), dtype=torch.int64, device=dev) + nf_all = torch.zeros((0,), dtype=torch.int64, device=dev) # virtual atoms (atype < 0) are excluded as centers AND neighbors — the # World-2 builder contract shared with the dense reference builder; the # geometric search above cannot know about them. at = torch.as_tensor(atype, device=dev).reshape(nf, nloc) keep = (at[nf_all, i_all] >= 0) & (at[nf_all, j_all] >= 0) - i_all, j_all, nf_all = i_all[keep], j_all[keep], nf_all[keep] - S_all = S_all[keep] + indices = torch.nonzero(keep, as_tuple=False).flatten() + i_all = i_all.index_select(0, indices) + j_all = j_all.index_select(0, indices) + nf_all = nf_all.index_select(0, indices) + S_all = S_all.index_select(0, indices) # i = center (dst), j = neighbor (src); pass ORIGINAL coord/box # (grad-carrying). Unlike the nv builder, vesin's cell list handles @@ -231,5 +191,81 @@ def build_neighbor_graph_vesin( at_flat = torch.as_tensor(atype, device=dev).reshape(-1) graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact) if with_csr or canonicalize: - graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize) + graph = attach_edge_csr( + graph, + nf * nloc, + canonicalize=canonicalize, + destination_sorted=search.sorted and (pair_excl is None or compact), + ) return graph + + +def build_canonical_graph_vesin( + coord: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + search: VesinNeighborSearch, + node_mask: torch.Tensor | None = None, +) -> CanonicalGraph: + """Build the compact deployment graph directly from Vesin displacements. + + Parameters + ---------- + coord : torch.Tensor + Coordinates in angstrom with shape (nf, nloc, 3). + box : torch.Tensor or None + Cells in angstrom with shape (nf, 3, 3), or non-periodic boundaries. + rcut : float + Neighbor cutoff radius in angstrom. + search : VesinNeighborSearch + Persistent search with independent candidate storage for each frame. + node_mask : torch.Tensor or None + Active-node mask with shape (nf, nloc), on the coordinate device. + None includes all nodes. Inactive nodes remain in the node axis but + have no incoming or outgoing edges. + + Returns + ------- + CanonicalGraph + Local graph with exact-cutoff float32 geometry and dual CSR. Geometry + is detached; this builder serves exported models with explicit forces. + """ + nf, nloc, _ = coord.shape + centers, sources, vectors = [], [], [] + for frame in range(nf): + if nloc == 0: + break + center, source, edge_vec = search.compute( + coord[frame], + box[frame] if box is not None else None, + rcut, + frame=frame, + return_vectors=True, + ) + if node_mask is not None: + active = node_mask[frame, center] & node_mask[frame, source] + indices = torch.nonzero(active, as_tuple=False).flatten() + center = center.index_select(0, indices) + source = source.index_select(0, indices) + edge_vec = edge_vec.index_select(0, indices) + if frame: + center = center + frame * nloc + source = source + frame * nloc + centers.append(center) + sources.append(source) + vectors.append(edge_vec) + if not centers: + center = torch.empty(0, dtype=torch.int64, device=coord.device) + source = torch.empty_like(center) + edge_vec = coord.new_empty((0, 3)) + elif nf > 1: + center = torch.cat(centers) + source = torch.cat(sources) + edge_vec = torch.cat(vectors) + return canonical_graph_from_pairs( + center, + source, + edge_vec, + torch.full((nf,), nloc, dtype=torch.int64, device=coord.device), + nf * nloc, + ) diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 7453176d32..a2d67b5754 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -36,6 +36,9 @@ edge_schema_from_ij_shifts, merge_frame_edge_schemas, ) +from deepmd.pt_expt.utils.vesin_search import ( + VesinNeighborSearch, +) def is_vesin_torch_available() -> bool: @@ -56,6 +59,9 @@ class VesinNeighborList(NeighborList): namespace and device of the input coordinates. """ + def __init__(self, skin: float = 0.0) -> None: + self._search = VesinNeighborSearch(skin=skin) + def build( self, coord: Any, @@ -119,6 +125,8 @@ def build( atype_t[ff], rcut, sel, + self._search, + ff, ) for ff in range(nframes) ] @@ -143,6 +151,8 @@ def build( atype_t[ff], rcut, sel, + self._search, + ff, ) for ff in range(nframes) ] @@ -193,6 +203,8 @@ def _build_single( atype: torch.Tensor, rcut: float, sel: list[int], + search: VesinNeighborSearch, + frame: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Single-frame ``(i,j,S)`` -> minimal-extended conversion. @@ -220,17 +232,7 @@ def _build_single( cell if periodic else torch.zeros((3, 3), dtype=positions.dtype, device=device) ) - # Delegate the raw search to the shared helper in vesin_graph_builder - # (function-level import: legacy module depends on graph module lazily to - # avoid a module-level cycle — vesin_graph_builder imports - # is_vesin_torch_available from this module). - from deepmd.pt_expt.utils.vesin_graph_builder import ( - vesin_search_ijs, - ) - - ii, jj, ss = vesin_search_ijs( - positions.detach(), cell if periodic else None, periodic, rcut, device - ) + ii, jj, ss = search.compute(positions, cell, rcut, frame=frame) # Phantom atoms (atype < 0) pad a mixed-nloc batch and have no physical # site; the geometric search above cannot know that. Their pairs are # dropped before the ``sum(sel)`` truncation below, so that a phantom can @@ -309,6 +311,8 @@ def _build_single_edges( atype: torch.Tensor, rcut: float, sel: list[int], + search: VesinNeighborSearch, + frame: int, ) -> EdgeNeighborList: """Single-frame ``vesin`` output converted directly to edge vectors.""" device = positions.device @@ -325,14 +329,7 @@ def _build_single_edges( rcut, ) - periodic = cell is not None - from deepmd.pt_expt.utils.vesin_graph_builder import ( - vesin_search_ijs, - ) - - ii, jj, ss = vesin_search_ijs( - positions.detach(), cell if periodic else None, periodic, rcut, device - ) + ii, jj, ss = search.compute(positions, cell, rcut, frame=frame) # ss is int64 from the helper; edge_schema_from_ij_shifts accepts int shifts. return edge_schema_from_ij_shifts( positions=positions, diff --git a/deepmd/pt_expt/utils/vesin_search.py b/deepmd/pt_expt/utils/vesin_search.py new file mode 100644 index 0000000000..04ff5f31dc --- /dev/null +++ b/deepmd/pt_expt/utils/vesin_search.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Persistent geometric neighbor searches shared by inference input schemas.""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + TYPE_CHECKING, +) + +import torch + +if TYPE_CHECKING: + from vesin_torch import ( + NeighborList, + ) + + +class VesinNeighborSearch: + """Reuse Vesin search storage and optional Verlet candidates. + + Each frame owns a calculator. Vesin refreshes displacements and filters + candidates at the requested cutoff on every call, and rebuilds candidates + when motion exceeds half the skin or the cell or point count changes. + Search results are copied so later calls cannot overwrite indices retained + by an earlier autograd graph. + + Parameters + ---------- + skin : float, default: 0.0 + Extra candidate-search radius in angstrom. Zero disables Verlet caching. + sorted : bool, default: False + Whether pairs are returned in ascending center-index order. + """ + + def __init__(self, skin: float = 0.0, *, sorted: bool = False) -> None: + if not math.isfinite(skin) or skin < 0: + raise ValueError( + f"neighbor_skin must be finite and non-negative, got {skin}" + ) + self.skin = float(skin) + self.sorted = sorted + self._key: tuple[float, torch.device] | None = None + self._frames: dict[int, tuple[NeighborList, torch.Tensor | None]] = {} + + def compute( + self, + positions: torch.Tensor, + cell: torch.Tensor | None, + rcut: float, + *, + frame: int = 0, + return_vectors: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return neighbor indices and image shifts or displacement vectors. + + Parameters + ---------- + positions : torch.Tensor + Coordinates in angstrom with shape (N, 3). The search detaches + these coordinates and runs on their device. + cell : torch.Tensor or None + Cell vectors in angstrom with shape (3, 3), with the same dtype + and device as ``positions``. None selects non-periodic boundaries. + rcut : float + Neighbor cutoff radius in angstrom. + frame : int, optional + Non-negative cache slot, by default 0. Each frame in a batch uses + a distinct slot, reused for its successive configurations. + return_vectors : bool, optional + Return neighbor-minus-center vectors in angstrom instead of + lattice shifts. These vectors are detached from the coordinates + and are intended for inference with explicit force outputs. + + Returns + ------- + ii : torch.Tensor + Local center indices with shape (E,), int64. + jj : torch.Tensor + Local neighbor indices with shape (E,), int64. + geometry : torch.Tensor + Integer lattice shifts with shape (E, 3), int64, or displacement + vectors with the dtype of ``positions`` when ``return_vectors`` + is True. All outputs reside on the device of ``positions`` and + retain their values across subsequent searches. + """ + from vesin_torch import ( + NeighborList, + ) + + device = positions.device + key = (float(rcut), device) + if self._key != key: + self._frames.clear() + self._key = key + calculator, previous_cell = self._frames.get(frame, (None, None)) + # Cell deformation can bring uncached periodic images inside the + # cutoff even when Vesin's single-cell displacement bound permits + # reuse. Comparing CUDA cells synchronizes this invalidation decision. + if calculator is None or ( + self.skin > 0 + and ( + (cell is None) != (previous_cell is None) + or ( + cell is not None + and previous_cell is not None + and not torch.equal(cell, previous_cell) + ) + ) + ): + calculator = NeighborList( + cutoff=float(rcut), + full_list=True, + sorted=self.sorted, + skin=self.skin, + ) + # An independent snapshot detects in-place updates to the input. + previous_cell = ( + cell.detach().clone() if self.skin > 0 and cell is not None else None + ) + self._frames[frame] = calculator, previous_cell + box = ( + cell + if cell is not None + else torch.zeros((3, 3), dtype=positions.dtype, device=device) + ) + with torch.device(device): + try: + ii, jj, geometry = calculator.compute( + points=positions.detach(), + box=box.detach(), + periodic=cell is not None, + quantities="ijD" if return_vectors else "ijS", + ) + except RuntimeError as error: + # An empty CUDA Verlet cache can expose a null pair buffer that + # PyTorch cannot wrap. An uncached search supplies valid empty + # tensors; nonempty results and unrelated errors still fail. + if ( + device.type != "cuda" + or self.skin == 0 + or str(error) + != "The specified pointer resides on host memory and is not " + "registered with any CUDA device." + ): + raise + ii, jj, geometry = NeighborList( + cutoff=float(rcut), full_list=True + ).compute( + points=positions.detach(), + box=box.detach(), + periodic=cell is not None, + quantities="ijD" if return_vectors else "ijS", + ) + if ii.numel() != 0: + raise + if not return_vectors: + geometry = geometry.to(torch.int64) + return ii.to(torch.int64), jj.to(torch.int64), geometry.reshape(-1, 3) diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md index 4564c2f460..14134055fb 100644 --- a/doc/model/dpa4c.md +++ b/doc/model/dpa4c.md @@ -332,24 +332,34 @@ now summed with. ## Model compression Compression is the deployment step. It replaces the analytic radial functions -and their type-pair modulation with tabulated splines evaluated by fused CUDA -kernels, and re-exports the model in the compact canonical graph form that the -fast inference path consumes. Because the radial map is analytically bounded and +and their type-pair modulation with tabulated splines evaluated by fused CPU or +CUDA kernels. Eligible CUDA models use the compact canonical graph form; +CPU models use the graph form with their dedicated kernels. Because the radial map is analytically bounded and vanishes at `rcut`, the table needs no extrapolation region and no overflow checking. -Train, freeze, then compress the frozen archive: +Freeze and compress a checkpoint in one operation: ```bash dp --pt-expt train input.json -dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph -dp --pt-expt compress -i frozen_model.pt2 -o compressed_model.pt2 +dp --pt-expt freeze -c model.ckpt.pt -o compressed_model --compress ``` -The two archives are not interchangeable. `frozen_model.pt2` carries the plain -graph lower and is the uncompressed intermediate; `compressed_model.pt2` carries -the compact canonical graph lower and is what you deploy. Compression selects -that lower on its own, so it takes no lower-kind option of its own. +This writes `compressed_model.pt2` directly using the default compression +settings, without compiling an intermediate uncompressed archive. Compression +selects the export schema appropriate to the model and target device. + +To compress an existing frozen archive or choose the spline spacing, use the +compression command: + +```bash +dp --pt-expt compress -i frozen_model.pt2 -o compressed_model.pt2 --step 0.01 +``` + +For ASE, `DP(model="model.pt", freeze=True)` automatically compresses supported +DPA4C checkpoints and exports a single-process `.pt2`. If compression is not +supported, ASE warns and retains the ordinary freeze path. See the +[ASE interface](../third-party/ase.md). Only `-s, --step` applies to DPA4C; it sets the uniform spline spacing in Å, and a smaller value means a finer table and a larger model. The `--extrapolate`, @@ -359,14 +369,15 @@ computed from data; DPA4C needs none of them and ignores them. Compression requires: -- the PyTorch Exportable backend on CUDA; +- the PyTorch Exportable backend; - `precision: "float32"`; - `channels`, `lmax` and `radial_modes` inside the compiled sets listed above; - an empty {ref}`exclude_types `, since the fused kernel has no type-exclusion branch. -`dp --pt-expt compress` reports an explicit error when any of these is not met. +`dp --pt-expt compress` and `freeze --compress` report an explicit error when +any of these is not met. A model that excludes type pairs still trains and runs; deploy it as the uncompressed graph archive. diff --git a/doc/third-party/ase.md b/doc/third-party/ase.md index f4c350fca1..498b41c4f7 100644 --- a/doc/third-party/ase.md +++ b/doc/third-party/ase.md @@ -60,6 +60,74 @@ print(water.get_forces()) :::: +Neighbor construction can be configured independently of checkpoint freezing: + +```python +calculator = DP( + model="model.pt", + neighbor_graph_method="auto", + neighbor_skin=1.0, +) +``` + +For PT-expt graph models, `neighbor_graph_method` accepts `auto`, `dense`, `ase`, +`cell`, `vesin`, and `nv`. On a GPU, `auto` prefers Vesin for the single frame +evaluated by ASE; batched `DeepPot` inference prefers NV when available. CPU +graph inference retains the native `cell` search. PT SeZM/DPA4 checkpoints and +their edge-vector exports support `auto`, `vesin`, and `nv`. Models with the +traditional extended-coordinate input use `nlist_backend` instead. + +The default `neighbor_skin=1.0`, in angstrom, enables candidate reuse for Vesin +searches across all these input formats and requires Vesin and its torch +bindings at version 0.6 or later (PyTorch 2.3 or later). +Older PyTorch targets and macOS Intel installations retain native neighbor +construction without requiring these bindings. +Candidates extend to the model cutoff plus the skin and are rebuilt when an +atom moves more than half the skin, or when the cell or atom count changes. +Each evaluation still updates geometry and selects neighbors at the model's +exact cutoff. Set `neighbor_skin=0.0` to disable candidate reuse. Other search +backends retain their own rebuild behavior. Element-to-model type mappings +are reused until atomic numbers, atom order, or the supplied `type_dict` changes. + +For PT DPA4/SeZM checkpoints and graph-eligible PT-expt checkpoints, enable +automatic freezing for production inference: + +```python +from deepmd.calculator import DP + +calculator = DP(model="model.pt", freeze=True) +``` + +Each calculator construction with `freeze=True` rebuilds the artifact using +the selected checkpoint head and the backend's normal freeze policy. The +calculator loads a private artifact retained for its evaluator's lifetime and +atomically publishes a separate copy as `model-ase_single.pt2` beside the +checkpoint. Concurrent processes retain their own selected heads; the last +publication replaces the shared copy. This requires one extra artifact copy +per live evaluator. Concurrent exports must use separate processes because +the backend compiler configuration is not thread-safe. + +The export omits the multi-rank communication artifact +because ASE inference runs in one process. For CUDA targets, DPA4 defaults to +`DP_CUDA_INFER=1` and `DP_TRITON_INFER=2`, with explicit environment settings +taking precedence; DPA4C keeps its required CUDA level of 2. + +DPA4C checkpoints also enable compression through the same export path as +`dp --pt-expt freeze --compress`. Supported models use the dedicated CPU or +CUDA compression kernels. If the configuration cannot be compressed, such as +a `float64` descriptor or one with type exclusions, ASE logs the reason and +continues with an ordinary frozen model. Other model families keep their +uncompressed freeze behavior. Explicit `freeze --compress` reports an error +for an unsupported compression configuration. + +The default `freeze=False` loads checkpoints directly and logs a warning for +these models: inference can be **substantially slower**. This mode is intended +for debugging and short runs that avoid compilation overhead. To reuse an +export without compiling again, pass the `.pt2` file directly after concurrent +exports have finished. Already frozen models, other PT models, and PT-expt +models without graph eligibility ignore `freeze` and keep their normal loading +path. + Optimization is also available: ```python diff --git a/pyproject.toml b/pyproject.toml index 5da7f9b1c6..774cd08fa5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,6 +79,7 @@ repository = "https://github.com/deepmodeling/deepmd-kit" # which can be read by the build backend. [tool.deepmd_build_backend.optional-dependencies] test = [ + "dependency_groups", "dpdata>=0.2.7", # ASE issue: https://gitlab.com/ase/ase/-/merge_requests/2843 # fixed in 3.23.0 diff --git a/source/tests/common/sezm_config.py b/source/tests/common/sezm_config.py new file mode 100644 index 0000000000..a4d50d1965 --- /dev/null +++ b/source/tests/common/sezm_config.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-independent SeZM configurations for export tests.""" + + +def _tiny_sezm_model_params() -> dict: + """Minimal fp64 SeZM config for self-contained export tests. + + ``precision="float64"`` is what unlocks the ``rtol=1e-10, atol=1e-10`` + parity pt_expt enforces; fp32 accumulation alone drifts in the 1e-6 + range. All other knobs keep ``make_fx`` tracing small. + """ + return { + "type": "SeZM", + "type_map": ["A", "B"], + "descriptor": { + "type": "SeZM", + "sel": [2, 2], + "rcut": 3.0, + "channels": 4, + "n_focus": 1, + "n_radial": 3, + "radial_mlp": [6], + "use_env_seed": True, + "l_schedule": [1, 0], + "mmax": 1, + "so2_norm": False, + "so2_layers": 1, + "n_atten_head": 1, + "sandwich_norm": [True, False, True, False], + "ffn_neurons": 8, + "ffn_blocks": 1, + "s2_activation": [False, True], + "mlp_bias": False, + "layer_scale": False, + "use_amp": False, + "activation_function": "silu", + "glu_activation": True, + "precision": "float64", + "seed": 7, + }, + "fitting_net": { + "neuron": [8], + "activation_function": "silu", + "precision": "float64", + "seed": 7, + }, + "use_compile": False, + } + + +def _tiny_sezm_native_spin_model_params() -> dict: + """Minimal fp64 native-spin SeZM config for freeze routing tests.""" + params = _tiny_sezm_model_params() + params["type_map"] = ["O", "H"] + params["spin"] = { + "use_spin": [True, False], + "scheme": "native", + } + return params diff --git a/source/tests/common/test_pytorch_requirements.py b/source/tests/common/test_pytorch_requirements.py new file mode 100644 index 0000000000..f04b277320 --- /dev/null +++ b/source/tests/common/test_pytorch_requirements.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""PyTorch extras keep Vesin bindings within their supported versions.""" + +from pathlib import ( + Path, +) + +import pytest +from packaging.requirements import ( + Requirement, +) + + +@pytest.mark.parametrize("version", ["2.1.0", "2.2.0", "2.3.0", ""]) +def test_vesin_requirements(version: str, monkeypatch: pytest.MonkeyPatch) -> None: + # The PEP 517 backend belongs to the source checkout, not the installed wheel. + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3])) + from backend.find_pytorch import ( + get_pt_requirement, + ) + + monkeypatch.delenv("PYTORCH_VERSION", raising=False) + monkeypatch.setenv("CIBUILDWHEEL", "0") + get_pt_requirement.cache_clear() + requirements = { + requirement.name: requirement + for requirement in map(Requirement, get_pt_requirement(version)["torch"]) + } + get_pt_requirement.cache_clear() + if version in ("2.1.0", "2.2.0"): + assert "vesin" not in requirements + assert "vesin-torch" not in requirements + return + + assert requirements["vesin"].extras == {"torch"} + for name in ("vesin", "vesin-torch"): + requirement = requirements[name] + assert "0.5.9" not in requirement.specifier + assert "0.6" in requirement.specifier + for system, machine, supported in ( + ("Darwin", "x86_64", False), + ("Darwin", "arm64", True), + ("Linux", "x86_64", True), + ("Windows", "AMD64", True), + ): + assert ( + requirement.marker.evaluate( + {"platform_system": system, "platform_machine": machine} + ) + is supported + ) diff --git a/source/tests/pt/model/test_nlist_backend.py b/source/tests/pt/model/test_nlist_backend.py index d38262dacd..85de3b1b40 100644 --- a/source/tests/pt/model/test_nlist_backend.py +++ b/source/tests/pt/model/test_nlist_backend.py @@ -10,6 +10,12 @@ """ import copy +from pathlib import ( + Path, +) +from unittest import ( + mock, +) import numpy as np import pytest @@ -24,6 +30,9 @@ from deepmd.pt.train.wrapper import ( ModelWrapper, ) +from deepmd.pt.utils.env import ( + DEVICE, +) from deepmd.pt.utils.nv_nlist import ( NvNeighborList, is_nv_available, @@ -33,6 +42,10 @@ is_vesin_torch_available, ) +from ...common.sezm_config import ( + _tiny_sezm_model_params, +) + # Each O(N) strategy: (backend name, builder class, availability skip mark). _BACKEND_MARKS = { "vesin": pytest.mark.skipif( @@ -143,6 +156,31 @@ def test_invalid_backend_raises(pt_files) -> None: DeepPot(pt_files["se_e2_a"], nlist_backend="bogus") +@pytest.mark.parametrize( + ("method", "backend", "message"), + [ + ("dense", "auto", "PT SeZM/DPA4 supports neighbor_graph_method"), + ("vesin", "native", "Select only one"), + ], +) +def test_self_built_neighbor_options_rejected( + tmp_path: Path, method: str, backend: str, message: str +) -> None: + path = str(tmp_path / "sezm.pt") + _save_pt(_tiny_sezm_model_params(), path) + with pytest.raises(ValueError, match=message): + DeepPot(path, neighbor_graph_method=method, nlist_backend=backend) + + +def test_neighbor_graph_method_rejected_for_dense_model( + pt_files: dict[str, str], +) -> None: + with pytest.raises( + ValueError, match="neighbor_graph_method applies to PT SeZM/DPA4" + ): + DeepPot(pt_files["se_e2_a"], neighbor_graph_method="vesin") + + def test_native_uses_no_strategy(pt_files) -> None: dp = DeepPot(pt_files["se_e2_a"], nlist_backend="native") assert dp.deep_eval._nlist_builder is None @@ -155,10 +193,31 @@ def test_explicit_backend_selects_builder(pt_files, backend: str) -> None: @_BACKEND_MARKS["vesin"] -def test_auto_prefers_vesin(pt_files) -> None: - # auto picks the first available O(N) builder; vesin is preferred. - builder = DeepPot(pt_files["se_e2_a"]).deep_eval._nlist_builder - assert isinstance(builder, VesinNeighborList) +@pytest.mark.parametrize("nframes", [1, 3]) +def test_auto_selects_backend_from_frame_count( + pt_files: dict[str, str], nframes: int +) -> None: + coords, atype, boxes = _multiframe_system(nframes) + dp = DeepPot(pt_files["se_e2_a"], auto_batch_size=False) + with ( + mock.patch.object( + VesinNeighborList, + "build", + autospec=True, + side_effect=VesinNeighborList.build, + ) as vesin_build, + mock.patch.object( + NvNeighborList, "build", autospec=True, side_effect=NvNeighborList.build + ) as nv_build, + ): + dp.eval(coords, boxes, atype) + + if DEVICE.type == "cuda" and nframes > 1 and is_nv_available(): + nv_build.assert_called_once() + vesin_build.assert_not_called() + else: + vesin_build.assert_called_once() + nv_build.assert_not_called() def test_self_built_model_forces_native(pt_files, monkeypatch) -> None: diff --git a/source/tests/pt/model/test_sezm_export.py b/source/tests/pt/model/test_sezm_export.py index 04a920d83f..ecd54923b2 100644 --- a/source/tests/pt/model/test_sezm_export.py +++ b/source/tests/pt/model/test_sezm_export.py @@ -46,6 +46,11 @@ ensure_comm_registered, ) +from ...common.sezm_config import ( + _tiny_sezm_model_params, + _tiny_sezm_native_spin_model_params, +) + _COMM_KEYS = ( "send_list", "send_proc", @@ -133,53 +138,6 @@ def _eager_parallel_forward( ) -def _tiny_sezm_model_params() -> dict: - """Minimal fp64 SeZM config for self-contained export tests. - - ``precision="float64"`` is what unlocks the ``rtol=1e-10, atol=1e-10`` - parity pt_expt enforces; fp32 accumulation alone drifts in the 1e-6 - range. All other knobs are tuned to keep ``make_fx`` tracing time - in the low-single-digit seconds. - """ - return { - "type": "SeZM", - "type_map": ["A", "B"], - "descriptor": { - "type": "SeZM", - "sel": [2, 2], - "rcut": 3.0, - "channels": 4, - "n_focus": 1, - "n_radial": 3, - "radial_mlp": [6], - "use_env_seed": True, - "l_schedule": [1, 0], - "mmax": 1, - "so2_norm": False, - "so2_layers": 1, - "n_atten_head": 1, - "sandwich_norm": [True, False, True, False], - "ffn_neurons": 8, - "ffn_blocks": 1, - "s2_activation": [False, True], - "mlp_bias": False, - "layer_scale": False, - "use_amp": False, - "activation_function": "silu", - "glu_activation": True, - "precision": "float64", - "seed": 7, - }, - "fitting_net": { - "neuron": [8], - "activation_function": "silu", - "precision": "float64", - "seed": 7, - }, - "use_compile": False, - } - - def _tiny_sezm_spin_model_params() -> dict: """Minimal fp64 SeZM spin config for freeze routing tests.""" params = copy.deepcopy(_tiny_sezm_model_params()) @@ -192,17 +150,6 @@ def _tiny_sezm_spin_model_params() -> dict: return params -def _tiny_sezm_native_spin_model_params() -> dict: - """Minimal fp64 native-spin SeZM config for freeze routing tests.""" - params = copy.deepcopy(_tiny_sezm_model_params()) - params["type_map"] = ["O", "H"] - params["spin"] = { - "use_spin": [True, False], - "scheme": "native", - } - return params - - def _build_tiny_sezm_model() -> torch.nn.Module: """Fresh tiny SeZM model on CPU, in eval mode.""" model = get_model(_tiny_sezm_model_params()) diff --git a/source/tests/pt/model/test_sezm_vacuum_freeze.py b/source/tests/pt/model/test_sezm_vacuum_freeze.py index dc3a9766ec..bb1bce9490 100644 --- a/source/tests/pt/model/test_sezm_vacuum_freeze.py +++ b/source/tests/pt/model/test_sezm_vacuum_freeze.py @@ -27,12 +27,14 @@ ModelWrapper, ) +from ...common.sezm_config import ( + _tiny_sezm_model_params, +) from .test_sezm_export import ( _CPU, _SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON, _clear_default_device, - _tiny_sezm_model_params, ) BIAS = np.array([[-3.0], [0.5]]) diff --git a/source/tests/pt/test_compile_compat.py b/source/tests/pt/test_compile_compat.py index 12772edcd7..8c300fa11e 100644 --- a/source/tests/pt/test_compile_compat.py +++ b/source/tests/pt/test_compile_compat.py @@ -43,6 +43,7 @@ def test_fusion_size_defaults_to_eight(monkeypatch) -> None: assert inference_options["max_fusion_size"] == 8 assert inference_options["cpp.min_chunk_size"] == 1 assert inference_options["cpp.dynamic_threads"] is True + assert "cpp.max_horizontal_fusion_size" not in inference_options def test_torch_211_inference_keeps_default_cpp_parallelism(monkeypatch) -> None: @@ -53,6 +54,9 @@ def test_torch_211_inference_keeps_default_cpp_parallelism(monkeypatch) -> None: assert "cpp.min_chunk_size" not in inference_options assert "cpp.dynamic_threads" not in inference_options + assert inference_options["cpp.max_horizontal_fusion_size"] == 1 + assert "cpp.simdlen" not in inference_options + assert "cpp.max_horizontal_fusion_size" not in build_inductor_compile_options() def test_fusion_size_environment_is_shared(monkeypatch) -> None: diff --git a/source/tests/pt/test_neighbor_backend_missing.py b/source/tests/pt/test_neighbor_backend_missing.py index dd608f5e51..ac341fc9f3 100644 --- a/source/tests/pt/test_neighbor_backend_missing.py +++ b/source/tests/pt/test_neighbor_backend_missing.py @@ -12,25 +12,35 @@ import torch -from deepmd.pt.model.model import ( - sezm_model, +from deepmd.pt_expt.utils import ( + auto_neighbor_list, ) class TestNeighborBackendMissing(unittest.TestCase): - def _message_without_backends(self) -> str: + def _message_without_backends(self, method: str = "native") -> str: """The error raised when neither backend is importable.""" - nv = sezm_model.is_nv_available - vesin = sezm_model.is_vesin_torch_available - sezm_model.is_nv_available = lambda: False - sezm_model.is_vesin_torch_available = lambda: False + nv = auto_neighbor_list.is_nv_available + vesin = auto_neighbor_list.is_vesin_torch_available + auto_neighbor_list.is_nv_available = lambda: False + auto_neighbor_list.is_vesin_torch_available = lambda: False try: with self.assertRaises(RuntimeError) as caught: - sezm_model._select_neighbor_builder(1, torch.device("cpu")) + builder = auto_neighbor_list.make_neighbor_list( + method, torch.device("cpu"), return_mode="edges" + ) + builder.build( + torch.zeros((1, 1, 3), device="cpu"), + torch.zeros((1, 1), dtype=torch.int64, device="cpu"), + None, + 1.0, + [1], + return_mode="edges", + ) return str(caught.exception) finally: - sezm_model.is_nv_available = nv - sezm_model.is_vesin_torch_available = vesin + auto_neighbor_list.is_nv_available = nv + auto_neighbor_list.is_vesin_torch_available = vesin def test_error_names_installable_packages(self) -> None: """The message must name pip-installable distributions, not imports.""" @@ -42,7 +52,7 @@ def test_error_names_installable_packages(self) -> None: def test_error_points_at_the_extra(self) -> None: """A user who wants the whole descriptor should be told about it.""" - self.assertIn("deepmd-kit[torch]", self._message_without_backends()) + self.assertIn("deepmd-kit[torch]", self._message_without_backends("auto")) if __name__ == "__main__": diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index 57fdd021ce..27b2bf7c2d 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -229,6 +229,23 @@ def test_serialize_returns_model_tree(self) -> None: # single bool and fail with an ambiguous truth-value error. np.testing.assert_equal(data, serialize_from_file(self.tmpfile.name)["model"]) + def test_pte_load_is_independent_of_default_device(self) -> None: + """Archive loading and inference retain the recorded tensor devices.""" + rng = np.random.default_rng(GLOBAL_SEED) + coords = rng.random((1, 5, 3)) * 8.0 + cells = np.eye(3).reshape(1, 9) * 10.0 + atom_types = np.array([0, 1, 0, 1, 0], dtype=np.int32) + expected = self.dp.eval(coords, cells, atom_types, atomic=True) + for default_device in ("meta", "cuda:9999999"): + with self.subTest(default_device=default_device): + with torch.device(default_device): + dp = DeepPot(self.tmpfile.name) + data = serialize_from_file(self.tmpfile.name) + actual = dp.eval(coords, cells, atom_types, atomic=True) + np.testing.assert_equal(data["model"], self.model_data["model"]) + for value, reference in zip(actual, expected, strict=True): + np.testing.assert_allclose(value, reference, rtol=1e-10, atol=1e-10) + def test_eval_consistency(self) -> None: """Test that DeepPot.eval gives same results as direct model forward.""" rng = np.random.default_rng(GLOBAL_SEED) @@ -2342,12 +2359,19 @@ def _vesin_available(self) -> bool: return is_vesin_torch_available() def test_default_is_auto(self) -> None: + from deepmd.pt.utils.nv_nlist import ( + is_nv_available, + ) + dp = DeepPot(self.tmpfile.name) - self.assertEqual(dp.deep_eval._use_vesin, self._vesin_available()) + self.assertEqual( + dp.deep_eval._nlist_builder is not None, + self._vesin_available() or (DEVICE.type == "cuda" and is_nv_available()), + ) def test_native_disables_vesin(self) -> None: dp = DeepPot(self.tmpfile.name, nlist_backend="native") - self.assertFalse(dp.deep_eval._use_vesin) + self.assertIsNone(dp.deep_eval._nlist_builder) def test_invalid_raises(self) -> None: with self.assertRaises(ValueError): @@ -2376,6 +2400,25 @@ def test_vesin_matches_native(self) -> None: a, b, rtol=1e-10, atol=1e-10, err_msg=f"vesin vs native: {name}" ) + @unittest.skipIf(not torch.cuda.is_available(), "CUDA is required") + def test_nv_matches_native_multiframe(self) -> None: + from deepmd.pt.utils.nv_nlist import ( + is_nv_available, + ) + + if not is_nv_available(): + self.skipTest("nvalchemi-toolkit-ops is required") + rng = np.random.default_rng(GLOBAL_SEED + 21) + coords = rng.random((3, 6, 3)) * 8.0 + atom_types = np.arange(6, dtype=np.int32) % self.nt + cells = np.tile(np.eye(3).reshape(1, 9) * 10.0, (3, 1)) + native = DeepPot(self.tmpfile.name, nlist_backend="native") + accelerated = DeepPot(self.tmpfile.name, nlist_backend="nv") + expected = native.eval(coords, cells, atom_types, atomic=True) + actual = accelerated.eval(coords, cells, atom_types, atomic=True) + for value, reference in zip(actual, expected, strict=True): + np.testing.assert_allclose(value, reference, rtol=1e-9, atol=1e-9) + def test_eval_descriptor_vesin_matches_native(self) -> None: """eval_descriptor bypasses forward_common_lower's format_nlist, so the vesin builder must apply the same type-distinguishing as the native 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 29e3cc9172..507f005b19 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 @@ -357,7 +357,7 @@ def test_nlist_controls_remain_available(self) -> None: output_def, nlist_backend="native", ) - with self.assertRaisesRegex(ValueError, "only applies to graph-routed"): + with self.assertRaisesRegex(ValueError, "only applies to graph or edge-vector"): PtExptDeepEval( self.pt_path, output_def, @@ -422,6 +422,92 @@ def test_unsupported_extension_raises(self) -> None: class TestNeighborGraphMethodResolution(unittest.TestCase): """Auto graph-builder selection must cover each host policy explicitly.""" + def test_auto_nlist_strategy_uses_frame_count(self) -> None: + from deepmd.pt_expt.utils import ( + auto_neighbor_list, + ) + + cases = ( + ("cuda", 1, True, True, "vesin"), + ("cuda", 4, True, True, "nv"), + ("cuda", 1, True, False, "nv"), + ("cuda", 4, False, True, "vesin"), + ("cpu", 1, True, True, "vesin"), + ("cuda", 1, False, False, None), + ("cpu", 1, False, True, "vesin"), + ("cpu", 1, True, False, None), + ("cpu", 1, False, False, None), + ) + for device, nf, nv_available, vesin_available, expected in cases: + with self.subTest(device=device, nf=nf, expected=expected): + coord = mock.Mock(device=torch.device(device)) + atype = np.zeros((nf, 2), dtype=np.int64) + with ( + mock.patch.object( + auto_neighbor_list, "is_nv_available", return_value=nv_available + ), + mock.patch.object( + auto_neighbor_list, + "is_vesin_torch_available", + return_value=vesin_available, + ), + ): + builder = auto_neighbor_list.make_neighbor_list( + "auto", coord.device, skin=1.0 + ) + if expected is None: + self.assertIsNone(builder) + continue + with ( + mock.patch.object( + builder._vesin, "build", return_value="vesin" + ), + mock.patch.object(builder._nv, "build", return_value="nv"), + ): + actual = builder.build(coord, atype, None, 3.0, [4]) + self.assertEqual(actual, expected) + + def test_native_edge_strategy_uses_frame_count(self) -> None: + """Native edge inputs use the shared carry-all strategy on every host.""" + from deepmd.pt_expt.utils import ( + auto_neighbor_list, + ) + + for device, nf, expected in ( + ("cpu", 1, "vesin"), + ("cpu", 4, "nv"), + ("cuda", 1, "nv"), + ("cuda", 4, "nv"), + ): + with ( + self.subTest(device=device, nf=nf), + mock.patch.object( + auto_neighbor_list, "is_nv_available", return_value=True + ), + mock.patch.object( + auto_neighbor_list, "is_vesin_torch_available", return_value=True + ), + mock.patch.object( + auto_neighbor_list.VesinNeighborList, "build", return_value="vesin" + ) as vesin_build, + mock.patch.object( + auto_neighbor_list.NvNeighborList, "build", return_value="nv" + ) as nv_build, + ): + coord = mock.Mock(device=torch.device(device)) + atype = np.zeros((nf, 2), dtype=np.int64) + builder = auto_neighbor_list.make_neighbor_list( + "native", coord.device, return_mode="edges" + ) + actual = builder.build( + coord, atype, None, 3.0, [1], return_mode="edges" + ) + selected = vesin_build if expected == "vesin" else nv_build + selected.assert_called_once_with( + coord, atype, None, 3.0, [1], return_mode="edges", pair_excl=None + ) + self.assertEqual(actual, expected) + def test_explicit_resolution(self) -> None: """Every implemented graph builder is accepted without rewriting.""" for method in ("dense", "ase", "cell", "vesin", "nv"): @@ -450,7 +536,8 @@ def test_auto_resolution(self) -> None: ("cpu", False, False, True, 1, "vesin", False), ("cpu", False, False, True, 4, "dense", False), ("cpu", False, False, False, 1, "dense", False), - ("cuda", True, True, True, 1, "nv", False), + ("cuda", True, True, True, 1, "vesin", False), + ("cuda", True, True, False, 1, "nv", False), ("cuda", True, True, True, 4, "nv", False), ("cuda", False, True, True, 1, "vesin", False), ("cuda", False, True, True, 4, "dense", True), diff --git a/source/tests/pt_expt/model/test_graph_builder_dispatch.py b/source/tests/pt_expt/model/test_graph_builder_dispatch.py index a6185aa6f9..8880282820 100644 --- a/source/tests/pt_expt/model/test_graph_builder_dispatch.py +++ b/source/tests/pt_expt/model/test_graph_builder_dispatch.py @@ -152,7 +152,7 @@ def test_explicit_nv_rejects_cpu(): ("cpu", False, False, True, 1, "vesin"), ("cpu", False, False, True, 4, "dense"), ("cpu", True, False, False, 1, "dense"), - ("cuda", True, True, True, 1, "nv"), + ("cuda", True, True, True, 1, "vesin"), ("cuda", True, True, True, 4, "nv"), ("cuda", False, True, True, 1, "vesin"), ("cuda", False, True, True, 4, "dense"), diff --git a/source/tests/pt_expt/test_calculator.py b/source/tests/pt_expt/test_calculator.py new file mode 100644 index 0000000000..1c3da6a6ad --- /dev/null +++ b/source/tests/pt_expt/test_calculator.py @@ -0,0 +1,591 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""ASE checkpoint export and eager-inference behavior.""" + +import contextlib +import copy +import gc +import json +import logging +import os +import zipfile +from collections.abc import ( + Callable, + Iterator, +) +from concurrent.futures import ( + ThreadPoolExecutor, +) +from functools import ( + partial, +) +from pathlib import ( + Path, +) +from threading import ( + Barrier, +) +from unittest import ( + mock, +) + +import numpy as np +import pytest +import torch +from ase import ( + Atoms, +) + +from deepmd import calculator as calculator_module +from deepmd.calculator import ( + DP, +) +from deepmd.infer import ( + DeepPot, +) +from deepmd.pt.entrypoints import ( + freeze_pt2, +) +from deepmd.pt.utils import env as pt_env +from deepmd.pt_expt.utils import ( + env, + serialization, +) + +from ..common.sezm_config import ( + _tiny_sezm_model_params, + _tiny_sezm_native_spin_model_params, +) +from .model.test_dpa4_export import ( + _DPA4_CONFIG, +) +from .model.test_dpa4c_graph_lower import _config as dpa4c_config +from .model.test_export_with_comm import ( + _DPA3_CONFIG, +) +from .model.test_graph_export_with_comm import ( + DPA1_CONFIG, + DPA2_CONFIG, +) +from .test_dp_freeze import ( + model_se_e2_a, +) + +_LEVELS = ("DP_CUDA_INFER", "DP_TRITON_INFER", "DP_CUTILE_INFER", "DP_CUTE_INFER") + + +@pytest.fixture(autouse=True, params=["cpu"]) +def inference_device( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch +) -> str: + device = request.param + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is required") + for module in (pt_env, env, freeze_pt2): + monkeypatch.setattr(module, "DEVICE", torch.device(device)) + for name in (*_LEVELS, "DP_COMPILE_INFER"): + monkeypatch.delenv(name, raising=False) + return device + + +def _config(kind: str) -> dict: + if kind in ("pt", "pt-dpa4"): + config = _tiny_sezm_model_params() + if kind == "pt-dpa4": + config["type"] = config["descriptor"]["type"] = "dpa4" + elif kind == "dpa4": + config = copy.deepcopy(_DPA4_CONFIG) + config["descriptor"].update(channels=4, n_radial=3, lmax=1, n_blocks=1) + config["fitting_net"]["neuron"] = [8] + elif kind == "dpa4c": + config = dpa4c_config() + config["descriptor"].update(channels=8, lmax=2, precision="float32") + config["fitting_net"]["neuron"] = [8] + config["fitting_net"]["precision"] = "float32" + else: + config = copy.deepcopy( + { + "dpa1": DPA1_CONFIG, + "dpa2": DPA2_CONFIG, + "dpa3": _DPA3_CONFIG, + "dense": model_se_e2_a, + "pt-dense": model_se_e2_a, + }[kind] + ) + config["type_map"] = ["O", "H"] + if kind in ("dense", "pt-dense"): + config["descriptor"]["sel"] = config["descriptor"]["sel"][:2] + if env.DEVICE.type == "cuda": + config["descriptor"]["precision"] = "float32" + config["fitting_net"]["precision"] = "float32" + if kind in ("pt", "pt-dpa4", "dpa4"): + config["descriptor"]["channels"] = 32 + return config + + +def _write_checkpoint( + path: Path, + kind: str, + head: str | None = None, + *, + config: dict | None = None, +) -> Path: + if kind.startswith("pt"): + from deepmd.pt.model.model import ( + get_model, + ) + from deepmd.pt.train.wrapper import ( + ModelWrapper, + ) + else: + from deepmd.pt_expt.model.get_model import ( + get_model, + ) + from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, + ) + + params = _config(kind) if config is None else copy.deepcopy(config) + model = get_model(copy.deepcopy(params)).to("cpu").eval() + generator = torch.Generator(device="cpu").manual_seed(17) + with torch.no_grad(): + for parameter in model.parameters(): + if torch.count_nonzero(parameter) == 0: + parameter.copy_( + 0.05 + * torch.randn( + parameter.shape, + generator=generator, + dtype=parameter.dtype, + device="cpu", + ) + ) + wrapper = ModelWrapper( + model if head is None else {head: model}, + model_params=params if head is None else {"model_dict": {head: params}}, + ) + torch.save({"model": wrapper.state_dict()}, path) + return path + + +@pytest.fixture +def file_freezer(monkeypatch: pytest.MonkeyPatch) -> mock.Mock: + """Exercise artifact ownership without running threaded compilation.""" + + class FileEvaluator: + def __init__(self, model_file: str) -> None: + self.model_path = model_file + + def serialize(self) -> str: + return Path(self.model_path).read_text() + + def load(model_file: str, **kwargs: object) -> mock.Mock: + return mock.Mock( + deep_eval=FileEvaluator(model_file), + get_type_map=mock.Mock(return_value=["O", "H"]), + get_ntypes=mock.Mock(return_value=2), + ) + + def export(source: str, output: str, *, head: str | None) -> None: + Path(output).write_text(head or "Default") + + freezer = mock.Mock(side_effect=export) + + def select( + evaluator: object, head: str | None, **kwargs: object + ) -> Callable[..., None]: + return partial(freezer, head=head) + + monkeypatch.setattr(calculator_module, "DeepPot", load) + monkeypatch.setattr(calculator_module, "_checkpoint_freezer", select) + return freezer + + +def test_automatic_freeze_retains_owned_artifact( + tmp_path: Path, file_freezer: mock.Mock +) -> None: + checkpoint = tmp_path / "model.pt" + first = DP(checkpoint, head="first", freeze=True) + second = DP(checkpoint, head="second", freeze=True) + evaluator = first.dp.deep_eval + private = Path(evaluator.model_path) + assert evaluator.serialize() == "first" + assert second.dp.deep_eval.serialize() == "second" + assert (tmp_path / "model-ase_single.pt2").read_text() == "second" + (tmp_path / "model-ase_single.pt2").write_text("external overwrite") + assert evaluator.serialize() == "first" + assert second.dp.deep_eval.serialize() == "second" + del first + gc.collect() + assert private.exists() + del evaluator + gc.collect() + assert not private.parent.exists() + + +def test_concurrent_automatic_freeze_loads_each_head( + tmp_path: Path, file_freezer: mock.Mock +) -> None: + barrier = Barrier(2) + + def export(source: str, output: str, *, head: str) -> None: + Path(output).write_text(head) + barrier.wait(timeout=10) + + file_freezer.side_effect = export + + def construct(head: str) -> DP: + return DP(tmp_path / "model.pt", head=head, freeze=True) + + with ThreadPoolExecutor(max_workers=2) as pool: + first, second = pool.map(construct, ["first", "second"]) + assert first.dp.deep_eval.serialize() == "first" + assert second.dp.deep_eval.serialize() == "second" + assert (tmp_path / "model-ase_single.pt2").read_text() in {"first", "second"} + + +@pytest.mark.parametrize("stage", ["export", "load", "publish"]) +def test_failed_automatic_freeze_preserves_published_artifact( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + file_freezer: mock.Mock, + stage: str, +) -> None: + output = tmp_path / "model-ase_single.pt2" + output.write_text("original") + + def fail_export(source: str, output: str, **kwargs: object) -> None: + Path(output).write_text("incomplete") + raise RuntimeError("export failed") + + load = calculator_module.DeepPot + + def fail_load(model_file: str, **kwargs: object) -> object: + if model_file.endswith(".pt2"): + raise RuntimeError("load failed") + return load(model_file, **kwargs) + + if stage == "export": + file_freezer.side_effect = fail_export + elif stage == "load": + monkeypatch.setattr(calculator_module, "DeepPot", fail_load) + else: + monkeypatch.setattr( + Path, "replace", mock.Mock(side_effect=RuntimeError("publish failed")) + ) + with pytest.raises(RuntimeError, match=f"{stage} failed"): + DP(tmp_path / "model.pt", freeze=True) + assert output.read_text() == "original" + assert list(tmp_path.iterdir()) == [output] + + +@pytest.mark.parametrize("skin", [None, 0.75]) +def test_graph_neighbor_options_reach_evaluator( + tmp_path: Path, skin: float | None +) -> None: + checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4") + options = {"neighbor_skin": skin} if skin is not None else {} + with mock.patch("deepmd.calculator.DeepPot", wraps=DeepPot) as deep_pot: + DP(checkpoint, neighbor_graph_method="vesin", **options) + deep_pot.assert_called_once_with( + str(checkpoint.resolve()), + neighbor_list=None, + head=None, + nlist_backend="auto", + neighbor_graph_method="vesin", + neighbor_skin=1.0 if skin is None else skin, + ) + + +def test_native_spin_edge_freeze_preserves_neighbor_options(tmp_path: Path) -> None: + config = _tiny_sezm_native_spin_model_params() + # Edge inputs retain all cutoff neighbors independently of dense capacity. + config["descriptor"]["sel"] = [1, 0] + checkpoint = _write_checkpoint(tmp_path / "native.pt", "pt", config=config) + original = DeepPot( + str(checkpoint), neighbor_graph_method="vesin", neighbor_skin=1.0 + ) + frozen = tmp_path / "native.pt2" + freeze_pt2.freeze_sezm_to_pt2(str(checkpoint), str(frozen), with_comm=False) + coords = np.array([[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [-0.3, 0.8, 0.0]]]) + cells = np.eye(3).reshape(1, 9) * 8.0 + types = np.array([0, 1, 0]) + spins = np.array([[[0.2, 0.3, 0.9], [0.0, 0.0, 0.0], [-0.1, 0.6, 0.7]]]) + for method in ("vesin", "auto"): + exported = DeepPot(str(frozen), neighbor_graph_method=method, neighbor_skin=1.0) + for displacement in (0.0, 0.1, 0.7): + moved = coords.copy() + moved[0, 1, 0] += displacement + expected = original.eval(moved, cells, types, spin=spins, atomic=True) + actual = exported.eval(moved, cells, types, spin=spins, atomic=True) + for value, reference in zip(actual, expected, strict=True): + np.testing.assert_allclose(value, reference, rtol=1e-6, atol=1e-8) + + native = DeepPot(str(frozen), nlist_backend="native", auto_batch_size=False) + batch_coords = np.repeat(coords, 2, axis=0) + batch_cells = np.repeat(cells, 2, axis=0) + batch_types = np.array([[0, 1, 0], [1, 0, 0]]) + batch_spins = np.repeat(spins, 2, axis=0) + actual = native.eval( + batch_coords, + batch_cells, + batch_types, + spin=batch_spins, + atomic=True, + mixed_type=True, + ) + frames = [ + original.eval( + batch_coords[frame : frame + 1], + batch_cells[frame : frame + 1], + batch_types[frame], + spin=batch_spins[frame : frame + 1], + atomic=True, + ) + for frame in range(2) + ] + for value, first, second in zip(actual, *frames, strict=True): + np.testing.assert_allclose( + value, np.concatenate((first, second), axis=0), rtol=1e-6, atol=1e-8 + ) + + +@pytest.mark.parametrize("skin", [-1.0, float("inf"), float("nan")]) +def test_invalid_neighbor_skin_rejected(skin: float) -> None: + with pytest.raises(ValueError, match="neighbor_skin"): + DP("unused.pt", neighbor_skin=skin) + + +@pytest.mark.parametrize( + "kind", ["pt", "pt-dpa4", "dpa4", "dpa4c", "dpa1", "dpa2", "pt-dense", "dense"] +) +@pytest.mark.parametrize("inference_device", ["cpu", "cuda"], indirect=True) +def test_cached_neighbors_and_types_match_fresh_evaluation( + tmp_path: Path, kind: str, inference_device: str +) -> None: + checkpoint = _write_checkpoint(tmp_path / "model.pt", kind) + options = ( + {"nlist_backend": "vesin"} + if kind in ("pt-dense", "dense") + else {"neighbor_graph_method": "vesin"} + ) + cached = DP(checkpoint, **options) + fresh = DP(checkpoint, neighbor_skin=0.0, **options) + positions = np.array([[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [-0.3, 0.8, 0.0]]) + for step in range(5): + points = positions.copy() + points[1, 0] += 0.1 * step + symbols = "OH2" if step < 2 else "HOH" + cell = [8.0, 8.0, 8.0] if step < 3 else [7.5, 8.0, 8.0] + if step == 4: + cached.type_dict = fresh.type_dict = {"O": 1, "H": 0} + atoms = Atoms(symbols, positions=points, cell=cell, pbc=True) + results = [] + for calculator in (cached, fresh): + atoms.calc = calculator + results.append( + (atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress()) + ) + fp32 = kind == "dpa4c" or inference_device == "cuda" + rtol, atol = (1e-4, 1e-5) if fp32 else (1e-7, 1e-8) + for actual, expected in zip(*results, strict=True): + np.testing.assert_allclose(actual, expected, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("kind", ["pt", "pt-dpa4", "dpa4", "dpa4c", "dpa1", "dpa2"]) +def test_checkpoint_default_warns_without_export( + tmp_path: Path, caplog: pytest.LogCaptureFixture, kind: str +) -> None: + checkpoint = _write_checkpoint(tmp_path / "model.pt", kind) + with caplog.at_level(logging.WARNING, logger="deepmd.calculator"): + calculator = DP(checkpoint) + + assert calculator.dp.deep_eval.model_path == str(checkpoint) + assert not checkpoint.with_name("model-ase_single.pt2").exists() + warnings = [ + record.getMessage() + for record in caplog.records + if record.name == "deepmd.calculator" + ] + assert len(warnings) == 1 + assert "SUBSTANTIALLY slower" in warnings[0] + assert "freeze=True" in warnings[0] + assert "manually" in warnings[0] + assert "debugging" in warnings[0] + + +@pytest.mark.parametrize( + ("kind", "head"), + [("pt", "branch"), ("dpa4", "branch"), ("dpa4c", "branch"), ("dpa4", None)], +) +@pytest.mark.parametrize("override", [None, ("0", "0"), ("2", "3")]) +def test_automatic_freeze_reuses_kernel_policy( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + kind: str, + head: str | None, + override: tuple[str, str] | None, +) -> None: + checkpoint = _write_checkpoint(tmp_path / "model.pt", kind, head=head or "Default") + monkeypatch.setattr(freeze_pt2, "DEVICE", torch.device("cuda")) + export_policy = serialization._fused_operators_for_export + + @contextlib.contextmanager + def cuda_export_policy() -> Iterator[None]: + # Only policy selection sees CUDA; lazy imports retain the real query. + with contextlib.ExitStack() as stack: + with mock.patch( + "deepmd.pt_expt.kernels.utils.backend_device_type", return_value="cuda" + ): + stack.enter_context(export_policy()) + yield + + monkeypatch.setattr( + serialization, "_fused_operators_for_export", cuda_export_policy + ) + if override is not None: + for name, value in zip(_LEVELS[:2], override, strict=True): + monkeypatch.setenv(name, value) + captured = {} + + def capture(*args: object, **kwargs: object) -> None: + captured["levels"] = tuple(os.environ.get(name) for name in _LEVELS[:2]) + captured["with_comm"] = kwargs["with_comm"] + Path(args[1] if kind == "pt" else args[0]).write_bytes(b"export") + + if kind == "pt": + monkeypatch.setattr(freeze_pt2, "_freeze_sezm_to_pt2", capture) + else: + monkeypatch.setattr(serialization, "_deserialize_to_file_pt2", capture) + + def load(model_file: str, **kwargs: object) -> object: + if model_file.endswith(".pt"): + return DeepPot(model_file, **kwargs) + result = mock.Mock() + result.get_type_map.return_value = ["O", "H"] + result.get_ntypes.return_value = 2 + return result + + with mock.patch("deepmd.calculator.DeepPot", side_effect=load) as deep_pot: + DP( + checkpoint, + freeze=True, + head=head, + neighbor_graph_method="vesin", + neighbor_skin=0.75, + ) + + expected = override or ("1", "2") + if kind == "dpa4c": + expected = ("2", None if override is None else override[1]) + assert captured == {"levels": expected, "with_comm": False} + assert tuple(os.environ.get(name) for name in _LEVELS[:2]) == ( + override or (None, None) + ) + assert deep_pot.call_count == 2 + private = Path(deep_pot.call_args.args[0]) + assert private.name == "model-ase_single.pt2" + assert private.parent != tmp_path + assert (tmp_path / private.name).read_bytes() == b"export" + deep_pot.assert_called_with( + str(private), + neighbor_list=None, + nlist_backend="auto", + neighbor_graph_method="vesin", + neighbor_skin=0.75, + ) + + +@pytest.mark.parametrize("kind", ["pt-dense", "dense", "dpa3"]) +@pytest.mark.parametrize("freeze", [False, True]) +def test_ineligible_checkpoint_keeps_eager_inference( + tmp_path: Path, caplog: pytest.LogCaptureFixture, kind: str, freeze: bool +) -> None: + checkpoint = _write_checkpoint(tmp_path / "model.pt", kind) + with caplog.at_level(logging.WARNING, logger="deepmd.calculator"): + calculator = DP(checkpoint, freeze=freeze) + assert calculator.dp.deep_eval.model_path == str(checkpoint) + assert not any(record.name == "deepmd.calculator" for record in caplog.records) + assert not checkpoint.with_name("model-ase_single.pt2").exists() + + +def test_single_task_pt_freeze_ignores_head( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Single-task head arguments retain the eager checkpoint's semantics.""" + checkpoint = _write_checkpoint(tmp_path / "model.pt", "pt") + loaded = DeepPot(str(checkpoint), head="Default") + + def write_artifact(source: str, output: str, **kwargs: object) -> None: + Path(output).write_bytes(b"export") + + export = mock.Mock(side_effect=write_artifact) + monkeypatch.setattr(freeze_pt2, "freeze_sezm_to_pt2", export) + with mock.patch("deepmd.calculator.DeepPot", return_value=loaded): + DP(checkpoint, head="Default", freeze=True) + export.assert_called_once_with( + str(checkpoint), + mock.ANY, + head=None, + with_comm=False, + ) + + +@pytest.mark.parametrize("kind", ["pt", "dpa4", "dpa4c", "dpa1", "dpa2"]) +@pytest.mark.parametrize("inference_device", ["cpu", "cuda"], indirect=True) +def test_automatic_freeze_single_artifact_matches_eager( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + kind: str, + inference_device: str, +) -> None: + config = _config(kind) + options = {} + if kind == "pt": + config["descriptor"]["sel"] = [1, 0] + options["nlist_backend"] = "native" + checkpoint = _write_checkpoint(tmp_path / "model.pt", kind, config=config) + original = checkpoint.read_bytes() + atoms = Atoms( + "OH2", + positions=[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [-0.3, 0.8, 0.0]], + cell=[8.0, 8.0, 8.0], + pbc=True, + ) + atoms.calc = DP(checkpoint, **options) + expected = (atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress()) + caplog.clear() + with mock.patch( + "torch._inductor.aoti_compile_and_package", + wraps=torch._inductor.aoti_compile_and_package, + ) as compile_package: + atoms.calc = DP(checkpoint, freeze=True, **options) + assert compile_package.call_count == 1 + if kind == "dpa4c": + exported = compile_package.call_args.args[0] + targets = {str(node.target) for node in exported.graph_module.graph.nodes} + assert any("dpa4c_" in target and "compress" in target for target in targets) + frozen = tmp_path / "model-ase_single.pt2" + assert checkpoint.read_bytes() == original + private = Path(atoms.calc.dp.deep_eval.model_path) + assert private != frozen + assert private.read_bytes() == frozen.read_bytes() + with zipfile.ZipFile(frozen) as archive: + metadata = json.loads(archive.read("model/extra/metadata.json")) + assert metadata["has_comm_artifact"] is False + assert "model/extra/forward_lower_with_comm.pt2" not in archive.namelist() + actual = (atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress()) + assert np.max(np.abs(expected[1])) > 1e-8 + fp32 = inference_device == "cuda" or kind == "dpa4c" + rtol, atol = (1e-4, 1e-5) if fp32 else (1e-6, 1e-8) + for reference, result in zip(expected, actual, strict=True): + np.testing.assert_allclose(result, reference, rtol=rtol, atol=atol) + assert not any( + record.name == "deepmd.calculator" and record.levelno >= logging.WARNING + for record in caplog.records + ) + atoms.calc = DP(frozen, freeze=True) + reload_tol = 1e-6 if fp32 else 1e-12 + np.testing.assert_allclose( + atoms.get_forces(), actual[1], rtol=reload_tol, atol=reload_tol + ) + assert not (tmp_path / "model-ase_single-ase_single.pt2").exists() diff --git a/source/tests/pt_expt/test_freeze_compress.py b/source/tests/pt_expt/test_freeze_compress.py new file mode 100644 index 0000000000..091f884530 --- /dev/null +++ b/source/tests/pt_expt/test_freeze_compress.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Single-pass compressed export through the CLI and ASE calculator.""" + +import json +import logging +import zipfile +from pathlib import ( + Path, +) +from unittest import ( + mock, +) + +import numpy as np +import pytest +import torch +from ase import ( + Atoms, +) + +from deepmd.calculator import ( + DP, +) +from deepmd.infer import ( + DeepPot, +) +from deepmd.main import ( + parse_args, +) +from deepmd.pt_expt.entrypoints import compress as compression +from deepmd.pt_expt.entrypoints import main as entrypoints +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, +) +from deepmd.pt_expt.utils import ( + env, + serialization, +) +from deepmd.utils.vacuum_reference import ( + reference_charge_spin, +) + +from .test_calculator import ( + _config, + _write_checkpoint, +) + + +@pytest.fixture(autouse=True) +def cpu_export(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(env, "DEVICE", torch.device("cpu")) + for name in ("DP_CUDA_INFER", "DP_TRITON_INFER", "DP_COMPILE_INFER"): + monkeypatch.delenv(name, raising=False) + + +@pytest.mark.parametrize("compress", [False, True]) +def test_freeze_cli_compression_option(tmp_path: Path, compress: bool) -> None: + checkpoint = tmp_path / "model.pt" + checkpoint.touch() + args = ["--pt-expt", "freeze", "-c", str(checkpoint)] + if compress: + args.append("--compress") + flags = parse_args(args) + with mock.patch.object(entrypoints, "freeze") as freeze: + entrypoints.main(flags) + assert freeze.call_args.kwargs["compress"] is compress + + +@pytest.mark.parametrize("backend", ["pt", "tf", "pd", "jax"]) +def test_freeze_cli_rejects_unsupported_compression( + backend: str, capsys: pytest.CaptureFixture[str] +) -> None: + """Unsupported backends reject compression instead of ignoring the flag.""" + with pytest.raises(SystemExit) as error: + parse_args([f"--{backend}", "freeze", "--compress"]) + assert error.value.code == 2 + assert "freeze --compress requires --pt-expt" in capsys.readouterr().err + + +def test_freeze_missing_neighbor_statistics_explains_compression_command( + tmp_path: Path, +) -> None: + """Models needing data statistics point to the CLI that accepts them.""" + checkpoint = _write_checkpoint(tmp_path / "model.pt", "dense") + output = tmp_path / "model.pte" + with pytest.raises(ValueError, match="dp --pt-expt compress") as error: + entrypoints.freeze(str(checkpoint), str(output), compress=True) + assert "without --compress" in str(error.value) + assert not output.exists() + + +@pytest.mark.parametrize("compress", [False, True]) +def test_freeze_compression_changes_snapshot_before_export( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, compress: bool +) -> None: + config = _config("dpa4c") + config["descriptor"].update(add_chg_spin_ebd=True, default_chg_spin=[1, 2]) + config["fitting_net"]["vacuum_ref"] = True + config["preset_out_bias"] = {"energy": {"O": 0.0, "H": 0.0}} + checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4c", config=config) + original = checkpoint.read_bytes() + reference = get_model(config).eval() + ModelWrapper(reference, model_params=config).load_state_dict( + torch.load(checkpoint, map_location="cpu", weights_only=True)["model"] + ) + fitting = reference.atomic_model.fitting_net + original_bias = fitting.bias_atom_e.detach().numpy().copy() + fitting.vacuum_ref = False + with torch.no_grad(): + fitting.bias_atom_e.zero_() + reference.atomic_model.out_bias.zero_() + coord = torch.zeros((2, 1, 3), dtype=torch.float64) + atype = torch.arange(2, dtype=torch.int64).reshape(2, 1) + neutral = ( + reference.call_common( + coord, + atype, + None, + charge_spin=torch.as_tensor(reference_charge_spin(config["type_map"])), + )["energy"] + .detach() + .numpy() + .reshape(2, 1) + ) + deployment = ( + reference.call_common(coord, atype, None)["energy"] + .detach() + .numpy() + .reshape(2, 1) + ) + # The deployment state must differ from the neutral isolated-atom reference. + assert np.max(np.abs(deployment - neutral)) > 1e-5 + enable_compression = type(reference).enable_compression + + def check_reference_before_compression( + model: torch.nn.Module, *args: object, **kwargs: object + ) -> None: + np.testing.assert_allclose( + model.atomic_model.fitting_net.bias_atom_e.detach().numpy(), + original_bias - neutral, + rtol=1e-6, + atol=1e-7, + ) + enable_compression(model, *args, **kwargs) + + monkeypatch.setattr( + type(reference), "enable_compression", check_reference_before_compression + ) + captured = {} + + def capture(model_file: str, data: dict, **kwargs: object) -> None: + captured.update(data=data, options=kwargs) + + monkeypatch.setattr(serialization, "deserialize_to_file", capture) + monkeypatch.setattr(compression, "deserialize_to_file", capture) + entrypoints.freeze( + str(checkpoint), + str(tmp_path / "model.pt2"), + compress=compress, + with_comm=False, + ) + model = captured["data"]["model"] + assert ("compress" in model["descriptor"]) is compress + assert model["fitting"]["vacuum_ref"] is (not compress) + np.testing.assert_allclose( + model["fitting"]["@variables"]["bias_atom_e"], + original_bias - neutral if compress else original_bias, + rtol=1e-6, + atol=1e-7, + ) + assert captured["options"]["with_comm"] is False + assert captured["options"]["lower_kind"] == ("auto" if compress else "graph") + assert checkpoint.read_bytes() == original + + +@pytest.mark.parametrize("case", ["precision", "exclusions", "charge_state"]) +def test_ase_falls_back_but_explicit_compression_reports_ineligibility( + tmp_path: Path, caplog: pytest.LogCaptureFixture, case: str +) -> None: + config = _config("dpa4c") + if case == "precision": + config["descriptor"]["precision"] = "float64" + config["fitting_net"]["precision"] = "float64" + error = "requires descriptor precision" + elif case == "exclusions": + config["descriptor"]["exclude_types"] = [[0, 1]] + error = "non-empty `exclude_types`" + else: + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = None + error = "requires `default_chg_spin`" + checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4c", config=config) + loaded = DeepPot(str(checkpoint)) + + def write_artifact(source: str, output: str, **kwargs: object) -> None: + Path(output).write_bytes(b"export") + + with ( + mock.patch("deepmd.calculator.DeepPot", return_value=loaded), + mock.patch.object(entrypoints, "freeze", side_effect=write_artifact) as freeze, + caplog.at_level(logging.WARNING, logger="deepmd.calculator"), + ): + DP(checkpoint, freeze=True) + assert freeze.call_args.kwargs["compress"] is False + assert "Freezing without compression" in caplog.text + assert error in caplog.text + output = tmp_path / "explicit.pt2" + with pytest.raises(ValueError, match=error): + entrypoints.freeze(str(checkpoint), str(output), compress=True) + assert not output.exists() + + +def test_compressed_freeze_compiles_one_cpu_artifact(tmp_path: Path) -> None: + checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4c") + output = tmp_path / "model.pt2" + + def compile_package( + exported: torch.export.ExportedProgram, package_path: str, **kwargs: object + ) -> None: + with zipfile.ZipFile(package_path, "w") as archive: + archive.writestr("model/data.pkl", b"") + + with mock.patch( + "torch._inductor.aoti_compile_and_package", side_effect=compile_package + ) as compile_mock: + entrypoints.freeze(str(checkpoint), str(output), compress=True, with_comm=False) + assert compile_mock.call_count == 1 + with zipfile.ZipFile(output) as archive: + metadata = json.loads(archive.read("model/extra/metadata.json")) + assert metadata["has_comm_artifact"] is False + assert metadata["lower_input_kind"] == "graph" + assert "model/extra/forward_lower_with_comm.pt2" not in archive.namelist() + snapshot = serialization.serialize_from_file(str(output)) + assert "compress" in snapshot["model"]["descriptor"] + + +def test_ase_float64_fallback_exports_and_evaluates( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + config = _config("dpa4c") + config["descriptor"]["precision"] = "float64" + config["fitting_net"]["precision"] = "float64" + checkpoint = _write_checkpoint(tmp_path / "model.pt", "dpa4c", config=config) + atoms = Atoms( + "OH2", + positions=[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [-0.3, 0.8, 0.0]], + cell=[8.0, 8.0, 8.0], + pbc=True, + calculator=DP(checkpoint), + ) + expected = atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress() + caplog.clear() + with caplog.at_level(logging.WARNING, logger="deepmd.calculator"): + atoms.calc = DP(checkpoint, freeze=True) + actual = atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress() + assert "Freezing without compression" in caplog.text + snapshot = serialization.serialize_from_file(str(tmp_path / "model-ase_single.pt2")) + assert "compress" not in snapshot["model"]["descriptor"] + for reference, result in zip(expected, actual, strict=True): + np.testing.assert_allclose(result, reference, rtol=1e-6, atol=1e-8) diff --git a/source/tests/pt_expt/utils/test_canonical_graph.py b/source/tests/pt_expt/utils/test_canonical_graph.py index 3128db789a..03b2eee505 100644 --- a/source/tests/pt_expt/utils/test_canonical_graph.py +++ b/source/tests/pt_expt/utils/test_canonical_graph.py @@ -9,6 +9,7 @@ UINT32_MAX, CanonicalGraph, canonical_graph_from_neighbor_graph, + canonical_graph_from_pairs, validate_canonical_graph_shapes, ) @@ -73,3 +74,59 @@ def test_storage_exceeding_uint32_range_is_rejected() -> None: ) with pytest.raises(ValueError, match="exceeds the uint32 range"): validate_canonical_graph_shapes(graph, 1) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("edge_count", [0, 1, 5]) +def test_pairs_preserve_geometry_and_both_csr_views( + device: str, edge_count: int +) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is required") + center = torch.tensor([3, 1, 0, 3, 1], device=device)[:edge_count] + source = torch.tensor([2, 0, 1, 2, 1], device=device)[:edge_count] + vectors = torch.arange(15, dtype=torch.float64, device=device).reshape(5, 3) + graph = canonical_graph_from_pairs( + center, + source, + vectors[:edge_count], + torch.tensor([2, 99, 2, 99], device=device)[::2], + 4, + ) + validate_canonical_graph_shapes(graph, 4) + torch.testing.assert_close( + graph.destination_row_ptr.diff(), torch.bincount(center, minlength=4) + ) + torch.testing.assert_close( + graph.source_row_ptr.diff(), torch.bincount(source, minlength=4) + ) + destination = torch.repeat_interleave( + torch.arange(4, device=device), graph.destination_row_ptr.diff() + ) + payload = ( + torch.cat( + ( + destination[:, None], + graph.source[:edge_count].to(torch.int64)[:, None], + graph.edge_vec[:edge_count], + ), + dim=1, + ) + .cpu() + .tolist() + ) + expected = ( + torch.cat((center[:, None], source[:, None], vectors[:edge_count]), dim=1) + .cpu() + .tolist() + ) + assert sorted(payload) == sorted(expected) + by_source = graph.source.to(torch.int64).index_select( + 0, graph.source_order[:edge_count].to(torch.int64) + ) + torch.testing.assert_close( + by_source, + torch.repeat_interleave( + torch.arange(4, device=device), graph.source_row_ptr.diff() + ), + ) 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 74d7b17246..2f47624296 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -21,9 +21,16 @@ import pytest import torch +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + ModelOutputDef, +) from deepmd.entrypoints.convert_backend import ( convert_backend, ) +from deepmd.pt_expt.infer.deep_eval import ( + DeepEval, +) from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, ) @@ -508,12 +515,40 @@ def test_neighbor_graph_method_rejected_on_nlist_artifact(dpa1_dpmodel_data) -> copy.deepcopy(dpa1_dpmodel_data), do_atomic_virial=True, ) - with pytest.raises(ValueError, match="graph-form"): + with pytest.raises(ValueError, match="use nlist_backend"): DeepPot(p, neighbor_graph_method="vesin") # The default remains valid for nlist-form artifacts. DeepPot(p) +@pytest.mark.parametrize( + ("method", "backend", "message"), + [ + ("dense", "auto", "Edge-vector artifacts support neighbor_graph_method"), + ("vesin", "native", "Select only one"), + ], +) +def test_edge_vec_neighbor_options_rejected( + monkeypatch: pytest.MonkeyPatch, method: str, backend: str, message: str +) -> None: + """The constructor validates edge options independently of compiled loading.""" + + def load_metadata(evaluator: DeepEval, model_file: str) -> None: + evaluator.metadata = {"lower_input_kind": "edge_vec"} + + monkeypatch.setattr(DeepEval, "_load_pt2", load_metadata) + with ( + pytest.warns(DeprecationWarning, match="deprecated edge_vec lower schema"), + pytest.raises(ValueError, match=message), + ): + DeepEval( + "edge.pt2", + ModelOutputDef(FittingOutputDef([])), + neighbor_graph_method=method, + nlist_backend=backend, + ) + + class _FakeDesc: def __init__(self, n_attn: int) -> None: self._n = n_attn diff --git a/source/tests/pt_expt/utils/test_serialization_kernel_levels.py b/source/tests/pt_expt/utils/test_serialization_kernel_levels.py index 53651103ca..0816908353 100644 --- a/source/tests/pt_expt/utils/test_serialization_kernel_levels.py +++ b/source/tests/pt_expt/utils/test_serialization_kernel_levels.py @@ -1,11 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Kernel-level selection for pt_expt serialization.""" +import json import os import sys +import zipfile +from pathlib import ( + Path, +) from types import ( ModuleType, ) +from unittest import ( + mock, +) import pytest import torch @@ -261,3 +269,41 @@ def test_target_policy_only_suppresses_incompatible_dpa4_accelerators( assert tuple(os.environ[name] for name in _INFER_LEVELS) == expected assert all(os.environ[name] == "1" for name in _INFER_LEVELS) + + +@pytest.mark.parametrize("single_process", [False, True]) +def test_communication_artifact_is_optional( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, single_process: bool +) -> None: + """The default includes communication; a single-process export omits it.""" + data = {"model": {"type": "dpa4"}} + + def trace(*args: object, **kwargs: object) -> tuple: + return ( + None, + {"has_comm_artifact": True, "has_chg_spin_ebd": False}, + data, + ["energy"], + ) + + def compile_package(exported: object, package_path: str, **kwargs: object) -> None: + with zipfile.ZipFile(package_path, "w") as archive: + archive.writestr("model/data.pkl", b"") + + monkeypatch.setattr(serialization, "_trace_and_export", trace) + output = tmp_path / "model.pt2" + kwargs = {"with_comm": False} if single_process else {} + with mock.patch( + "torch._inductor.aoti_compile_and_package", side_effect=compile_package + ) as compile_mock: + serialization.deserialize_to_file( + str(output), data, lower_kind="graph", **kwargs + ) + + assert compile_mock.call_count == (1 if single_process else 2) + with zipfile.ZipFile(output) as archive: + metadata = json.loads(archive.read("model/extra/metadata.json")) + assert metadata["has_comm_artifact"] is (not single_process) + assert ("model/extra/forward_lower_with_comm.pt2" in archive.namelist()) is ( + not single_process + ) diff --git a/source/tests/pt_expt/utils/test_tabulate_ops_lazy.py b/source/tests/pt_expt/utils/test_tabulate_ops_lazy.py index 841ada65c3..9b52b9b103 100644 --- a/source/tests/pt_expt/utils/test_tabulate_ops_lazy.py +++ b/source/tests/pt_expt/utils/test_tabulate_ops_lazy.py @@ -23,6 +23,9 @@ import textwrap import torch +from torch._subclasses.fake_tensor import ( + FakeTensorMode, +) from deepmd.pt_expt.utils import ( tabulate_ops, @@ -105,3 +108,21 @@ def _fallback(*args, **kwargs): ns.__dict__.pop(op_name, None) tabulate_ops._registered.clear() tabulate_ops._registered.update(saved_registered) + + +def test_composite_operator_uses_its_decomposition_for_fake_tensors() -> None: + """Composite operators retain PyTorch's fake-tensor decomposition.""" + with torch.library._scoped_library("deepmd_test_tabulate", "DEF") as library: + library.define("identity(Tensor x) -> Tensor") + library.impl("identity", lambda x: x + 1, "CompositeImplicitAutograd") + + def fake(value: torch.Tensor) -> torch.Tensor: + raise AssertionError("Composite operators must use their decomposition") + + tabulate_ops._try_register_fake("deepmd_test_tabulate::identity", fake) + with FakeTensorMode(): + value = torch.ones((3, 4), device="cpu") + result = torch.ops.deepmd_test_tabulate.identity(value) + assert result.shape == value.shape + assert result.dtype == value.dtype + assert result.device == value.device diff --git a/source/tests/pt_expt/utils/test_vesin_graph_builder.py b/source/tests/pt_expt/utils/test_vesin_graph_builder.py index c216676994..6777f71d45 100644 --- a/source/tests/pt_expt/utils/test_vesin_graph_builder.py +++ b/source/tests/pt_expt/utils/test_vesin_graph_builder.py @@ -1,4 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from dataclasses import ( + replace, +) +from unittest import ( + mock, +) + import numpy as np import pytest import torch @@ -10,11 +17,18 @@ apply_pair_exclusion, build_neighbor_graph, ) +from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + validate_canonical_graph_shapes, +) vesin_builder = pytest.importorskip("deepmd.pt_expt.utils.vesin_graph_builder") from deepmd.pt_expt.utils.vesin_neighbor_list import ( is_vesin_torch_available, ) +from deepmd.pt_expt.utils.vesin_search import ( + VesinNeighborSearch, +) pytestmark = pytest.mark.skipif( not is_vesin_torch_available(), reason="vesin[torch] not installed" @@ -113,9 +127,9 @@ def test_vesin_excludes_virtual_atoms_like_dense(): def _valid_edge_set(ng): """Return the set of (src, dst, rounded edge_vec) for all real edges.""" - ei = np.asarray(ng.edge_index) - ev = np.asarray(ng.edge_vec) - em = np.asarray(ng.edge_mask) + ei = ng.edge_index.detach().cpu().numpy() + ev = ng.edge_vec.detach().cpu().numpy() + em = ng.edge_mask.detach().cpu().numpy() return { (int(ei[0, k]), int(ei[1, k]), tuple(np.round(ev[k], 6))) for k in range(ei.shape[1]) @@ -172,3 +186,200 @@ def test_vesin_nlist_edges_pair_excl_raises(): nl = VesinNeighborList() with pytest.raises(NotImplementedError, match="return_mode='edges'"): nl.build(coord, atype, None, 2.0, [4], return_mode="edges", pair_excl=pe) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_cached_graph_tracks_geometry_and_atom_changes(device: str) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is required") + coord, atype, box = _system(True) + coord = coord.reshape(1, 4, 3).to(device) + atype = atype.to(device) + box = box.reshape(1, 3, 3).to(device) + search = VesinNeighborSearch(skin=1.0, sorted=True) + for step in range(9): + points = coord.clone() + types = atype.clone() + cell = box.clone() + if step == 1: + points[0, 1, 0] += 0.2 + elif step == 2: + points[0, 1, 0] += 0.8 + elif step == 3: + cell *= 1.15 + elif step == 4: + cell = None + elif step == 5: + types[0, 1] = -1 + elif step == 6: + points = torch.cat((points, points[:, :1] + 0.4), dim=1) + types = torch.cat((types, types[:, :1]), dim=1) + elif step == 7: + cell[0, 0, 1] += 0.3 + elif step == 8: + points[0, 1] += cell[0, 0] + actual = vesin_builder.build_neighbor_graph_vesin( + points, types, cell, 2.0, search=search, canonicalize=True + ) + expected = build_neighbor_graph(points, types, cell, 2.0, canonicalize=True) + assert _valid_edge_set(actual) == _valid_edge_set(expected) + torch.testing.assert_close(actual.n_node, expected.n_node) + assert actual.destination_sorted + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_cached_search_preserves_earlier_autograd_inputs(device: str) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is required") + coord, atype, box = _system(True) + coord = coord.reshape(1, 4, 3).to(device).requires_grad_() + atype = atype.to(device) + box = box.reshape(1, 3, 3).to(device) + search = VesinNeighborSearch(skin=1.0, sorted=True) + graph = vesin_builder.build_neighbor_graph_vesin( + coord, atype, box, 2.0, search=search + ) + moved = coord.detach().clone() + moved[0, 1, 0] += 0.3 + vesin_builder.build_neighbor_graph_vesin(moved, atype, box, 2.0, search=search) + actual = torch.autograd.grad(graph.edge_vec.square().sum(), coord)[0] + reference = build_neighbor_graph(coord, atype, box, 2.0) + expected = torch.autograd.grad(reference.edge_vec.square().sum(), coord)[0] + torch.testing.assert_close(actual, expected) + + +def test_cached_search_keeps_frames_independent() -> None: + coord, atype, box = _system(True) + coord = coord.unsqueeze(0).repeat(2, 1, 1) + atype = atype.repeat(2, 1) + box = box.unsqueeze(0).repeat(2, 1, 1) + box[1] *= 1.5 + search = VesinNeighborSearch(skin=1.0, sorted=True) + for _ in range(3): + actual = vesin_builder.build_neighbor_graph_vesin( + coord, atype, box, 2.0, search=search, canonicalize=True + ) + expected = build_neighbor_graph(coord, atype, box, 2.0, canonicalize=True) + assert _valid_edge_set(actual) == _valid_edge_set(expected) + coord[1, 1, 0] += 0.2 + + +@pytest.mark.parametrize("return_vectors", [False, True]) +def test_cached_search_rebuilds_for_cell_contraction(return_vectors: bool) -> None: + """Cell contraction includes images absent from the previous candidates.""" + from vesin_torch import ( + NeighborList, + ) + + points = torch.zeros((1, 3), dtype=torch.float64) + cell = torch.diag(torch.tensor([2.35, 10.0, 10.0], dtype=torch.float64)) + other_cell = cell.clone() + search = VesinNeighborSearch(skin=1.0) + with mock.patch("vesin_torch.NeighborList", wraps=NeighborList) as constructor: + for frame in (0, 1): + ii, _, _ = search.compute(points, cell, 6.0, frame=frame) + assert ii.numel() == 4 + search.compute(points, cell.clone(), 6.0) + assert constructor.call_count == 2 + + # The third image moves from 7.05 to 5.85 angstrom, crossing both + # the old candidate radius (7.0) and the physical cutoff (6.0). + cell[0, 0] = 1.95 + ii, jj, geometry = search.compute( + points, cell, 6.0, return_vectors=return_vectors + ) + assert ii.numel() == 6 + assert torch.count_nonzero(ii) == torch.count_nonzero(jj) == 0 + expected = torch.tensor([-3, -2, -1, 1, 2, 3], dtype=geometry.dtype) + if return_vectors: + expected *= cell[0, 0] + torch.testing.assert_close(geometry[:, 0].sort().values, expected) + assert torch.count_nonzero(geometry[:, 1:]) == 0 + assert constructor.call_count == 3 + + ii, _, _ = search.compute(points, other_cell, 6.0, frame=1) + assert ii.numel() == 4 + assert constructor.call_count == 3 + ii, _, _ = search.compute(points, None, 6.0) + assert ii.numel() == 0 + ii, _, _ = search.compute(points, cell, 6.0) + assert ii.numel() == 6 + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("periodic", [False, True]) +@pytest.mark.parametrize("nloc", [0, 1, 4]) +@pytest.mark.parametrize("nframes", [1, 2]) +def test_canonical_vesin_matches_dense( + device: str, periodic: bool, nloc: int, nframes: int +) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is required") + coord, atype, box = _system(periodic) + coord = coord[:nloc].unsqueeze(0).repeat(nframes, 1, 1).to(device) + atype = atype[:, :nloc].repeat(nframes, 1).to(device) + box = box.unsqueeze(0).repeat(nframes, 1, 1).to(device) if periodic else None + if nloc == 4: + atype[0, 1] = -1 + search = VesinNeighborSearch(skin=1.0) + for step in range(3): + points = coord.clone() + if nloc: + points[-1, 0, 0] += step * 0.3 + actual = vesin_builder.build_canonical_graph_vesin( + points, box, 2.0, search, node_mask=atype >= 0 if nloc == 4 else None + ) + reference = build_neighbor_graph(points, atype, box, 2.0, canonicalize=True) + expected = canonical_graph_from_neighbor_graph( + replace(reference, n_local=reference.n_node) + ) + validate_canonical_graph_shapes(actual, nframes * nloc) + torch.testing.assert_close(actual.n_node, expected.n_node) + torch.testing.assert_close( + actual.destination_row_ptr, expected.destination_row_ptr + ) + torch.testing.assert_close(actual.source_row_ptr, expected.source_row_ptr) + payloads = [] + for graph in (actual, expected): + edge_count = int(graph.destination_row_ptr[-1]) + destination = torch.repeat_interleave( + torch.arange(nframes * nloc, device=device), + graph.destination_row_ptr.diff(), + ) + values = ( + torch.cat( + ( + destination[:, None], + graph.source[:edge_count].to(torch.int64)[:, None], + graph.edge_vec[:edge_count], + ), + dim=1, + ) + .cpu() + .numpy() + ) + order = np.lexsort(values[:, ::-1].T) + payloads.append(values[order]) + np.testing.assert_allclose(*payloads, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("return_vectors", [False, True]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_cached_search_handles_empty_nonempty_transition( + device: str, return_vectors: bool +) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is required") + search = VesinNeighborSearch(skin=1.0) + for distance in (8.0, 0.8, 8.0): + positions = torch.tensor( + [[0.0, 0.0, 0.0], [distance, 0.0, 0.0]], + dtype=torch.float64, + device=device, + ) + ii, jj, geometry = search.compute( + positions, None, 2.0, return_vectors=return_vectors + ) + assert ii.numel() == (2 if distance < 2.0 else 0) + expected = positions[jj] - positions[ii] if return_vectors else geometry * 0 + torch.testing.assert_close(geometry, expected)