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

Replace the hard-coded None→dense default with a shared availability-
probing ladder (CUDA: nv→vesin→dense; CPU: vesin→dense) used by the
model default-flip, DeepEval, and compiled training's eager builder.

Signed-off-by: shaurya2k06 <shaurya2k06@gmail.com>

Author: Shaurya2k06

deepmd/pt_expt/infer/deep_eval.py

@@ -139,15 +139,18 @@ class DeepEval(DeepEvalBackend):
         Neighbor-list builder for the NLIST/extended lower path (``.pte`` and
         nlist-form ``.pt2``): ``"auto"`` / ``"vesin"`` / ``"native"``. Not
         used by graph-form ``.pt2`` artifacts.
-    neighbor_graph_method : str, default: "dense"
+    neighbor_graph_method : str, default: "auto"
         Carry-all graph builder for GRAPH-FORM ``.pt2`` artifacts ONLY
-        (``metadata["lower_input_kind"] == "graph"``): ``"dense"`` / ``"ase"``
-        (backend-agnostic) or ``"vesin"`` / ``"nv"`` (on-device O(N)). A
-        non-default value on any other artifact raises at construction — 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.
+        (``metadata["lower_input_kind"] == "graph"``): ``"auto"`` (availability-
+        probed O(N) default; see
+        :func:`~deepmd.pt_expt.utils.neighbor_graph_method.resolve_auto_graph_builder`),
+        ``"dense"`` / ``"ase"`` (backend-agnostic), or ``"vesin"`` / ``"nv"``
+        (on-device O(N)). A non-default explicit value on any other artifact
+        raises at construction — 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.
     **kwargs : dict
         Keyword arguments.
     """
@@ -160,14 +163,14 @@ def __init__(
         auto_batch_size: bool | int | AutoBatchSize = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
         nlist_backend: str = "auto",
-        neighbor_graph_method: str = "dense",
+        neighbor_graph_method: str = "auto",
         **kwargs: Any,
     ) -> None:
         self.output_def = output_def
         self.model_path = model_file
         self.neighbor_list = neighbor_list
         # World-2 graph-form ``.pt2`` (lower_input_kind == "graph") builder select:
-        # "dense"/"ase" (backend-agnostic) or "vesin"/"nv" (on-device O(N)).
+        # "auto" (probed) / "dense"/"ase" (backend-agnostic) / "vesin"/"nv" (O(N)).
         self._neighbor_graph_method = neighbor_graph_method
         self._is_pt2 = model_file.endswith(".pt2")
 
@@ -185,11 +188,13 @@ def __init__(
             )
 
         # neighbor_graph_method is consumed ONLY by graph-form .pt2 eval
-        # (_eval_model_graph); fail fast instead of silently ignoring it on
-        # nlist-form artifacts (there, the builder knob is nlist_backend).
-        if neighbor_graph_method != "dense" and getattr(self, "metadata", {}).get(
-            "lower_input_kind"
-        ) not in ("graph", "dpa1_canonical"):
+        # (_eval_model_graph); fail fast instead of silently ignoring an
+        # EXPLICIT builder on nlist-form artifacts (there, the builder knob
+        # is nlist_backend). "auto" / "dense" are allowed as no-ops so the
+        # default constructor works for every artifact kind.
+        if neighbor_graph_method not in ("auto", "dense") and getattr(
+            self, "metadata", {}
+        ).get("lower_input_kind") not in ("graph", "dpa1_canonical"):
             raise ValueError(
                 f"neighbor_graph_method={neighbor_graph_method!r} only applies to "
                 "graph-form .pt2 artifacts (lower_input_kind == 'graph'); this "
@@ -1930,14 +1935,21 @@ def _build_eval_graph(
     ) -> "NeighborGraph":
         """Build the carry-all NeighborGraph for graph-form ``.pt2`` 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`` resolves via
+        :func:`~deepmd.pt_expt.utils.neighbor_graph_method.resolve_auto_graph_builder`;
+        ``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.
         """
+        from deepmd.pt_expt.utils.neighbor_graph_method import (
+            resolve_auto_graph_builder,
+        )
+
         method = self._neighbor_graph_method
+        if method == "auto":
+            method = resolve_auto_graph_builder(device)
         # 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++
@@ -2006,7 +2018,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":

deepmd/pt_expt/model/make_model.py

@@ -266,11 +266,19 @@ def _build_graph_for_method(
     Single owning site for the graph-builder dispatch shared by
     :meth:`_call_common_graph` and the graph Hessian wrapper
     (:class:`_WrapperForwardEnergyGraph`), so both build the graph identically.
+    ``"auto"`` is resolved device-aware via
+    :func:`~deepmd.pt_expt.utils.neighbor_graph_method.resolve_auto_graph_builder`.
     """
     from deepmd.dpmodel.utils.neighbor_graph import (
         build_neighbor_graph,
         build_neighbor_graph_ase,
     )
