Skip to content

Commit a1dd9eb

Browse files
feat(pt_expt): call-time DeepEval auto ladder with nf==1 vesin gate (deepmodeling#5903)
## Summary - Add shared `resolve_auto_graph_builder(device, nf)` for **inference / DeepEval only** (training keeps `resolve_neighbor_graph_method` from deepmodeling#5913). - Resolve `neighbor_graph_method="auto"` **at eval call time** with the batch frame count: CUDA prefers `nv`; `vesin` only when `nf == 1` and importable; otherwise `dense`. Matches `_select_neighbor_builder`. - Multi-frame `auto_batch_size` / `dp test` batches therefore stay off vesin's per-frame Python loop. - Parametrize vesin (and nv) vs dense energy/force parity over `nf in {1, 4}`. This is **not** a model-level / training default flip — those already shipped in deepmodeling#5912 / deepmodeling#5913. The remaining change is the inference auto ladder: re-introducing vesin only under the `nf == 1` gate that review asked for on deepmodeling#5912. ## Validation - `ruff check` / `ruff format` on touched files - `pytest` resolver ladder + DeepEval resolution + vesin parity for `nf=1` and `nf=4` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Improved automatic neighbor-graph builder selection during inference. - CUDA prioritizes optimized GPU processing, with Vesin used for eligible single-frame cases and dense processing as a fallback. - CPU inference uses Vesin for single-frame cases when available; otherwise, it uses dense processing. - Training continues to use dense processing for automatic selection. - **Bug Fixes** - Ensured automatic and unspecified builder settings produce consistent energy and force results. - **Tests** - Expanded coverage for single- and multi-frame backend selection and CPU/CUDA fallback behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: shaurya2k06 <shaurya2k06@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent bc902da commit a1dd9eb

5 files changed

Lines changed: 243 additions & 52 deletions

File tree

deepmd/pt_expt/infer/deep_eval.py

Lines changed: 34 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
# SPDX-License-Identifier: LGPL-3.0-or-later
22
import json
3-
import logging
43
import warnings
54
from collections.abc import (
65
Callable,
@@ -80,8 +79,6 @@
8079
NeighborGraph,
8180
)
8281

83-
log = logging.getLogger(__name__)
84-
8582

8683
# Public output keys emitted by graph-lower forwards, keyed by the
8784
# output-variable category that ``request_defs`` carries. The graph path is
@@ -192,15 +189,17 @@ class DeepEval(DeepEvalBackend):
192189
neighbor_graph_method : str, default: "auto"
193190
Carry-all graph builder for graph-form ``.pt2`` artifacts and
194191
graph-routed ``.pt`` checkpoints
195-
(``metadata["lower_input_kind"] == "graph"``): ``"auto"`` selects
196-
``"nv"`` on CUDA when nvalchemiops is available and otherwise falls
197-
back to ``"dense"``. ``"vesin"`` remains explicit opt-in because it
198-
loops over frames in Python. Explicit
192+
(``metadata["lower_input_kind"] == "graph"``): ``"auto"`` selects via
193+
:func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`
194+
at each eval call (CUDA: ``nv`` if importable; else ``vesin`` only when
195+
``nf == 1`` and importable; else ``dense``). Explicit
199196
``"dense"`` / ``"ase"`` / ``"vesin"`` / ``"nv"`` choices are preserved.
200197
A non-default value on any other artifact raises at construction because
201198
the knob would silently do nothing there; use ``nlist_backend`` for the
202199
nlist path instead. All builders emit the same neighbor set, so the
203-
choice is performance-only. Consolidating the two knobs into a single
200+
choice is performance-only. Training keeps a separate auto policy
201+
(:func:`~deepmd.pt_expt.utils.graph_builder.resolve_neighbor_graph_method`)
202+
that never selects ``vesin``. Consolidating the two knobs into a single
204203
backend-selection API is deferred to the dense-nlist deprecation.
205204
**kwargs : dict
206205
Keyword arguments.
@@ -271,33 +270,31 @@ def __init__(
271270
raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize")
272271

273272
@staticmethod
274-
def _resolve_neighbor_graph_method(method: str) -> str:
275-
"""Resolve the graph builder once for the active device."""
273+
def _resolve_neighbor_graph_method(method: str, nf: int | None = None) -> str:
274+
"""Validate and optionally resolve the graph builder for the active device.
275+
276+
``"auto"`` is left unresolved when ``nf`` is omitted so construction-
277+
time setup can defer to :meth:`_build_eval_graph`, where the frame
278+
count is known and vesin can be gated on ``nf == 1``.
279+
"""
276280
if method not in ("auto", "dense", "ase", "vesin", "nv"):
277281
raise ValueError(
278282
f"Unknown neighbor_graph_method {method!r}; "
279283
"expected 'auto', 'dense', 'ase', 'vesin', or 'nv'."
280284
)
281285
if method != "auto":
282286
return method
287+
if nf is None:
288+
return "auto"
283289

284-
from deepmd.pt.utils.nv_nlist import (
285-
is_nv_available,
286-
)
287290
from deepmd.pt_expt.utils.env import (
288291
DEVICE,
289292
)
293+
from deepmd.pt_expt.utils.graph_builder import (
294+
resolve_auto_graph_builder,
295+
)
290296

291-
if DEVICE.type == "cuda":
292-
if is_nv_available():
293-
return "nv"
294-
log.warning(
295-
"nvalchemi-toolkit-ops is unavailable; falling back from "
296-
"neighbor_graph_method='auto' to the dense graph builder. "
297-
"Install it with `pip install nvalchemi-toolkit-ops` to enable "
298-
"the NV graph builder."
299-
)
300-
return "dense"
297+
return resolve_auto_graph_builder(DEVICE, nf)
301298

302299
def _setup_neighbor_backend(self, nlist_backend: str) -> None:
303300
"""Resolve the graph or neighbor-list construction strategy.
@@ -2316,14 +2313,21 @@ def _build_eval_graph(
23162313
) -> "NeighborGraph":
23172314
"""Build the carry-all NeighborGraph for graph-lower inference.
23182315
2319-
Dispatches on ``self._neighbor_graph_method``: ``dense``/``ase`` run
2320-
backend-agnostic (numpy); ``vesin``/``nv`` run on-device (torch, O(N)).
2321-
All backends emit the SAME neighbor set (carry-all, sel-free), so the
2322-
selection is a pure performance choice and results are unchanged. The
2323-
result is canonicalized to the destination-major graph-form ``.pt2``
2324-
ABI after construction.
2316+
Dispatches on ``self._neighbor_graph_method``: ``auto`` is resolved
2317+
call-time via
2318+
:func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`
2319+
using the batch frame count (vesin only when ``nf == 1``);
2320+
``dense``/``ase`` run backend-agnostic (numpy); ``vesin``/``nv`` run
2321+
on-device (torch, O(N)). All backends emit the SAME neighbor set
2322+
(carry-all, sel-free), so the selection is a pure performance choice
2323+
and results are unchanged. The result is canonicalized to the
2324+
destination-major graph-form ``.pt2`` ABI after construction.
23252325
"""
23262326
method = self._neighbor_graph_method
2327+
if method == "auto":
2328+
coord_arr = np.asarray(coord_input)
2329+
nf = int(coord_arr.shape[0]) if coord_arr.ndim >= 2 else 1
2330+
method = self._resolve_neighbor_graph_method("auto", nf=nf)
23272331
# Model-level ``pair_exclude_types`` is a graph-BUILD transform
23282332
# (decision #18): apply it here so the exported ``.pt2`` lower consumes a
23292333
# pre-excluded ``edge_mask`` and never re-applies it (mirrors the C++
@@ -2392,7 +2396,7 @@ def _build_eval_graph(
23922396
)
23932397
raise ValueError(
23942398
f"unknown neighbor_graph_method {method!r}; "
2395-
"use 'dense', 'ase', 'vesin', or 'nv'"
2399+
"use 'auto', 'dense', 'ase', 'vesin', or 'nv'"
23962400
)
23972401

23982402
def _model_pair_excl(self) -> "PairExcludeMask | None":

deepmd/pt_expt/utils/graph_builder.py

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,93 @@
1818

1919
log = logging.getLogger(__name__)
2020

21+
# Warn once per process: resolve_auto_graph_builder runs per batch after
22+
# call-time resolution, but "install nvalchemi-toolkit-ops" is a one-shot
23+
# action for the user.
24+
_warned_auto_no_nv = False
25+
26+
27+
def resolve_auto_graph_builder(
28+
device: torch.device | str,
29+
nf: int = 1,
30+
) -> str:
31+
"""Resolve ``neighbor_graph_method="auto"`` to a concrete inference builder.
32+
33+
Single owner of the inference / DeepEval auto ladder. Training uses
34+
:func:`resolve_neighbor_graph_method`, which never selects ``vesin``.
35+
36+
Mirrors :func:`deepmd.pt.model.model.sezm_model._select_neighbor_builder`:
37+
``vesin`` is eligible only for a single-frame batch (``nf == 1``), because
38+
its API loops frames in Python (~1 ms/frame). Multi-frame batches stay on
39+
``nv`` (CUDA) or ``dense`` so ``auto_batch_size`` / ``dp test`` do not
40+
regress to the per-frame loop.
41+
42+
Policy
43+
------
44+
* CUDA + ``nvalchemiops``: ``nv`` (any ``nf``).
45+
* ``nf == 1`` + ``vesin.torch``: ``vesin``.
46+
* otherwise: ``dense``.
47+
48+
``ase`` is never chosen automatically. All builders emit the same carry-all
49+
neighbor set; the choice is performance-only. Builders run eagerly outside
50+
traced / compiled regions, so this does not change ``.pt2`` artifacts.
51+
52+
Parameters
53+
----------
54+
device : torch.device or str
55+
Device the coordinates live on (or will be moved to). Controls whether
56+
the CUDA-only ``nv`` builder is eligible.
57+
nf : int, default: 1
58+
Number of frames in the batch. ``vesin`` is selected only when
59+
``nf == 1`` and ``vesin.torch`` is importable.
60+
61+
Returns
62+
-------
63+
str
64+
One of ``"nv"``, ``"vesin"``, or ``"dense"``.
65+
66+
Raises
67+
------
68+
ValueError
69+
If ``nf`` is not a positive ``int`` (``bool`` is rejected).
70+
"""
71+
global _warned_auto_no_nv
72+
73+
from deepmd.pt.utils.nv_nlist import (
74+
is_nv_available,
75+
)
76+
from deepmd.pt_expt.utils.vesin_neighbor_list import (
77+
is_vesin_torch_available,
78+
)
79+
80+
# ``bool`` is a subclass of ``int``; reject it explicitly.
81+
if type(nf) is not int:
82+
raise ValueError(f"nf must be a positive int, got {nf!r}")
83+
if nf < 1:
84+
raise ValueError(f"nf must be >= 1, got {nf}")
85+
86+
dev = torch.device(device)
87+
nv_available = is_nv_available()
88+
if dev.type == "cuda" and nv_available:
89+
return "nv"
90+
if nf == 1 and is_vesin_torch_available():
91+
return "vesin"
92+
if dev.type == "cuda" and not nv_available:
93+
if not _warned_auto_no_nv:
94+
_warned_auto_no_nv = True
95+
log.warning(
96+
"nvalchemi-toolkit-ops is unavailable; falling back from "
97+
"neighbor_graph_method='auto' to the dense graph builder"
98+
+ (
99+
""
100+
if nf == 1
101+
else " (vesin is not used for nf>1; its API loops frames in Python)"
102+
)
103+
+ ". Install it with `pip install nvalchemi-toolkit-ops` to enable "
104+
"the NV graph builder."
105+
)
106+
return "dense"
107+
21108

22109
def resolve_neighbor_graph_method(
23110
requested: str,
@@ -36,6 +123,8 @@ def resolve_neighbor_graph_method(
36123
-------
37124
str
38125
The concrete builder name, either ``"dense"`` or ``"nv"``.
126+
Training auto never selects ``vesin`` (per-frame Python loop); use
127+
:func:`resolve_auto_graph_builder` for inference auto selection.
39128
40129
Raises
41130
------

deepmd/pt_expt/utils/vesin_graph_builder.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,12 @@
88
99
Scope note: ``vesin.torch``'s API is single-system, so this builder LOOPS over
1010
frames in Python (~1 ms/frame call overhead measured on GPU). It is intended
11-
for ``nf == 1`` inference and CPU use. It is never on a default hot path:
12-
``neighbor_graph_method=None`` resolves to the ``"dense"`` converter, and
13-
vesin is explicit opt-in only. For batched multi-frame GPU work prefer
14-
``nv`` (:mod:`.nv_graph_builder`), which batches all frames in one kernel.
11+
for ``nf == 1`` inference and CPU use. Inference ``neighbor_graph_method="auto"``
12+
(:func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`) selects
13+
vesin only when ``nf == 1`` and ``vesin.torch`` is importable (and ``nv`` is
14+
unavailable on CUDA); multi-frame batches stay on ``nv``/``dense``. Training
15+
auto never selects vesin. Prefer ``nv`` (:mod:`.nv_graph_builder`) for batched
16+
multi-frame GPU work, which batches all frames in one kernel.
1517
"""
1618

1719
from __future__ import (

source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -422,16 +422,31 @@ def test_unsupported_extension_raises(self) -> None:
422422
class TestNeighborGraphMethodResolution(unittest.TestCase):
423423
"""Auto graph-builder selection must cover each host policy explicitly."""
424424

425+
def test_auto_deferred_until_nf_known(self) -> None:
426+
"""Construction-time resolve leaves ``auto`` unresolved without ``nf``."""
427+
self.assertEqual(
428+
PtExptDeepEval._resolve_neighbor_graph_method("auto"),
429+
"auto",
430+
)
431+
425432
def test_auto_resolution(self) -> None:
433+
# (device, nv, vesin, nf, expected, warns)
426434
cases = (
427-
("cpu", False, "dense", False),
428-
("cuda", True, "nv", False),
429-
("cuda", False, "dense", True),
430-
)
431-
for device_type, nv_available, expected, warns in cases:
435+
("cpu", False, True, 1, "vesin", False),
436+
("cpu", False, True, 4, "dense", False),
437+
("cpu", False, False, 1, "dense", False),
438+
("cuda", True, True, 1, "nv", False),
439+
("cuda", True, True, 4, "nv", False),
440+
("cuda", False, True, 1, "vesin", False),
441+
("cuda", False, True, 4, "dense", True),
442+
("cuda", False, False, 1, "dense", True),
443+
)
444+
for device_type, nv_available, vesin_available, nf, expected, warns in cases:
432445
with self.subTest(
433446
device_type=device_type,
434447
nv_available=nv_available,
448+
vesin_available=vesin_available,
449+
nf=nf,
435450
):
436451
with (
437452
mock.patch(
@@ -442,17 +457,27 @@ def test_auto_resolution(self) -> None:
442457
"deepmd.pt.utils.nv_nlist.is_nv_available",
443458
return_value=nv_available,
444459
),
460+
mock.patch(
461+
"deepmd.pt_expt.utils.vesin_neighbor_list.is_vesin_torch_available",
462+
return_value=vesin_available,
463+
),
445464
):
446465
if warns:
466+
# reset warn-once so each case can assert the message
467+
import deepmd.pt_expt.utils.graph_builder as gb
468+
469+
gb._warned_auto_no_nv = False
447470
with self.assertLogs(
448-
"deepmd.pt_expt.infer.deep_eval",
471+
"deepmd.pt_expt.utils.graph_builder",
449472
level="WARNING",
450473
):
451474
actual = PtExptDeepEval._resolve_neighbor_graph_method(
452-
"auto"
475+
"auto", nf=nf
453476
)
454477
else:
455-
actual = PtExptDeepEval._resolve_neighbor_graph_method("auto")
478+
actual = PtExptDeepEval._resolve_neighbor_graph_method(
479+
"auto", nf=nf
480+
)
456481
self.assertEqual(actual, expected)
457482

458483

0 commit comments

Comments
 (0)