From abe2e033a20af33c4bee341af7d59e418b69be64 Mon Sep 17 00:00:00 2001 From: shaurya2k06 Date: Thu, 6 Aug 2026 11:44:55 +0530 Subject: [PATCH] feat(pt_expt): call-time DeepEval auto ladder with nf==1 vesin gate Extract resolve_auto_graph_builder for inference and select vesin only for single-frame batches, matching _select_neighbor_builder. Multi-frame auto stays on nv/dense so auto_batch_size does not hit the per-frame loop. Signed-off-by: shaurya2k06 --- deepmd/pt_expt/infer/deep_eval.py | 64 ++++++++------- deepmd/pt_expt/utils/graph_builder.py | 79 +++++++++++++++++++ deepmd/pt_expt/utils/vesin_graph_builder.py | 10 ++- .../infer/test_deep_eval_pt_checkpoint.py | 37 +++++++-- .../model/test_graph_builder_dispatch.py | 69 +++++++++++++--- 5 files changed, 207 insertions(+), 52 deletions(-) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 9d6396d990..b266159c59 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,15 +189,17 @@ 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` + at each eval call (CUDA: ``nv`` if importable; else ``vesin`` only when + ``nf == 1`` and 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 + 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. **kwargs : dict Keyword arguments. @@ -271,8 +270,13 @@ def __init__( raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize") @staticmethod - def _resolve_neighbor_graph_method(method: str) -> str: - """Resolve the graph builder once for the active device.""" + def _resolve_neighbor_graph_method(method: str, nf: int | None = None) -> str: + """Validate and optionally resolve the graph builder for the active device. + + ``"auto"`` is left unresolved when ``nf`` is omitted so construction- + time setup can defer to :meth:`_build_eval_graph`, where the frame + count is known and vesin can be gated on ``nf == 1``. + """ if method not in ("auto", "dense", "ase", "vesin", "nv"): raise ValueError( f"Unknown neighbor_graph_method {method!r}; " @@ -280,24 +284,17 @@ def _resolve_neighbor_graph_method(method: str) -> str: ) if method != "auto": return method + if nf is None: + return "auto" - 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, nf) def _setup_neighbor_backend(self, nlist_backend: str) -> None: """Resolve the graph or neighbor-list construction strategy. @@ -2316,14 +2313,21 @@ def _build_eval_graph( ) -> "NeighborGraph": """Build the carry-all NeighborGraph for graph-lower inference. - Dispatches on ``self._neighbor_graph_method``: ``dense``/``ase`` run - backend-agnostic (numpy); ``vesin``/``nv`` run on-device (torch, O(N)). - All backends emit the SAME neighbor set (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. + Dispatches on ``self._neighbor_graph_method``: ``auto`` is resolved + call-time via + :func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder` + using the batch frame count (vesin only when ``nf == 1``); + ``dense``/``ase`` run backend-agnostic (numpy); ``vesin``/``nv`` run + on-device (torch, O(N)). All backends emit the SAME neighbor set + (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. """ method = self._neighbor_graph_method + if method == "auto": + coord_arr = np.asarray(coord_input) + nf = int(coord_arr.shape[0]) if coord_arr.ndim >= 2 else 1 + method = self._resolve_neighbor_graph_method("auto", nf=nf) # Model-level ``pair_exclude_types`` is a graph-BUILD transform # (decision #18): apply it here so the exported ``.pt2`` lower consumes a # pre-excluded ``edge_mask`` and never re-applies it (mirrors the C++ @@ -2392,7 +2396,7 @@ def _build_eval_graph( ) raise ValueError( f"unknown neighbor_graph_method {method!r}; " - "use 'dense', 'ase', 'vesin', or 'nv'" + "use 'auto', 'dense', 'ase', 'vesin', or 'nv'" ) def _model_pair_excl(self) -> "PairExcludeMask | None": diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index b074a3af8f..9e845a19ed 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -19,6 +19,83 @@ log = logging.getLogger(__name__) +def resolve_auto_graph_builder( + device: torch.device | str, + nf: int = 1, +) -> 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 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. + + Policy + ------ + * CUDA + ``nvalchemiops``: ``nv`` (any ``nf``). + * ``nf == 1`` + ``vesin.torch``: ``vesin``. + * otherwise: ``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. + + Parameters + ---------- + device : torch.device or str + Device the coordinates live on (or will be moved to). Controls whether + the CUDA-only ``nv`` builder is eligible. + nf : int, default: 1 + Number of frames in the batch. ``vesin`` is selected only when + ``nf == 1`` and ``vesin.torch`` is importable. + + Returns + ------- + str + One of ``"nv"``, ``"vesin"``, or ``"dense"``. + + Raises + ------ + ValueError + If ``nf`` is not a positive ``int`` (``bool`` is rejected). + """ + from deepmd.pt.utils.nv_nlist import ( + is_nv_available, + ) + from deepmd.pt_expt.utils.vesin_neighbor_list import ( + is_vesin_torch_available, + ) + + # ``bool`` is a subclass of ``int``; reject it explicitly. + if type(nf) is not int: + raise ValueError(f"nf must be a positive int, got {nf!r}") + if nf < 1: + raise ValueError(f"nf must be >= 1, got {nf}") + + dev = torch.device(device) + if dev.type == "cuda" and is_nv_available(): + return "nv" + if nf == 1 and is_vesin_torch_available(): + return "vesin" + if dev.type == "cuda" and not is_nv_available(): + log.warning( + "nvalchemi-toolkit-ops is unavailable; falling back from " + "neighbor_graph_method='auto' to the dense graph builder" + + ( + "" + if nf == 1 + else " (vesin is not used for nf>1; its API loops frames in Python)" + ) + + ". Install it with `pip install nvalchemi-toolkit-ops` to enable " + "the NV graph builder." + ) + return "dense" + + def resolve_neighbor_graph_method( requested: str, device: torch.device, @@ -36,6 +113,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..874e979449 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 ``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 +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..28efc012d7 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 @@ -422,16 +422,31 @@ def test_unsupported_extension_raises(self) -> None: class TestNeighborGraphMethodResolution(unittest.TestCase): """Auto graph-builder selection must cover each host policy explicitly.""" + def test_auto_deferred_until_nf_known(self) -> None: + """Construction-time resolve leaves ``auto`` unresolved without ``nf``.""" + self.assertEqual( + PtExptDeepEval._resolve_neighbor_graph_method("auto"), + "auto", + ) + def test_auto_resolution(self) -> None: + # (device, nv, vesin, nf, expected, warns) cases = ( - ("cpu", False, "dense", False), - ("cuda", True, "nv", False), - ("cuda", False, "dense", True), - ) - for device_type, nv_available, expected, warns in cases: + ("cpu", False, True, 1, "vesin", False), + ("cpu", False, True, 4, "dense", False), + ("cpu", False, False, 1, "dense", False), + ("cuda", True, True, 1, "nv", False), + ("cuda", True, True, 4, "nv", False), + ("cuda", False, True, 1, "vesin", False), + ("cuda", False, True, 4, "dense", True), + ("cuda", False, False, 1, "dense", True), + ) + for device_type, nv_available, vesin_available, nf, expected, warns in cases: with self.subTest( device_type=device_type, nv_available=nv_available, + vesin_available=vesin_available, + nf=nf, ): with ( mock.patch( @@ -442,17 +457,23 @@ 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( - "auto" + "auto", nf=nf ) else: - actual = PtExptDeepEval._resolve_neighbor_graph_method("auto") + actual = PtExptDeepEval._resolve_neighbor_graph_method( + "auto", nf=nf + ) self.assertEqual(actual, expected) 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..7095f1ff66 100644 --- a/source/tests/pt_expt/model/test_graph_builder_dispatch.py +++ b/source/tests/pt_expt/model/test_graph_builder_dispatch.py @@ -67,13 +67,20 @@ def _make_model(): return EnergyModel(ds, ft, type_map=["O", "H"]).to(env.DEVICE) -def _eval(model, method): +def _eval(model, method, nf: int = 1): rng = np.random.default_rng(0) coord = torch.tensor( - rng.random((1, 6, 3)) * 4.0, dtype=torch.float64, device=env.DEVICE + rng.random((nf, 6, 3)) * 4.0, dtype=torch.float64, device=env.DEVICE + ) + atype = torch.tensor( + [[0, 1, 1, 0, 1, 1]] * nf, dtype=torch.int64, device=env.DEVICE + ) + box = ( + (torch.eye(3, dtype=torch.float64, device=env.DEVICE) * 6.0) + .reshape(1, 3, 3) + .expand(nf, 3, 3) + .clone() ) - atype = torch.tensor([[0, 1, 1, 0, 1, 1]], dtype=torch.int64, device=env.DEVICE) - box = (torch.eye(3, dtype=torch.float64, device=env.DEVICE) * 6.0).reshape(1, 3, 3) ret = model.forward_common(coord, atype, box, neighbor_graph_method=method) # graph path returns the output-agnostic dict (no translated force/virial); # energy_redu = total energy, energy_derv_r = d energy / d coord (force parity) @@ -138,12 +145,53 @@ def test_explicit_nv_rejects_cpu(): resolve_neighbor_graph_method("nv", torch.device("cpu")) +@pytest.mark.parametrize( + ("device", "nv", "vesin", "nf", "expected"), + [ + ("cpu", False, True, 1, "vesin"), + ("cpu", False, True, 4, "dense"), + ("cpu", True, False, 1, "dense"), + ("cuda", True, True, 1, "nv"), + ("cuda", True, True, 4, "nv"), + ("cuda", False, True, 1, "vesin"), + ("cuda", False, True, 4, "dense"), + ("cuda", False, False, 1, "dense"), + ], +) +def test_resolve_auto_graph_builder_ladder( + device: str, nv: bool, vesin: bool, nf: int, 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, nf=nf) == expected + + +@pytest.mark.parametrize("nf", [None, 1.5, True, False, 0, -1]) +def test_resolve_auto_graph_builder_rejects_invalid_nf(nf) -> None: + from deepmd.pt_expt.utils.graph_builder import ( + resolve_auto_graph_builder, + ) + + with pytest.raises(ValueError, match="nf must"): + resolve_auto_graph_builder("cpu", nf=nf) + + @pytest.mark.skipif(not is_vesin_torch_available(), reason="vesin[torch] not installed") -def test_vesin_matches_dense_energy_force(): +@pytest.mark.parametrize("nf", [1, 4]) +def test_vesin_matches_dense_energy_force(nf: int): torch.manual_seed(0) model = _make_model() - e_d, f_d = _eval(model, "dense") - e_v, f_v = _eval(model, "vesin") + e_d, f_d = _eval(model, "dense", nf=nf) + e_v, f_v = _eval(model, "vesin", nf=nf) tol = 1e-12 if env.DEVICE.type == "cpu" else 1e-10 torch.testing.assert_close(e_v, e_d, rtol=tol, atol=tol) torch.testing.assert_close(f_v, f_d, rtol=tol, atol=tol) @@ -153,11 +201,12 @@ def test_vesin_matches_dense_energy_force(): not (torch.cuda.is_available() and is_nv_available()), reason="nvalchemiops requires CUDA + nvalchemi-toolkit-ops", ) -def test_nv_matches_dense_energy_force(): +@pytest.mark.parametrize("nf", [1, 4]) +def test_nv_matches_dense_energy_force(nf: int): torch.manual_seed(0) model = _make_model() - e_d, f_d = _eval(model, "dense") - e_n, f_n = _eval(model, "nv") + e_d, f_d = _eval(model, "dense", nf=nf) + e_n, f_n = _eval(model, "nv", nf=nf) tol = 1e-10 # CUDA fp64: absorbs scatter-atomic / index_add nondeterminism torch.testing.assert_close(e_n, e_d, rtol=tol, atol=tol) torch.testing.assert_close(f_n, f_d, rtol=tol, atol=tol)