+    from deepmd.pt_expt.utils.neighbor_graph_method import (
+        resolve_auto_graph_builder,
+    )
+
+    if method == "auto":
+        method = resolve_auto_graph_builder(coord.device)
 
     if method == "dense":
         return build_neighbor_graph(
@@ -297,7 +305,7 @@ def _build_graph_for_method(
             coord, atype, box, rcut, with_csr=with_csr, pair_excl=pair_excl
         )
     raise ValueError(
-        f"unknown neighbor_graph_method {method!r}; use 'dense', 'ase', "
+        f"unknown neighbor_graph_method {method!r}; use 'auto', 'dense', 'ase', "
         "'vesin', or 'nv'"
     )
 
@@ -665,27 +673,35 @@ def forward_common_lower_graph(
         def _resolve_graph_method(
             self, neighbor_graph_method: str | None
         ) -> str | None:
-            """pt_expt default-flip (decision #17): ``None`` => carry-all graph for
-            graph-eligible mixed_types descriptors, else dense. Unlike dpmodel/jax,
-            pt_expt has the autograd ``forward_common_lower_graph`` that produces
-            force/virial on the graph, so the graph can be the DEFAULT here.
-            ``"legacy"`` forces dense; explicit ``"dense"``/``"ase"`` force the graph.
+            """pt_expt default-flip (decision #17): ``None`` / ``"auto"`` =>
+            carry-all graph for graph-eligible mixed_types descriptors, with
+            the concrete builder chosen by
+            :func:`~deepmd.pt_expt.utils.neighbor_graph_method.resolve_auto_graph_builder`
+            at build time. Unlike dpmodel/jax, pt_expt has the autograd
+            ``forward_common_lower_graph`` that produces force/virial on the
+            graph, so the graph can be the DEFAULT here. ``"legacy"`` forces
+            the dense-nlist path; explicit ``"dense"``/``"ase"``/``"vesin"``/
+            ``"nv"`` force that graph builder.
 
             Parameters
             ----------
             neighbor_graph_method
-                The user-requested method: ``None`` (default-flip), ``"legacy"``
-                (force dense), or ``"dense"``/``"ase"`` (force the graph builder).
+                The user-requested method: ``None`` / ``"auto"`` (default-flip
+                to the availability-probed O(N) builder), ``"legacy"`` (force
+                dense nlist), or an explicit graph builder name.
 
             Returns
             -------
             method
                 The resolved method passed to :meth:`_call_common_graph`, or
-                ``None`` to take the dense path.
+                ``None`` to take the dense-nlist path. Eligible defaults return
+                ``"auto"`` (resolved device-aware in :func:`_build_graph_for_method`).
             """
             if neighbor_graph_method == "legacy":
                 return None
-            if neighbor_graph_method is not None:
+            # ``"auto"`` is synonymous with ``None`` for the eligibility gate;
+            # both become the concrete builder only at graph-build time.
+            if neighbor_graph_method not in (None, "auto"):
                 return neighbor_graph_method
             # The DEFAULT-flip is gated to ENERGY-output models: the graph
             # lower itself is output-agnostic, but the compiled-training
@@ -703,7 +719,7 @@ def _resolve_graph_method(
             descriptor = getattr(self.atomic_model, "descriptor", None)
             uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False)
             if self.mixed_types() and uses_graph_lower():
-                return "dense"
+                return "auto"
             return None
 
         def _call_common_graph(
@@ -736,7 +752,8 @@ def _call_common_graph(
             ap
                 the atomic parameter. nf x nloc x nda
             method
-                the carry-all builder, ``"dense"`` or ``"ase"``.
+                the carry-all builder: ``"auto"``, ``"dense"``, ``"ase"``,
+                ``"vesin"``, or ``"nv"``.
             do_atomic_virial
                 whether to calculate the atomic virial.
 

deepmd/pt_expt/train/training.py

@@ -1123,15 +1123,16 @@ def _forward_graph(
     ) -> dict[str, torch.Tensor]:
         """Carry-all GRAPH forward -> compiled ``forward_common_lower_graph``.
 
-        Builds the carry-all NeighborGraph eagerly (the SAME builder the eager
-        uncompiled default-flip uses, so the graph tensors are bit-identical),
-        then calls the compiled graph lower.  The graph force is per-LOCAL-node
-        ``(N, 3)`` with ``N == nframes * nloc`` for a single-rank carry-all graph,
-        so no extended->local scatter is needed; only the flat ``(N, *)`` node
-        keys are unravelled to ``(nf, nloc, *)`` at the I/O boundary.
+        Builds the carry-all NeighborGraph eagerly (the SAME auto-resolved
+        builder the eager uncompiled default-flip uses, so the graph tensors
+        match up to fp addition order), then calls the compiled graph lower.
+        The graph force is per-LOCAL-node ``(N, 3)`` with ``N == nframes * nloc``
+        for a single-rank carry-all graph, so no extended->local scatter is
+        needed; only the flat ``(N, *)`` node keys are unravelled to
+        ``(nf, nloc, *)`` at the I/O boundary.
         """
-        from deepmd.dpmodel.utils.neighbor_graph import (
-            build_neighbor_graph,
+        from deepmd.pt_expt.model.make_model import (
+            _build_graph_for_method,
         )
 
         _model = self.original_model
@@ -1181,8 +1182,9 @@ def _forward_graph(
         # level pair_exclude is a graph-BUILD transform (decision #18): fold it
         # into edge_mask here so the compiled lower consumes a pre-excluded graph
         # (the lower no longer re-applies it), matching the eager path exactly.
+        # ``"auto"`` resolves via resolve_auto_graph_builder(coord.device).
         pair_excl = getattr(_model.atomic_model, "pair_excl", None)
-        ng = build_neighbor_graph(coord_3d, atype, box_flat, rcut, pair_excl=pair_excl)
+        ng = _build_graph_for_method("auto", coord_3d, atype, box_flat, rcut, pair_excl)
         atype_flat = atype.reshape(nframes * nloc)
 
         # Lazy compile of the GRAPH lower (cached per structure key).

deepmd/pt_expt/utils/neighbor_graph_method.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Shared NeighborGraph builder selection for pt_expt.
+
+Single owner of the ``None`` / ``"auto"`` resolution ladder used by the
+model-level default-flip (:meth:`deepmd.pt_expt.model.make_model` /
+``_resolve_graph_method``), DeepEval graph-form ``.pt2`` inference, and
+compiled training's eager builder outside the traced lower.
+
+Policy
+------
+* CUDA device: ``nv`` if importable, else ``vesin`` if importable, else
+  ``dense``.
+* CPU device: ``vesin`` if importable, else ``dense``.
+
+``ase`` is never chosen automatically (explicit opt-in only). All builders
+emit the same carry-all neighbor set; the choice is performance-only.
+Builders run eagerly outside traced / compiled regions (export and training
+compile use synthetic dense graph inputs), so flipping the default does not
+change ``.pt2`` artifacts.
+
+Perf note: at small benchmark sizes the dense all-pairs builder is typically
+not the bottleneck; the O(N) win appears for large systems (N of a few
+thousand atoms and up). Time builders manually on a large system to document
+the crossover before relying on auto for production throughput.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+import torch
+
+from deepmd.pt.utils.nv_nlist import (
+    is_nv_available,
+)
+from deepmd.pt_expt.utils.vesin_neighbor_list import (
+    is_vesin_torch_available,
+)
+
+
+def resolve_auto_graph_builder(
+    device: torch.device | str,
+) -> str:
+    """Resolve ``neighbor_graph_method`` ``None`` / ``"auto"`` to a concrete builder.
+
+    Parameters
+    ----------
+    device
+        Device the coordinates live on (or will be moved to). Controls whether
+        the CUDA-only ``nv`` builder is eligible.
+
+    Returns
+    -------
+    str
+        One of ``"nv"``, ``"vesin"``, or ``"dense"``.
+    """
+    dev = torch.device(device)
+    if dev.type == "cuda":
+        if is_nv_available():
+            return "nv"
+        if is_vesin_torch_available():
+            return "vesin"
+        return "dense"
+    if is_vesin_torch_available():
+        return "vesin"
+    return "dense"

deepmd/pt_expt/utils/nv_graph_builder.py

@@ -11,6 +11,10 @@
 Unlike the vesin builder, nvalchemiops batches natively over frames via
 ``batch_idx``/``batch_ptr`` -- a single GPU kernel handles all ``nf`` frames,
 so there is NO per-frame Python loop. CUDA-only ⇒ this module lives in pt_expt.
+On CUDA it is the first choice of
+:func:`~deepmd.pt_expt.utils.neighbor_graph_method.resolve_auto_graph_builder`
+when ``nvalchemiops`` is importable. Export / training compile still use
+synthetic dense graph inputs — the builder choice does not affect ``.pt2``.
 
 The matrix decode mirrors :func:`deepmd.pt.utils.nv_nlist._matrix_to_extended_inputs`
 (the authoritative, tested extraction) but stops at the sparse ``(i, j, S)``

deepmd/pt_expt/utils/vesin_graph_builder.py

@@ -8,10 +8,13 @@
 
 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. On CPU (and on CUDA when ``nv`` is
+unavailable) it is selected by the shared
+:func:`~deepmd.pt_expt.utils.neighbor_graph_method.resolve_auto_graph_builder`
+default ladder; otherwise prefer ``nv`` (:mod:`.nv_graph_builder`) for batched
+multi-frame GPU work, which batches all frames in one kernel. Export and
+training compile still use synthetic dense graph inputs — the builder choice
+does not affect ``.pt2`` artifacts.
 """
 
 from __future__ import (

source/tests/pt_expt/infer/test_graph_deepeval.py

@@ -252,7 +252,7 @@ def test_graph_pt2_deepeval_parity(graph_pt2, pbc, system) -> None:
 @pytest.mark.parametrize("pbc", [True, False])  # periodic vs non-periodic
 def test_graph_pt2_deepeval_vesin_matches_dense(graph_pt2, pbc) -> None:
     """Selecting neighbor_graph_method='vesin' at DeepEval yields identical
-    energy/force/virial to the default 'dense' builder on the SAME graph ``.pt2``
+    energy/force/virial to the default 'auto' builder on the SAME graph ``.pt2``
     (the builder is a pure perf choice; neighbor sets are equal).
     """
     pt2_path, _ = graph_pt2
@@ -261,7 +261,7 @@ def test_graph_pt2_deepeval_vesin_matches_dense(graph_pt2, pbc) -> None:
     max_nn = _max_neighbors(coords, box, atype)
     assert max_nn < SEL, "test system must be non-binding for carry-all parity"
 
-    dp_dense = DeepPot(pt2_path)  # default neighbor_graph_method == "dense"
+    dp_dense = DeepPot(pt2_path, neighbor_graph_method="dense")
     dp_vesin = DeepPot(pt2_path, neighbor_graph_method="vesin")
     assert dp_vesin.deep_eval._neighbor_graph_method == "vesin"
 
@@ -272,6 +272,37 @@ def test_graph_pt2_deepeval_vesin_matches_dense(graph_pt2, pbc) -> None:
     np.testing.assert_allclose(v_v, v_d, rtol=1e-10, atol=1e-10, err_msg="virial")
 
 
+@pytest.mark.parametrize("pbc", [True, False])
+def test_graph_pt2_deepeval_auto_matches_dense(graph_pt2, pbc) -> None:
+    """Default / ``neighbor_graph_method='auto'`` match explicit ``'dense'``.
+
+    Builders are value-transparent; auto only changes which O(N) backend is
+    used when available. Export still uses synthetic dense inputs, so the
+    ``.pt2`` itself is unaffected by the auto default.
+    """
+    pt2_path, _ = graph_pt2
+    coords, cells, atype = _build_system(**_SYSTEMS["small_8"])
+    box = cells if pbc else None
+    max_nn = _max_neighbors(coords, box, atype)
+    assert max_nn < SEL, "test system must be non-binding for carry-all parity"
+
+    dp_default = DeepPot(pt2_path)
+    dp_auto = DeepPot(pt2_path, neighbor_graph_method="auto")
+    dp_dense = DeepPot(pt2_path, neighbor_graph_method="dense")
+    assert dp_default.deep_eval._neighbor_graph_method == "auto"
+    assert dp_auto.deep_eval._neighbor_graph_method == "auto"
+
+    e_def, f_def, v_def = dp_default.eval(coords, box, atype)
+    e_auto, f_auto, v_auto = dp_auto.eval(coords, box, atype)
+    e_d, f_d, v_d = dp_dense.eval(coords, box, atype)
+    np.testing.assert_allclose(e_def, e_d, rtol=1e-10, atol=1e-10, err_msg="energy")
+    np.testing.assert_allclose(f_def, f_d, rtol=1e-10, atol=1e-10, err_msg="force")
+    np.testing.assert_allclose(v_def, v_d, rtol=1e-10, atol=1e-10, err_msg="virial")
+    np.testing.assert_allclose(e_auto, e_d, rtol=1e-10, atol=1e-10, err_msg="energy")
+    np.testing.assert_allclose(f_auto, f_d, rtol=1e-10, atol=1e-10, err_msg="force")
+    np.testing.assert_allclose(v_auto, v_d, rtol=1e-10, atol=1e-10, err_msg="virial")
+
+
 def test_graph_pt2_single_atom_no_edges(graph_pt2) -> None:
     """A single isolated atom (zero real edges) evaluates through the ``.pt2``.