From 01253d981038d81c8bd464e02a2388a4c705d2db Mon Sep 17 00:00:00 2001 From: shaurya2k06 Date: Sat, 1 Aug 2026 14:25:02 +0530 Subject: [PATCH] feat(pt_expt): include vesin in DeepEval auto graph-builder ladder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Training keeps CPU on dense (vesin loops frames). Inference auto now shares resolve_auto_graph_builder: CUDA nv→vesin→dense, CPU vesin→dense. Signed-off-by: shaurya2k06 --- deepmd/pt_expt/infer/deep_eval.py | 35 ++++++--------- deepmd/pt_expt/utils/graph_builder.py | 43 +++++++++++++++++++ deepmd/pt_expt/utils/vesin_graph_builder.py | 10 +++-- .../infer/test_deep_eval_pt_checkpoint.py | 18 +++++--- .../model/test_graph_builder_dispatch.py | 27 ++++++++++++ 5 files changed, 102 insertions(+), 31 deletions(-) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index b8b85d2a4b..af1cc65c9f 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json -import logging import warnings from collections.abc import ( Callable, @@ -80,8 +79,6 @@ NeighborGraph, ) -log = logging.getLogger(__name__) - # Public output keys emitted by graph-lower forwards, keyed by the # output-variable category that ``request_defs`` carries. The graph path is @@ -192,16 +189,19 @@ class DeepEval(DeepEvalBackend): neighbor_graph_method : str, default: "auto" Carry-all graph builder for graph-form ``.pt2`` artifacts and graph-routed ``.pt`` checkpoints - (``metadata["lower_input_kind"] == "graph"``): ``"auto"`` selects - ``"nv"`` on CUDA when nvalchemiops is available and otherwise falls - back to ``"dense"``. ``"vesin"`` remains explicit opt-in because it - loops over frames in Python. Explicit + (``metadata["lower_input_kind"] == "graph"``): ``"auto"`` selects via + :func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder` + (CUDA: ``nv`` if importable, else ``vesin`` if importable, else + ``dense``; CPU: ``vesin`` if importable, else ``dense``). Explicit ``"dense"`` / ``"ase"`` / ``"vesin"`` / ``"nv"`` choices are preserved. A non-default value on any other artifact raises at construction because the knob would silently do nothing there; use ``nlist_backend`` for the nlist path instead. All builders emit the same neighbor set, so the - choice is performance-only. Consolidating the two knobs into a single - backend-selection API is deferred to the dense-nlist deprecation. + choice is performance-only. Training keeps a separate CPU-dense auto + policy (:func:`~deepmd.pt_expt.utils.graph_builder.resolve_neighbor_graph_method`) + because vesin's per-frame Python loop is not a multi-frame training + default. Consolidating the two knobs into a single backend-selection API + is deferred to the dense-nlist deprecation. **kwargs : dict Keyword arguments. """ @@ -281,23 +281,14 @@ def _resolve_neighbor_graph_method(method: str) -> str: if method != "auto": return method - from deepmd.pt.utils.nv_nlist import ( - is_nv_available, - ) from deepmd.pt_expt.utils.env import ( DEVICE, ) + from deepmd.pt_expt.utils.graph_builder import ( + resolve_auto_graph_builder, + ) - if DEVICE.type == "cuda": - if is_nv_available(): - return "nv" - log.warning( - "nvalchemi-toolkit-ops is unavailable; falling back from " - "neighbor_graph_method='auto' to the dense graph builder. " - "Install it with `pip install nvalchemi-toolkit-ops` to enable " - "the NV graph builder." - ) - return "dense" + return resolve_auto_graph_builder(DEVICE) def _setup_neighbor_backend(self, nlist_backend: str) -> None: """Resolve the graph or neighbor-list construction strategy. diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index b074a3af8f..d7a373f8fd 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -19,6 +19,47 @@ log = logging.getLogger(__name__) +def resolve_auto_graph_builder( + device: torch.device | str, +) -> str: + """Resolve ``neighbor_graph_method="auto"`` to a concrete inference builder. + + Single owner of the inference / DeepEval auto ladder (training uses + :func:`resolve_neighbor_graph_method`, which keeps CPU on ``dense`` because + vesin loops frames in Python and is not safe as a multi-frame training + default): + + * CUDA: ``nv`` if ``nvalchemiops`` is importable, else ``vesin`` if + ``vesin.torch`` is importable, else ``dense``. + * CPU: ``vesin`` if ``vesin.torch`` is importable, else ``dense``. + + ``ase`` is never chosen automatically. All builders emit the same carry-all + neighbor set; the choice is performance-only. Builders run eagerly outside + traced / compiled regions, so this does not change ``.pt2`` artifacts. + """ + from deepmd.pt.utils.nv_nlist import ( + is_nv_available, + ) + from deepmd.pt_expt.utils.vesin_neighbor_list import ( + is_vesin_torch_available, + ) + + dev = torch.device(device) + if dev.type == "cuda": + if is_nv_available(): + return "nv" + if is_vesin_torch_available(): + return "vesin" + log.warning( + "nvalchemi-toolkit-ops and vesin[torch] are unavailable; falling " + "back from neighbor_graph_method='auto' to the dense graph builder." + ) + return "dense" + if is_vesin_torch_available(): + return "vesin" + return "dense" + + def resolve_neighbor_graph_method( requested: str, device: torch.device, @@ -36,6 +77,8 @@ def resolve_neighbor_graph_method( ------- str The concrete builder name, either ``"dense"`` or ``"nv"``. + Training auto never selects ``vesin`` (per-frame Python loop); use + :func:`resolve_auto_graph_builder` for inference auto selection. Raises ------ diff --git a/deepmd/pt_expt/utils/vesin_graph_builder.py b/deepmd/pt_expt/utils/vesin_graph_builder.py index a715189ae5..7e88f86a06 100644 --- a/deepmd/pt_expt/utils/vesin_graph_builder.py +++ b/deepmd/pt_expt/utils/vesin_graph_builder.py @@ -8,10 +8,12 @@ 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. It is never on a default hot path: -``neighbor_graph_method=None`` resolves to the ``"dense"`` converter, and -vesin is explicit opt-in only. For batched multi-frame GPU work prefer -``nv`` (:mod:`.nv_graph_builder`), which batches all frames in one kernel. +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 ``vesin.torch`` is importable (CPU always; CUDA only when ``nv`` +is unavailable); otherwise it falls back to ``dense``. Training auto keeps CPU +on ``dense`` and never selects vesin. Prefer ``nv`` (:mod:`.nv_graph_builder`) +for batched multi-frame GPU work, which batches all frames in one kernel. """ from __future__ import ( 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 fb1f86afa3..e8b04a1bbc 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 @@ -423,15 +423,19 @@ class TestNeighborGraphMethodResolution(unittest.TestCase): """Auto graph-builder selection must cover each host policy explicitly.""" def test_auto_resolution(self) -> None: + # (device, nv, vesin, expected, warns) cases = ( - ("cpu", False, "dense", False), - ("cuda", True, "nv", False), - ("cuda", False, "dense", True), + ("cpu", False, True, "vesin", False), + ("cpu", False, False, "dense", False), + ("cuda", True, True, "nv", False), + ("cuda", False, True, "vesin", False), + ("cuda", False, False, "dense", True), ) - for device_type, nv_available, expected, warns in cases: + for device_type, nv_available, vesin_available, expected, warns in cases: with self.subTest( device_type=device_type, nv_available=nv_available, + vesin_available=vesin_available, ): with ( mock.patch( @@ -442,10 +446,14 @@ def test_auto_resolution(self) -> None: "deepmd.pt.utils.nv_nlist.is_nv_available", return_value=nv_available, ), + mock.patch( + "deepmd.pt_expt.utils.vesin_neighbor_list.is_vesin_torch_available", + return_value=vesin_available, + ), ): if warns: with self.assertLogs( - "deepmd.pt_expt.infer.deep_eval", + "deepmd.pt_expt.utils.graph_builder", level="WARNING", ): actual = PtExptDeepEval._resolve_neighbor_graph_method( 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 b4a7b903d6..510b62d8f5 100644 --- a/source/tests/pt_expt/model/test_graph_builder_dispatch.py +++ b/source/tests/pt_expt/model/test_graph_builder_dispatch.py @@ -138,6 +138,33 @@ def test_explicit_nv_rejects_cpu(): resolve_neighbor_graph_method("nv", torch.device("cpu")) +@pytest.mark.parametrize( + ("device", "nv", "vesin", "expected"), + [ + ("cpu", False, True, "vesin"), + ("cpu", True, False, "dense"), + ("cuda", True, True, "nv"), + ("cuda", False, True, "vesin"), + ("cuda", False, False, "dense"), + ], +) +def test_resolve_auto_graph_builder_ladder( + device: str, nv: bool, vesin: bool, expected: str +) -> None: + from deepmd.pt_expt.utils.graph_builder import ( + resolve_auto_graph_builder, + ) + + with ( + patch("deepmd.pt.utils.nv_nlist.is_nv_available", return_value=nv), + patch( + "deepmd.pt_expt.utils.vesin_neighbor_list.is_vesin_torch_available", + return_value=vesin, + ), + ): + assert resolve_auto_graph_builder(device) == expected + + @pytest.mark.skipif(not is_vesin_torch_available(), reason="vesin[torch] not installed") def test_vesin_matches_dense_energy_force(): torch.manual_seed(0)