From 31271df161cfc7b2a9ed3ce64a135ce758d36e5d Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 5 Oct 2026 14:26:21 +0800 Subject: [PATCH 1/5] feat(ase): optionally freeze graph checkpoints for inference Add DP(..., freeze=True) for PT DPA4/SeZM checkpoints and graph-eligible PT-expt models. Determine eligibility from the loaded model and selected head, release the eager instance, export -ase_single.pt2 beside the checkpoint, and load the resulting AOTInductor artifact. Other model families and already frozen files retain their existing loading paths. Each opted-in calculator construction rebuilds the artifact; callers can load the resulting .pt2 directly to reuse it. Reuse the existing backend freeze entry points and kernel policies: - CUDA DPA4 defaults to DP_CUDA_INFER=1 and DP_TRITON_INFER=2, while explicit environment settings take precedence. - DPA4C graph export retains its required CUDA inference level of 2. - Add an optional with_comm argument to the exporters, defaulting to True for existing callers. ASE passes False, compiles only the regular artifact, and records has_comm_artifact=False. - Preserve multi-task head selection and normalize single-task PT heads to None, matching the eager loader's treatment of head arguments. Keep freeze=False as the default. Eligible checkpoints log a warning that direct inference can be SUBSTANTIALLY slower and recommend automatic or manual freezing. Document the eager path as useful for debugging and short runs that avoid the initial compilation overhead. Validation runs: - 44 checkpoint, head-selection, kernel-policy, and export-policy checks. - 43 CPU export and compatibility checks, including single-artifact ASE parity and preservation of the default communication artifact. - 2 existing ASE calculator tests, including frame/atomic parameters. - 5 real CUDA export/ASE round trips: PT SeZM/DPA4 and PT-expt DPA4, DPA4C, DPA1, and DPA2. Check energy, forces, stress, artifact reload, unchanged checkpoints, and exactly one compilation per ASE export. - Real dpa4c-air-v20260911 checkpoint: automatic export, numerical parity, artifact reload, and absence of the communication artifact verified. - Repository pre-commit checks passed. Preset dpa4c-air-v20260911 benchmark (float32, vacuum_ref enabled): NVIDIA RTX PRO 6000 Blackwell Server Edition, one GPU, PyTorch 2.13.0+cu130. Use three OMat24 validation structures and periodic supercells. Each mode runs 20 warmup steps and five timed blocks of 40 steps, alternating mode order. Change coordinates on every step to invalidate the ASE cache and synchronize the GPU around each block. Report median end-to-end ASE force evaluation time, including neighbor construction and data transfers; initialization, compilation, and warmup are excluded. Inference environment overrides are unset, so both modes use their default execution policies. Atoms Eager ms/step Frozen ms/step Speedup 16 15.85 4.54 3.49x 64 16.23 4.63 3.51x 128 16.35 4.68 3.50x 1024 17.20 5.91 2.91x 8192 24.17 11.86 2.04x In the benchmark process, eager initialization took 4.38 s, automatic freezing and loading took 50.91 s, and reloading the generated .pt2 took 0.26 s. Maximum frozen/eager differences across the sampled configurations were 4.47e-7 eV/atom in energy, 1.67e-5 eV/Angstrom in force components, and 2.91e-8 eV/Angstrom^3 in stress. The source checkpoint remained unchanged. feat(pt-expt): compress checkpoints during freeze Add --compress to the pt-expt freeze command and compress=True to its Python entry point. Share the model compression and export pipeline with the existing compress command, including neighbor-distance handling and default table settings. Build the compressed model before export without first compiling an intermediate uncompressed artifact. Fold vacuum references before specializing charge-dependent tables. ASE freeze=True requests compression for supported DPA4C models. A shared eligibility check supplies both explicit compression errors and automatic fallback reasons: unsupported precision, excluded type pairs, unsupported compiled profiles, and missing initial charge states. Automatic compression logs a warning and retains ordinary freezing when unsupported; explicit --compress reports the error. Other model families retain their existing ASE export behavior. Reuse the exporter's kernel policy and lower-schema selection: eligible CUDA models use the compact canonical ABI, CPU models use the graph ABI, and ASE still omits communication artifacts. Respect PyTorch's dispatch contract by leaving CompositeImplicitAutograd operators to their decompositions instead of registering independent fake/Meta kernels. Validation: - 57 focused checks passed, covering CLI dispatch, shared compression, vacuum folding, automatic fallback, explicit errors, head and kernel policies, fake-op registration, and the existing SE compression command. - Actual float64 ASE fallback export and numerical parity passed. - Actual compressed CPU and CUDA ASE exports passed energy/force/stress and reload comparisons, with compression operators present in the exported graphs and exactly one compiled artifact per export. - A trained dpa4c-air-v20260911 checkpoint passed CPU/CUDA automatic compression, parity, artifact reload, and checkpoint-integrity checks. - Repository pre-commit checks passed. Preset dpa4c-air-v20260911 benchmark: float32 with vacuum_ref enabled, PyTorch 2.13.0+cu130, inference environment overrides unset. Compare ordinary frozen artifacts with compressed frozen artifacts using the same OMat24 validation structures and periodic supercells. Change coordinates each step to invalidate ASE caching and alternate execution order. Timings are median end-to-end ASE force evaluations, including neighbor construction and transfers; export and warmup are excluded. CUDA: one NVIDIA RTX PRO 6000 Blackwell Server Edition, 20 warmup steps, five blocks of 40 steps, with GPU synchronization around each block. Atoms Frozen ms/step Compressed ms/step Speedup 16 4.67 3.64 1.28x 64 4.73 3.70 1.28x 128 4.82 3.75 1.28x 1024 6.03 4.99 1.21x 8192 12.00 6.19 1.94x CPU: Intel Xeon Platinum 8558P, OMP_NUM_THREADS=1 and one PyTorch compute thread, 10 warmup steps and three blocks of 20 steps. Atoms Frozen ms/step Compressed ms/step Speedup 16 41.66 1.55 26.91x 128 64.57 6.67 9.68x 1024 242.45 44.86 5.40x Maximum compressed/eager differences across these samples: - CUDA: energy 5.66e-7 eV/atom, force components 1.78e-5 eV/Angstrom, stress 7.25e-8 eV/Angstrom^3. - CPU: energy 6.29e-7 eV/atom, force components 1.43e-5 eV/Angstrom, stress 6.54e-8 eV/Angstrom^3. Compressed export and loading took 26.60 s on CUDA and 13.60 s on CPU in these benchmark processes. Compression trades artifact size for evaluation speed: the CUDA artifact grew from about 15 MiB to 145 MiB with its frozen tables. The checkpoints and previous exported artifacts were preserved. --- deepmd/calculator.py | 123 ++++++- deepmd/main.py | 7 + deepmd/pt/entrypoints/freeze_pt2.py | 16 +- deepmd/pt_expt/entrypoints/compress.py | 99 ++++-- deepmd/pt_expt/entrypoints/main.py | 26 ++ .../pt_expt/kernels/dpa4c/graph_compress.py | 73 ++-- deepmd/pt_expt/utils/serialization.py | 23 +- deepmd/pt_expt/utils/tabulate_ops.py | 6 + doc/model/dpa4c.md | 35 +- doc/third-party/ase.md | 32 ++ source/tests/pt_expt/test_calculator.py | 314 ++++++++++++++++++ source/tests/pt_expt/test_freeze_compress.py | 173 ++++++++++ .../utils/test_serialization_kernel_levels.py | 46 +++ .../pt_expt/utils/test_tabulate_ops_lazy.py | 21 ++ 14 files changed, 930 insertions(+), 64 deletions(-) create mode 100644 source/tests/pt_expt/test_calculator.py create mode 100644 source/tests/pt_expt/test_freeze_compress.py diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 0eca4776ed..1075a57690 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """ASE calculator interface module.""" +import logging +from functools import ( + partial, +) from pathlib import ( Path, ) @@ -23,6 +27,10 @@ ) if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + from ase import ( Atoms, ) @@ -30,8 +38,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 +135,18 @@ 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 + 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 + exports ``-ase_single.pt2`` beside the checkpoint, + using the backend's freeze-time kernel policy and omitting 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,15 +188,40 @@ def __init__( neighbor_list: Optional["NeighborList"] = None, head: str | None = None, nlist_backend: str = "auto", + freeze: bool = False, **kwargs: Any, ) -> None: Calculator.__init__(self, label=label, **kwargs) + 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, ) + 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) + freezer(str(model_path), str(output)) + self.dp = DeepPot( + str(output), neighbor_list=neighbor_list, nlist_backend=nlist_backend + ) + 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: diff --git a/deepmd/main.py b/deepmd/main.py index 1206b2fa0e..137bd8ac27 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( diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index b3cd33798a..708faa55f6 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -917,6 +917,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 +939,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 +959,7 @@ def freeze_sezm_to_pt2( target_device=target_device, head=head, atomic_virial=atomic_virial, + with_comm=with_comm, ) @@ -965,6 +970,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 +1092,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 +1114,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_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 9aeae29ca3..d70a9b5210 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; " @@ -95,8 +147,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 +167,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 +183,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/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/serialization.py b/deepmd/pt_expt/utils/serialization.py index b3edf999ad..c1df31c6dd 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1647,6 +1647,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. @@ -1690,6 +1692,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( @@ -1717,6 +1723,7 @@ def deserialize_to_file( model_json_override, do_atomic_virial, lower_kind, + with_comm=with_comm, ) else: _deserialize_to_file_pte( @@ -2587,6 +2594,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). @@ -2597,11 +2606,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) @@ -2630,6 +2640,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/doc/model/dpa4c.md b/doc/model/dpa4c.md index 7e79d8098e..a9f2865794 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..c7bc92f5ee 100644 --- a/doc/third-party/ase.md +++ b/doc/third-party/ase.md @@ -60,6 +60,38 @@ print(water.get_forces()) :::: +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) +``` + +This exports `model-ase_single.pt2` beside the checkpoint and loads it through +the compiled inference backend. Each calculator construction with `freeze=True` +rebuilds the artifact using the selected checkpoint head and the backend's +normal freeze policy. 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. 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/source/tests/pt_expt/test_calculator.py b/source/tests/pt_expt/test_calculator.py new file mode 100644 index 0000000000..6b22e3e079 --- /dev/null +++ b/source/tests/pt_expt/test_calculator.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""ASE checkpoint export and eager-inference behavior.""" + +import copy +import json +import logging +import os +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.pt.entrypoints import ( + freeze_pt2, +) +from deepmd.pt.utils import env as pt_env +from deepmd.pt_expt.utils import ( + env, + serialization, +) + +from ..pt.model.test_sezm_export import ( + _tiny_sezm_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.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")) + monkeypatch.setattr( + "deepmd.pt_expt.kernels.utils.backend_device_type", lambda: "cuda" + ) + 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"] + + 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) + + 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 + deep_pot.assert_called_with( + str(tmp_path / "model-ase_single.pt2"), + neighbor_list=None, + nlist_backend="auto", + ) + + +@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") + export = mock.Mock() + 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), + str(tmp_path / "model-ase_single.pt2"), + 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: + checkpoint = _write_checkpoint(tmp_path / "model.pt", kind) + 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) + 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) + 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 + assert atoms.calc.dp.deep_eval.model_path == str(frozen) + 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..d389e1f7a7 --- /dev/null +++ b/source/tests/pt_expt/test_freeze_compress.py @@ -0,0 +1,173 @@ +# 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.utils import ( + env, + serialization, +) + +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("compress", [False, True]) +def test_freeze_compression_changes_snapshot_before_export( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, compress: bool +) -> None: + config = _config("dpa4c") + 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() + 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) + 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)) + with ( + mock.patch("deepmd.calculator.DeepPot", return_value=loaded), + mock.patch.object(entrypoints, "freeze") 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_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 From b4b2996f0f3cdc0285da984002e2a32095ce874f Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 6 Oct 2026 09:35:53 +0800 Subject: [PATCH 2/5] perf(ase): reuse neighbor searches and prepare compact graphs directly ASE repeatedly evaluates one structure, but inference rebuilt neighbor-search state and passed graph data through redundant input adapters. Compressed DPA1/DPA4C kernels made that preparation a substantial fraction of each step. Expose the search controls at the calculator boundary and share persistent strategies across PT, PT-expt and their supported exported input schemas. Neighbor-search policy and cache ownership: - Forward neighbor_graph_method and neighbor_skin through DeepPot, including the reload after automatic checkpoint freezing. ASE defaults to a 1.0 angstrom skin; direct DeepEval callers retain a zero-skin default. - Prefer Vesin for single-frame CUDA graph inference and Toolkit-Ops for multi-frame CUDA batches. Preserve the native CPU graph search, explicit backend choices and the fallback when optional dependencies are absent. - Share the extended/edge-vector strategy factory between PT and PT-expt. Configure SeZM/DPA4 inference inside its geometry stage so native forward and compiled core execution remain intact. Training keeps its own policy. - Keep one Vesin calculator per frame and reuse search storage and Verlet candidates. Refresh geometry and enforce the exact model cutoff on every evaluation, rebuilding candidates when required by motion, cell or atom count changes. Copy search results to preserve earlier autograd graphs. - Cache ASE element-to-model type mappings and invalidate them when atomic numbers, atom order or type_dict change. Graph and input preparation: - Construct compressed DPA1/DPA4C canonical inputs directly from Vesin displacement vectors. Build destination-major payloads and both CSR views with uint32 source/order arrays, float32 geometry and the required guard storage for fewer than two physical edges. Avoid the generic graph conversion and its tensor-value reads on the host. - Use PyTorch destination sorting on CUDA instead of Vesin sorting. Preserve differentiable coordinate/cell geometry in generic graph construction, compute the virtual-endpoint selection once, and reuse the known destination ordering of NV output where exclusion semantics permit it. - Avoid per-edge cell replication and frame-offset work for one frame, and pass a single edge-vector frame through without concatenating its arrays. - Centralize graph input packing for ordinary and native-spin inference, preserving the distinct generic/canonical ABIs, optional fitting inputs and compressed charge-state handling. Keep generic conversion available for other neighbor builders and pair-excluded graphs. Correctness and compatibility: - Make local spin samples contiguous before PT edge-vector export so traced frame strides are independent of ghost storage and work for batched input. Regenerate affected native-spin edge-vector exports to use these strides. - Restrict the Vesin CUDA empty-cache workaround to the exact known pointer error with skin enabled. Accept an uncached retry only for an empty result; unrelated errors and nonempty retries still raise the original exception. - Stage PTE archive loading on CPU at both inference and serialization boundaries. PyTorch 2.13 constructs a storage wrapper without an explicit device; the scoped context prevents accidental default-device placement while preserving the tensor devices recorded in the archive. - Require both vesin and vesin-torch >= 0.6 for supported PyTorch targets. Keep older PyTorch targets and macOS x86_64 on native neighbor construction. - Document calculator controls, cache invalidation and supported input schemas in the public ASE guide. Validation: - 53 graph utility checks passed on CPU/CUDA, covering dual CSR, zero/one edges, empty systems, virtual atoms, multiple frames, periodic geometry, cache updates, output lifetime, gradients and empty/nonempty transitions. - 28 CPU inference checks and 13 CPU edge/neighbor checks passed; 15 CUDA-conditional cases were skipped in those CPU runs. Three compressed charge-state route checks passed on CUDA. - All 22 frame/atomic-parameter and spin-parameter checks passed after the archive-loading fix, resolving the 13 prior initialization errors. The new default-device regression passed on CPU and CUDA under both meta and invalid-CUDA defaults; both cases failed before the fix. - 120 real exported-model comparisons passed across ordinary/compressed Air, DPA1 canonical and DPA4C native-spin canonical inputs, including multiple frames, isolated/padded atoms and periodic/nonperiodic configurations. - All 24 paired MD initial-state comparisons passed. The maximum force difference was 1.074e-5 eV/angstrom and the maximum energy difference was 1.202e-7 eV/atom. Ruff and whitespace checks passed. Measured graph-preparation impact with public preset models: - Hardware: RTX PRO 6000 Blackwell Server Edition and Xeon Platinum 8558P; PyTorch 2.13.0+cu130, one CPU thread, identical frozen artifacts. - Full ASE NVE steps at 1 fs, 300 K initial velocities, 20 warmup steps and 100 timed steps per run; three alternating before/after repetitions. Compilation, loading and warmup are excluded. Shared-host load varied. - The reference already uses cached Vesin searches, so this comparison quantifies the graph/input preparation changes within the shared design. - Median milliseconds per step, 512 / 8000 atoms: preset Mini, ordinary: 9.16 -> 9.22 / 92.56 -> 92.50 preset Air, ordinary: 6.67 -> 5.81 / 36.50 -> 32.60 preset Air, compressed: 4.22 -> 2.61 / 10.04 -> 4.69 - Compressed Air improves by 1.62x / 2.14x. Mini shows no stable end-to-end improvement. For compressed Air at 8000 atoms, the measured ranges are 8.33-10.11 ms before and 3.58-5.06 ms after; these are workload-specific results under the observed shared-host load. Full ASE optimization and deployment scaling with public presets: - The before/after labels below compare the parent of this change with the full ASE optimization. The cached-search experiment above isolates graph preparation within the shared design. - Hardware: one RTX PRO 6000 Blackwell Server Edition GPU per run and Xeon Platinum 8558P CPUs; PyTorch 2.13.0+cu130, TF32 disabled and one MPI rank for LAMMPS. - Diamond carbon with identical coordinates and 300 K initial velocities; NVE with a 1 fs timestep, 20 warmup steps and 100 timed steps, repeated three times. Report median complete MD throughput in atoms/s, rounded to the nearest integer; export, compression, loading and warmup are excluded. - Checkpoint rows use PT for DPA4 and PT-expt for DPA4C. Freeze rows use uncompressed PT-expt .pt2 artifacts; freeze + compress rows use compressed artifacts. Each frozen artifact is shared across its frontend comparisons. - Ordinary LAMMPS uses the allocated CPU thread count. Kokkos GPU execution uses one LAMMPS host thread with the same CPU allocation and DeePMD intra-op limit. Kokkos rows use device MD/neighbor-list execution and the shared PT-expt frozen model. DPA4 preset Mini: 8 physical CPU cores per run, 512 / 3136 / 9680 atoms. Independent runs execute concurrently on separate GPUs and disjoint CPU cores local to each GPU. All variants in a repetition use the same GPU. Median throughput (atoms/s): Variant 512 3,136 9,680 ASE .pt before 13,188 17,505 18,214 ASE .pt after 13,480 18,500 18,479 ASE freeze before 56,141 77,661 80,307 ASE freeze after 65,248 85,011 87,357 LAMMPS freeze 76,718 87,101 87,988 LAMMPS Kokkos freeze 75,165 87,353 88,945 DPA4C preset Air: 24 physical CPU cores local to the GPU per run; 512 / 3136 / 9680 / 49248 atoms. All runs execute serially. Median throughput (atoms/s): Variant 512 3,136 9,680 49,248 ASE .pt before 29,867 93,962 116,377 118,342 ASE .pt after 32,049 110,497 122,495 120,628 ASE freeze before 83,740 175,846 230,227 239,501 ASE freeze after 117,170 220,846 252,397 249,470 ASE freeze + compress before 113,548 456,463 1,052,187 1,510,584 ASE freeze + compress after 253,795 1,268,573 2,199,040 2,698,941 LAMMPS freeze 156,523 198,797 203,997 180,545 LAMMPS Kokkos freeze 169,509 247,170 269,813 269,648 LAMMPS freeze + compress 310,375 794,813 927,621 753,016 LAMMPS Kokkos freeze + compress 585,428 2,836,777 4,239,738 6,439,693 All 54 Mini and 120 Air initial energy/force/stress comparisons passed. Across these matrices, the maximum initial force difference is 1.848e-5 eV/angstrom and the maximum energy difference is 5.886e-7 eV/atom. Model, exported-artifact and runtime-library hashes were verified. --- backend/find_pytorch.py | 19 +- deepmd/calculator.py | 50 +- .../dpmodel/utils/neighbor_graph/from_ijs.py | 11 +- deepmd/dpmodel/utils/neighbor_list.py | 5 +- deepmd/pt/entrypoints/freeze_pt2.py | 3 +- deepmd/pt/infer/deep_eval.py | 92 +-- deepmd/pt/model/model/sezm_model.py | 18 +- deepmd/pt_expt/infer/deep_eval.py | 777 ++++++------------ deepmd/pt_expt/utils/auto_neighbor_list.py | 89 ++ deepmd/pt_expt/utils/canonical_graph.py | 65 ++ deepmd/pt_expt/utils/edge_schema.py | 2 + deepmd/pt_expt/utils/graph_builder.py | 15 +- deepmd/pt_expt/utils/nv_graph_builder.py | 9 +- deepmd/pt_expt/utils/serialization.py | 5 +- deepmd/pt_expt/utils/vesin_graph_builder.py | 214 +++-- deepmd/pt_expt/utils/vesin_neighbor_list.py | 35 +- deepmd/pt_expt/utils/vesin_search.py | 144 ++++ doc/third-party/ase.md | 29 + source/tests/pt/model/test_nlist_backend.py | 35 +- source/tests/pt_expt/infer/test_deep_eval.py | 42 +- .../infer/test_deep_eval_pt_checkpoint.py | 37 +- source/tests/pt_expt/test_calculator.py | 124 ++- .../pt_expt/utils/test_canonical_graph.py | 57 ++ .../pt_expt/utils/test_vesin_graph_builder.py | 172 +++- 24 files changed, 1340 insertions(+), 709 deletions(-) create mode 100644 deepmd/pt_expt/utils/auto_neighbor_list.py create mode 100644 deepmd/pt_expt/utils/vesin_search.py 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 1075a57690..a5877eac3c 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -2,6 +2,7 @@ """ASE calculator interface module.""" import logging +import math from functools import ( partial, ) @@ -16,6 +17,7 @@ Union, ) +import numpy as np from ase.calculators.calculator import ( Calculator, PropertyNotImplementedError, @@ -135,6 +137,19 @@ 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. @@ -189,15 +204,23 @@ def __init__( 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(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) @@ -212,7 +235,11 @@ def __init__( log.info("Freezing ASE checkpoint %s to %s", model_path, output) freezer(str(model_path), str(output)) self.dp = DeepPot( - str(output), neighbor_list=neighbor_list, nlist_backend=nlist_backend + str(output), + neighbor_list=neighbor_list, + nlist_backend=nlist_backend, + neighbor_graph_method=neighbor_graph_method, + neighbor_skin=neighbor_skin, ) elif freezer is not None: log.warning( @@ -228,6 +255,9 @@ def __init__( 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, @@ -255,8 +285,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/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index 708faa55f6..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, diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index f468379565..97099ff33c 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,38 @@ 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 + 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 +349,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 +360,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..7b8b47aafa 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -719,6 +719,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 +2765,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, nf: int, device: torch.device) -> NeighborList: + if not self.training and self._inference_neighbor_list is not None: + return self._inference_neighbor_list + return _select_neighbor_builder(nf, device) + def build_neighbor_list( self, coord: Float[Tensor, "nf nloc 3"] | Float[Tensor, "nf nloc_x3"], @@ -2772,7 +2786,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(nf, coord.device).build( coord.view(nf, nloc, 3), atype, box, @@ -2799,7 +2813,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(nf, coord.device).build( coord.view(nf, nloc, 3), atype, box, diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 94c1e1f635..ad4b9254b8 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,48 @@ 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 = VesinNeighborList() if self._use_vesin else None + self._nlist_builder = ( + make_neighbor_list(nlist_backend, DEVICE, self._neighbor_skin) + 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 +504,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 +1673,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 +1731,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 +1831,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 +1876,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 +1900,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 +1997,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 +2074,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, - ) - - 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, - ) + """Evaluate native-spin graph or compact canonical inputs. - 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 +2112,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 +2227,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 +2281,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 +2294,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 +2402,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 +2426,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/utils/auto_neighbor_list.py b/deepmd/pt_expt/utils/auto_neighbor_list.py new file mode 100644 index 0000000000..eda85b6897 --- /dev/null +++ b/deepmd/pt_expt/utils/auto_neighbor_list.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Batch-aware selection of persistent torch neighbor-list strategies.""" + +from typing import ( + Any, +) + +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 AutoNeighborList(NeighborList): + """Prefer Vesin for one frame and Toolkit-Ops for CUDA batches.""" + + def __init__(self, skin: float = 0.0) -> None: + self._vesin = VesinNeighborList(skin=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() + if atype.shape[0] == 1 and vesin_available: + builder = self._vesin + elif coord.device.type == "cuda" 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 vesin[torch] or " + "nvalchemi-toolkit-ops; 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 +) -> NeighborList | None: + """Select an inference strategy, or native construction when auto has none.""" + if method == "native": + return None + if method == "auto": + if is_vesin_torch_available() or (device.type == "cuda" and is_nv_available()): + return AutoNeighborList(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 c1df31c6dd..923e2cf8c0 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1320,7 +1320,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 = ( 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..6471b40001 --- /dev/null +++ b/deepmd/pt_expt/utils/vesin_search.py @@ -0,0 +1,144 @@ +# 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._calculators: list[NeighborList] = [] + + 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._calculators.clear() + self._key = key + while len(self._calculators) <= frame: + self._calculators.append( + NeighborList( + cutoff=float(rcut), + full_list=True, + sorted=self.sorted, + skin=self.skin, + ) + ) + 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 = self._calculators[frame].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/third-party/ase.md b/doc/third-party/ase.md index c7bc92f5ee..9eef85f2ba 100644 --- a/doc/third-party/ase.md +++ b/doc/third-party/ase.md @@ -60,6 +60,35 @@ 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: diff --git a/source/tests/pt/model/test_nlist_backend.py b/source/tests/pt/model/test_nlist_backend.py index d38262dacd..8a04cc6a01 100644 --- a/source/tests/pt/model/test_nlist_backend.py +++ b/source/tests/pt/model/test_nlist_backend.py @@ -10,6 +10,9 @@ """ import copy +from unittest import ( + mock, +) import numpy as np import pytest @@ -24,6 +27,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, @@ -155,10 +161,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_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index 57fdd021ce..3655b42e93 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) @@ -2343,11 +2360,13 @@ def _vesin_available(self) -> bool: def test_default_is_auto(self) -> None: 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() + ) 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 +2395,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 00eab95d62..9efad937a9 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,38 @@ 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"), + ) + for device, nf, nv_available, vesin_available, expected in cases: + with self.subTest(device=device, nf=nf, expected=expected): + builder = auto_neighbor_list.AutoNeighborList(skin=1.0) + 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, + ), + 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_explicit_resolution(self) -> None: """Every implemented graph builder is accepted without rewriting.""" for method in ("dense", "ase", "cell", "vesin", "nv"): @@ -450,7 +482,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/test_calculator.py b/source/tests/pt_expt/test_calculator.py index 6b22e3e079..f55b1d84dd 100644 --- a/source/tests/pt_expt/test_calculator.py +++ b/source/tests/pt_expt/test_calculator.py @@ -37,6 +37,7 @@ from ..pt.model.test_sezm_export import ( _tiny_sezm_model_params, + _tiny_sezm_native_spin_model_params, ) from .model.test_dpa4_export import ( _DPA4_CONFIG, @@ -150,6 +151,119 @@ def _write_checkpoint( return path +@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 = [ + native.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 @@ -212,7 +326,13 @@ def load(model_file: str, **kwargs: object) -> object: return result with mock.patch("deepmd.calculator.DeepPot", side_effect=load) as deep_pot: - DP(checkpoint, freeze=True, head=head) + DP( + checkpoint, + freeze=True, + head=head, + neighbor_graph_method="vesin", + neighbor_skin=0.75, + ) expected = override or ("1", "2") if kind == "dpa4c": @@ -226,6 +346,8 @@ def load(model_file: str, **kwargs: object) -> object: str(tmp_path / "model-ase_single.pt2"), neighbor_list=None, nlist_backend="auto", + neighbor_graph_method="vesin", + neighbor_skin=0.75, ) 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_vesin_graph_builder.py b/source/tests/pt_expt/utils/test_vesin_graph_builder.py index c216676994..22ed5f4ed8 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,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from dataclasses import ( + replace, +) + import numpy as np import pytest import torch @@ -10,11 +14,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 +124,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 +183,158 @@ 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("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) From 946fbae02d77d04e547846850c5a76216c39c272 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 18 Oct 2026 13:51:02 +0800 Subject: [PATCH 3/5] test(ase): isolate export fixtures and device policy mocks Share tiny SeZM configurations without importing another backend test package and its process-global thread/device initialization. Limit synthetic CUDA detection to entering the real export policy context so lazy canonical imports retain the actual backend query and subsequent CPU exports remain independent. --- source/tests/common/sezm_config.py | 59 +++++++++++++++++ source/tests/pt/model/test_sezm_export.py | 63 ++----------------- .../tests/pt/model/test_sezm_vacuum_freeze.py | 4 +- source/tests/pt_expt/test_calculator.py | 20 +++++- 4 files changed, 85 insertions(+), 61 deletions(-) create mode 100644 source/tests/common/sezm_config.py 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/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_expt/test_calculator.py b/source/tests/pt_expt/test_calculator.py index f55b1d84dd..dfb6ff273a 100644 --- a/source/tests/pt_expt/test_calculator.py +++ b/source/tests/pt_expt/test_calculator.py @@ -1,11 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """ASE checkpoint export and eager-inference behavior.""" +import contextlib import copy import json import logging import os import zipfile +from collections.abc import ( + Iterator, +) from pathlib import ( Path, ) @@ -35,7 +39,7 @@ serialization, ) -from ..pt.model.test_sezm_export import ( +from ..common.sezm_config import ( _tiny_sezm_model_params, _tiny_sezm_native_spin_model_params, ) @@ -300,8 +304,20 @@ def test_automatic_freeze_reuses_kernel_policy( ) -> 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( - "deepmd.pt_expt.kernels.utils.backend_device_type", lambda: "cuda" + serialization, "_fused_operators_for_export", cuda_export_policy ) if override is not None: for name, value in zip(_LEVELS[:2], override, strict=True): From 5885aab7589d0f569b00c453be3d212c0b2a0fc5 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 19 Oct 2026 17:54:02 +0800 Subject: [PATCH 4/5] fix(ase): isolate frozen artifacts and resolve review regressions Keep automatic exports in private evaluator-owned directories and publish separate complete copies atomically. Preserve the previous public artifact on failure, selected heads across concurrent processes, and later evaluator serialization after another export or public-file overwrite. Restrict PyTorch 2.11 CPU inference horizontal fusion to avoid a gather reusing a scatter index array before its declaration. Keep SIMD and vertical fusion enabled and retain real float64 export and numerical checks. Reject freeze --compress for unsupported backends, provide an actionable compression command when neighbor statistics are absent, and align the existing dispatch and nlist-message tests with their intended behavior. Validate artifact lifetime, deterministic concurrent publication, failure cleanup, CLI handling, compiler options, and actual CPU freeze/reload parity. Validation: PyTorch 2.13 CPU related suites 161 passed, 15 skipped; PyTorch 2.11 CPU AVX2 freeze/compression and compiler-option suites 23 passed. Scoped pre-commit hooks, repository Ruff, and whitespace checks pass. --- deepmd/calculator.py | 46 ++++-- deepmd/main.py | 6 + deepmd/pt/utils/compile_compat.py | 5 + deepmd/pt_expt/entrypoints/compress.py | 7 +- doc/third-party/ase.md | 21 ++- source/tests/pt/test_compile_compat.py | 4 + .../model/test_graph_builder_dispatch.py | 2 +- source/tests/pt_expt/test_calculator.py | 142 +++++++++++++++++- source/tests/pt_expt/test_freeze_compress.py | 29 +++- .../pt_expt/utils/test_graph_pt2_metadata.py | 2 +- 10 files changed, 237 insertions(+), 27 deletions(-) diff --git a/deepmd/calculator.py b/deepmd/calculator.py index a5877eac3c..402e09ba77 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -3,12 +3,22 @@ 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, @@ -154,8 +164,11 @@ class DP(Calculator): Export an eligible ``.pt`` checkpoint to ``.pt2`` for inference, by default False. Supports PT DPA4/SeZM and graph-eligible PT-expt models. Each call - exports ``-ase_single.pt2`` beside the checkpoint, - using the backend's freeze-time kernel policy and omitting the + 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 @@ -233,14 +246,27 @@ def __init__( self.dp.close() del self.dp log.info("Freezing ASE checkpoint %s to %s", model_path, output) - freezer(str(model_path), str(output)) - self.dp = DeepPot( - str(output), - neighbor_list=neighbor_list, - nlist_backend=nlist_backend, - neighbor_graph_method=neighbor_graph_method, - neighbor_skin=neighbor_skin, - ) + 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 " diff --git a/deepmd/main.py b/deepmd/main.py index 137bd8ac27..6d967553fc 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -1105,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/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 d70a9b5210..00a5cc219c 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -116,9 +116,10 @@ def compress_model( ) 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, diff --git a/doc/third-party/ase.md b/doc/third-party/ase.md index 9eef85f2ba..498b41c4f7 100644 --- a/doc/third-party/ase.md +++ b/doc/third-party/ase.md @@ -98,10 +98,16 @@ from deepmd.calculator import DP calculator = DP(model="model.pt", freeze=True) ``` -This exports `model-ase_single.pt2` beside the checkpoint and loads it through -the compiled inference backend. Each calculator construction with `freeze=True` -rebuilds the artifact using the selected checkpoint head and the backend's -normal freeze policy. The export omits the multi-rank communication artifact +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. @@ -117,9 +123,10 @@ 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. Already frozen -models, other PT models, and PT-expt models without graph eligibility ignore -`freeze` and keep their normal loading path. +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: 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_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 index dfb6ff273a..3f0d9ab03d 100644 --- a/source/tests/pt_expt/test_calculator.py +++ b/source/tests/pt_expt/test_calculator.py @@ -3,16 +3,27 @@ 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, ) @@ -24,6 +35,7 @@ Atoms, ) +from deepmd import calculator as calculator_module from deepmd.calculator import ( DP, ) @@ -155,6 +167,117 @@ def _write_checkpoint( 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 @@ -327,6 +450,7 @@ def cuda_export_policy() -> Iterator[None]: 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) @@ -358,8 +482,12 @@ def load(model_file: str, **kwargs: object) -> object: 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(tmp_path / "model-ase_single.pt2"), + str(private), neighbor_list=None, nlist_backend="auto", neighbor_graph_method="vesin", @@ -386,13 +514,17 @@ def test_single_task_pt_freeze_ignores_head( """Single-task head arguments retain the eager checkpoint's semantics.""" checkpoint = _write_checkpoint(tmp_path / "model.pt", "pt") loaded = DeepPot(str(checkpoint), head="Default") - export = mock.Mock() + + 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), - str(tmp_path / "model-ase_single.pt2"), + mock.ANY, head=None, with_comm=False, ) @@ -429,7 +561,9 @@ def test_automatic_freeze_single_artifact_matches_eager( 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 - assert atoms.calc.dp.deep_eval.model_path == str(frozen) + 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 diff --git a/source/tests/pt_expt/test_freeze_compress.py b/source/tests/pt_expt/test_freeze_compress.py index d389e1f7a7..a01567715b 100644 --- a/source/tests/pt_expt/test_freeze_compress.py +++ b/source/tests/pt_expt/test_freeze_compress.py @@ -60,6 +60,29 @@ def test_freeze_cli_compression_option(tmp_path: Path, compress: bool) -> None: 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 @@ -108,9 +131,13 @@ def test_ase_falls_back_but_explicit_compression_reports_ineligibility( 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") as freeze, + mock.patch.object(entrypoints, "freeze", side_effect=write_artifact) as freeze, caplog.at_level(logging.WARNING, logger="deepmd.calculator"), ): DP(checkpoint, freeze=True) 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..f82102c020 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -508,7 +508,7 @@ 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) From 7d57390482706e949dbbe4700e080fce5c2e4c47 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 21 Oct 2026 15:50:11 +0800 Subject: [PATCH 5/5] fix(ase): invalidate cached neighbors when the cell changes Retain each frame's calculator with an independent cell snapshot. Rebuild the frame's candidate search on cell or periodicity changes so cell contraction cannot introduce periodic images absent from the old candidates. Keep fixed-cell reuse and avoid cell comparisons when the skin is zero. Cover multi-image contraction for lattice shifts and displacement vectors, in-place cell updates, frame isolation, fixed-cell reuse, and PBC transitions. Correct the default neighbor-builder expectation for CUDA/NV without Vesin and exercise the availability policy through the real strategy factory. Validation: 87 CPU tests and 38 subtests passed; 26 conditional tests skipped. The two contraction regressions fail before the fix and pass afterward. Scoped pre-commit hooks, repository Ruff, and whitespace checks pass. --- deepmd/pt_expt/utils/vesin_search.py | 37 ++++++++++----- source/tests/pt_expt/infer/test_deep_eval.py | 7 ++- .../infer/test_deep_eval_pt_checkpoint.py | 21 +++++++-- .../pt_expt/utils/test_vesin_graph_builder.py | 45 +++++++++++++++++++ 4 files changed, 95 insertions(+), 15 deletions(-) diff --git a/deepmd/pt_expt/utils/vesin_search.py b/deepmd/pt_expt/utils/vesin_search.py index 6471b40001..04ff5f31dc 100644 --- a/deepmd/pt_expt/utils/vesin_search.py +++ b/deepmd/pt_expt/utils/vesin_search.py @@ -43,7 +43,7 @@ def __init__(self, skin: float = 0.0, *, sorted: bool = False) -> None: self.skin = float(skin) self.sorted = sorted self._key: tuple[float, torch.device] | None = None - self._calculators: list[NeighborList] = [] + self._frames: dict[int, tuple[NeighborList, torch.Tensor | None]] = {} def compute( self, @@ -93,17 +93,34 @@ def compute( device = positions.device key = (float(rcut), device) if self._key != key: - self._calculators.clear() + self._frames.clear() self._key = key - while len(self._calculators) <= frame: - self._calculators.append( - NeighborList( - cutoff=float(rcut), - full_list=True, - sorted=self.sorted, - skin=self.skin, + 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 @@ -111,7 +128,7 @@ def compute( ) with torch.device(device): try: - ii, jj, geometry = self._calculators[frame].compute( + ii, jj, geometry = calculator.compute( points=positions.detach(), box=box.detach(), periodic=cell is not None, diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index 3655b42e93..27b2bf7c2d 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -2359,9 +2359,14 @@ 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._nlist_builder is not None, self._vesin_available() + 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: 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 9efad937a9..92f7ea39a2 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 @@ -433,10 +433,13 @@ def test_auto_nlist_strategy_uses_frame_count(self) -> None: ("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): - builder = auto_neighbor_list.AutoNeighborList(skin=1.0) coord = mock.Mock(device=torch.device(device)) atype = np.zeros((nf, 2), dtype=np.int64) with ( @@ -448,10 +451,20 @@ def test_auto_nlist_strategy_uses_frame_count(self) -> None: "is_vesin_torch_available", return_value=vesin_available, ), - 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]) + 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_explicit_resolution(self) -> None: 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 22ed5f4ed8..6777f71d45 100644 --- a/source/tests/pt_expt/utils/test_vesin_graph_builder.py +++ b/source/tests/pt_expt/utils/test_vesin_graph_builder.py @@ -2,6 +2,9 @@ from dataclasses import ( replace, ) +from unittest import ( + mock, +) import numpy as np import pytest @@ -261,6 +264,48 @@ def test_cached_search_keeps_frames_independent() -> None: 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])