Skip to content

feat(pt_expt): auto-select O(N) NeighborGraph builder by default #5903

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Shaurya2k06 wants to merge 4 commits into
base: master
Choose a base branch
from
+102 −31
Open

feat(pt_expt): auto-select O(N) NeighborGraph builder by default #5903

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
01253d9
feat(pt_expt): include vesin in DeepEval auto graph-builder ladder
Shaurya2k06 Aug 1, 2026
00facb6
Merge branch 'master' into feat/pt_expt-auto-neighbor-graph-builder
OutisLi Aug 2, 2026
1838368
Merge branch 'master' into feat/pt_expt-auto-neighbor-graph-builder
Shaurya2k06 Aug 2, 2026
804d663
Merge branch 'master' into feat/pt_expt-auto-neighbor-graph-builder
Shaurya2k06 Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 13 additions & 22 deletions deepmd/pt_expt/infer/deep_eval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
# SPDX-License-Identifier: LGPL-3.0-or-later
import json
import logging
import warnings
from collections.abc import (
Callable,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
"""
Expand Down Expand Up @@ -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.
Expand Down
43 changes: 43 additions & 0 deletions deepmd/pt_expt/utils/graph_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,47 @@
log = logging.getLogger(__name__)


def resolve_auto_graph_builder(

@wanghan-iapcm wanghan-iapcm Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The summary describes two changes that are not in this diff.

It says the PR will "Flip the pt_expt model-level default (None / "auto") from hard-coded "dense" to that ladder" and "use the same helper in ... compiled training's eager _forward_graph". Neither deepmd/pt_expt/model/make_model.py nor deepmd/pt_expt/train/training.py appears in the changed files, and both still resolve the same way they did before:

  • make_model.py _resolve_graph_method still ends at getattr(self, "neighbor_graph_method", "dense"), so a model driven directly still defaults to dense on every device;
  • training.py still imports only resolve_neighbor_graph_method, and _forward_graph still reads getattr(_model, "neighbor_graph_method", "dense").

build_neighbor_graph_for_method also has no "auto" branch -- it raises ValueError on anything it does not recognise -- so "auto" could not reach it even if the model default did produce it.

I think this is stale text rather than a missing change: those two flips landed in #5912 and #5913, which are already ancestors of this branch. The "Why existing tests missed this" paragraph has the same problem, since it describes fixing compiled training's hardcoded dense. Worth rewriting the body to what the commit does -- add vesin to the inference auto ladder and extract the shared helper -- because as written a reviewer would look for a model-level behaviour change that is not here, and a bisect later would be misled about where the flip came from.

device: torch.device | str,
) -> str:
"""Resolve ``neighbor_graph_method="auto"`` to a concrete inference builder.

@wanghan-iapcm wanghan-iapcm Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing the numpydoc sections the rest of this module uses.

This is a public function -- no leading underscore, imported by deep_eval.py -- and its sibling resolve_neighbor_graph_method a few lines below carries full Parameters / Returns / Raises blocks. This one documents the ladder in prose and a bullet list with no Parameters for device and no Returns for the str.

_select_neighbor_builder in deepmd/pt/model/model/sezm_model.py, which does the analogous job, also documents nf and device under Parameters and its return under Returns. Matching that is a small edit and keeps the API docs uniform.


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():

@wanghan-iapcm wanghan-iapcm Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the change the PR really makes, and it re-adds what #5912 removed.

On #5912 I raised this against an identical ladder:

This branch puts vesin on the default path, which contradicts the builder's own documented contract. [...] After this change, any environment with vesin.torch installed and without CUDA + nvalchemiops gets that per-frame Python loop by default. That includes runs under auto_batch_size, which deliberately batches many frames into a single _eval_model_graph call - exactly where the vectorised dense builder handles all frames at once and the vesin loop is worst. [...] Either prefer dense when nframes > 1, or update the contract to say vesin is now a default and re-examine whether the per-frame loop is acceptable there.

That rung was dropped and #5912 merged without it. This commit restores it and takes the second half of that either/or -- the module docstring is rewritten -- but not the first: there is no nframes gate, and no measurement re-examining the loop.

The reason I do not think the docstring rewrite settles it is that the repository already has an answer to this exact question, and it goes the other way. _select_neighbor_builder picks between the same two builders:

if device.type == "cpu" and nf == 1 and is_vesin_torch_available():
    return VesinNeighborList()

with the rationale stated directly above it: "Every other case -- any CUDA input or any multi-frame batch -- uses nvalchemiops, whose batched kernel amortizes the launch cost across frames." It takes nf, and it is called per forward. This resolver takes only device and is called once from _setup_neighbor_backend inside DeepEval.__init__, before any frame count exists -- so it cannot express that policy even if we wanted it to. That is a structural difference, not a parameter we forgot to thread through.

Three things make the exposure wider than it looks. vesin[torch] is an unconditional dependency of the torch extra, so this is the CPU default for essentially every pt user rather than an opt-in for people who installed something extra. auto_batch_size defaults to True, and _eval_model_graph receives whole batches from execute_all, so multi-frame is the normal case for dp test and dp model-devi, not the exception. And the builder's own scope note puts the cost at "~1 ms/frame call overhead". For a small system over many frames the arithmetic points at a slowdown, which inverts the PR's stated goal.

What would resolve it, in order of preference: move the resolution to call time and gate on nf == 1 for vesin, matching _select_neighbor_builder; or keep construction-time resolution and drop the vesin rung, leaving it explicit opt-in as #5912 concluded; or keep it and post a benchmark over a realistic dp test batch showing dense is not faster. Any of the three is fine by me -- what I do not want is the decision being reversed silently.

return "vesin"
return "dense"


def resolve_neighbor_graph_method(
requested: str,
device: torch.device,
Expand All @@ -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
------
Expand Down
10 changes: 6 additions & 4 deletions deepmd/pt_expt/utils/vesin_graph_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
18 changes: 13 additions & 5 deletions source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand Down
27 changes: 27 additions & 0 deletions source/tests/pt_expt/model/test_graph_builder_dispatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,33 @@ def test_explicit_nv_rejects_cpu():
resolve_neighbor_graph_method("nv", torch.device("cpu"))


@pytest.mark.parametrize(

@wanghan-iapcm wanghan-iapcm Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These five cases pin the ladder, and I want to say they are the right shape -- patching availability and device rather than recomputing the cascade in the assertion is exactly what was missing from the equivalent test on #5912.

The gap is that they are the only new coverage, and they are pure resolver assertions: they never build a graph. The cross-builder numerical checks that do build one, test_vesin_matches_dense_energy_force and test_nv_matches_dense_energy_force further up this file, share the _eval helper, and it is single-frame:

coord = torch.tensor(rng.random((1, 6, 3)) * 4.0, ...)

So parity between vesin and dense is established only at nf == 1 -- the regime the vesin builder's own docstring says it is "intended for" -- while this PR makes vesin the default precisely for the batched regime, where nothing compares it against dense. That matters slightly more than a generic coverage note because the graph lower accumulates with segment_sum over edges and edge ordering after canonicalize is an independent implementation per builder, so multi-frame agreement does not follow from single-frame agreement; it needs its own case.

A parametrization of _eval over nf in {1, 4} would cover it, and it is the test I would want in place before the default moves, whichever way the vesin question above is settled.

("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)
Expand Down
Loading