From f0831247cba651cca203759bfe1e2c97f3aa28bb Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <198997215+OutisLi-Bot@users.noreply.github.com>
Date: Thu, 15 Oct 2026 17:04:34 +0800
Subject: [PATCH 1/5] fix(pt_expt): fold vacuum reference on CPU tracing device

Freezing with vacuum_ref ran fold_vacuum_reference before moving the
model to the CPU tracing device, so vacuum_descriptor could OOM on CUDA.
Moving to CPU first avoids that allocation on the GPU.

fold_vacuum_reference replaces bias_atom_e (and may set vacuum_table)
through numpy assignment. dpmodel_setattr previously always placed those
buffers on env.DEVICE, so a fold after Module.to("cpu") pulled tensors
back to CUDA and broke FakeTensor tracing. Place new buffers on the
module's current device instead.
---
 deepmd/pt_expt/common.py                      | 38 +++++++++++++++----
 deepmd/pt_expt/utils/serialization.py         | 12 +++---
 .../pt_expt/model/test_dpa4_vacuum_ref.py     | 19 ++++++++++
 source/tests/pt_expt/utils/test_common.py     | 32 ++++++++++++++++
 4 files changed, 89 insertions(+), 12 deletions(-)

diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py
index 42080a9142..1fd0a9df6f 100644
--- a/deepmd/pt_expt/common.py
+++ b/deepmd/pt_expt/common.py
@@ -238,6 +238,28 @@ def _try_convert_list(name: str, value: list) -> torch.nn.Module | None:
     return None
 
 
+def _module_buffer_device(obj: torch.nn.Module, name: str) -> torch.device:
+    """Device for a numpy array assigned as a buffer on ``obj``.
+
+    Prefer the existing buffer of the same name when it already holds a
+    tensor, then any parameter or non-None buffer on the module, then
+    ``env.DEVICE``. Assignments after ``Module.to`` therefore keep residency
+    with the module rather than jumping back to the process default device.
+    """
+    from deepmd.pt_expt.utils import env  # deferred - avoids circular import
+
+    if "_buffers" in obj.__dict__:
+        existing = obj._buffers.get(name)
+        if existing is not None:
+            return existing.device
+    for tensor in obj.parameters():
+        return tensor.device
+    for tensor in obj.buffers():
+        if tensor is not None:
+            return tensor.device
+    return env.DEVICE
+
+
 def dpmodel_setattr(obj: torch.nn.Module, name: str, value: Any) -> tuple[bool, Any]:
     """Common __setattr__ logic for pt_expt wrappers around dpmodel classes.
 
@@ -284,6 +306,10 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer
     - In dpmodel, all persistent arrays use `self.xxx = np.array(...)`. Scalars
       use `.item()`, lists use `.tolist()`. So `isinstance(value, np.ndarray)`
       reliably identifies buffer-worthy attributes.
+    - A numpy array becomes a buffer on the module's current device (see
+      :func:`_module_buffer_device`), not unconditionally on ``env.DEVICE``, so
+      a later assignment after ``Module.to`` does not pull tensors back to the
+      process default device.
     - torch.Tensor values assigned to existing buffers fall through to
       torch.nn.Module.__setattr__, which correctly updates them.
     - dpmodel objects are identified by registry lookup (exact type match), so
@@ -293,10 +319,10 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer
 
     **Circular import resolution:**
 
-    The function uses a deferred import `from deepmd.pt_expt.utils import env`
-    inside the function body. This breaks the circular dependency chain:
-    common.py → utils/__init__.py → exclude_mask.py → common.py. The import is
-    cached by Python after the first call, so there's no performance penalty.
+    Device resolution and related helpers defer
+    ``from deepmd.pt_expt.utils import env`` so this module does not import
+    ``utils`` at load time (common.py → utils/__init__.py → exclude_mask.py →
+    common.py). The import is cached after the first call.
 
     **Usage pattern:**
 
@@ -319,11 +345,9 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer
     >>> # Case 3: dpmodel object → pt_expt module
     >>> obj.emask = AtomExcludeMaskDP(...)  # becomes AtomExcludeMask module
     """
-    from deepmd.pt_expt.utils import env  # deferred - avoids circular import
-
     # numpy array → torch buffer
     if isinstance(value, np.ndarray) and "_buffers" in obj.__dict__:
-        tensor = torch.as_tensor(value, device=env.DEVICE)
+        tensor = torch.as_tensor(value, device=_module_buffer_device(obj, name))
         if name in obj._buffers:
             obj._buffers[name] = tensor
             return True, tensor
diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py
index aeb11d798a..e1afd50d50 100644
--- a/deepmd/pt_expt/utils/serialization.py
+++ b/deepmd/pt_expt/utils/serialization.py
@@ -1833,12 +1833,14 @@ def _trace_and_export_impl(
         # pt_expt BaseModel registry returns this backend's torch class.
         model = BaseModel.deserialize(data["model"])
     model.eval()
-    # The vacuum reference is resolved on the export object, folded into the
-    # fitting bias or stored as a per-type table, so the exported graph
-    # carries no reference atoms; the move to the tracing device follows, so
-    # the table lands there with the rest of the model.
-    model.fold_vacuum_reference()
+    # Fold the vacuum reference on the CPU tracing device. vacuum_descriptor
+    # evaluates an isolated atom of every type through the descriptor; on CUDA
+    # that path can allocate several GiB even when export already holds most of
+    # the GPU. Move first so the fold runs where tracing will, and so the
+    # folded bias or vacuum table stays on that device with the rest of the
+    # model.
     model.to("cpu")
+    model.fold_vacuum_reference()
 
     # Device-dependent Python branches resolve on the CPU tracing inputs, so
     # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2)
diff --git a/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py b/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py
index 2dc3147fe7..4de9f248f4 100644
--- a/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py
+++ b/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py
@@ -209,6 +209,25 @@ def test_fold_reproduces_the_referenced_model() -> None:
     )
 
 
+def test_fold_vacuum_reference_keeps_cpu_residency(monkeypatch) -> None:
+    """Folding on CPU must leave every parameter and buffer on CPU.
+
+    Export moves the model to the CPU tracing device before folding. The fold
+    replaces ``bias_atom_e`` (and may set ``vacuum_table``) through numpy
+    assignment; those buffers must not jump back to ``env.DEVICE``.
+    """
+    model = make_model(True).to("cpu")
+    monkeypatch.setattr(env, "DEVICE", torch.device("meta"))
+    model.fold_vacuum_reference()
+    assert not model.atomic_model.fitting_net.vacuum_ref
+    for name, tensor in model.named_parameters():
+        assert tensor.device.type == "cpu", name
+    for name, tensor in model.named_buffers():
+        if tensor is None:
+            continue
+        assert tensor.device.type == "cpu", name
+
+
 def test_dense_and_graph_routes_share_the_vacuum_descriptor() -> None:
     """The reference rows appended to a graph equal the single-atom-frame descriptor."""
     model = make_model(vacuum_ref=True)
diff --git a/source/tests/pt_expt/utils/test_common.py b/source/tests/pt_expt/utils/test_common.py
index 57ee153e54..f0fa2c524f 100644
--- a/source/tests/pt_expt/utils/test_common.py
+++ b/source/tests/pt_expt/utils/test_common.py
@@ -127,3 +127,35 @@ class MockModule(MockNativeOP):
     assert hasattr(module, "forward")
     # Should not have forward_lower method since there's no call_lower method
     assert not hasattr(module, "forward_lower")
+
+
+def test_dpmodel_setattr_keeps_module_device(monkeypatch) -> None:
+    """Numpy buffer assignment after Module.to keeps the module device.
+
+    ``env.DEVICE`` is patched to ``meta`` so a regression that always places
+    new buffers on the process default would land off-CPU and fail the assert.
+    """
+    from deepmd.dpmodel.common import (
+        NativeOP,
+    )
+    from deepmd.pt_expt.common import (
+        torch_module,
+    )
+
+    class HolderDP(NativeOP):
+        def __init__(self) -> None:
+            self.stat = np.zeros(3, dtype=np.float64)
+
+        def call(self, x: np.ndarray) -> np.ndarray:
+            return x
+
+    @torch_module
+    class Holder(HolderDP):
+        pass
+
+    module = Holder().to("cpu")
+    assert module.stat.device.type == "cpu"
+    monkeypatch.setattr(env, "DEVICE", torch.device("meta"))
+    module.stat = np.ones(3, dtype=np.float64)
+    assert module.stat.device.type == "cpu"
+    np.testing.assert_array_equal(module.stat.detach().cpu().numpy(), np.ones(3))

From db49062e8cd028bae55339af892f0574c04e8435 Mon Sep 17 00:00:00 2001
From: OutisLi <137472077+OutisLi@users.noreply.github.com>
Date: Thu, 15 Oct 2026 17:52:28 +0800
Subject: [PATCH 2/5] fix(vacuum): bound graph references and preserve export
 kernels

Compute graph-native isolated references without dense sel padding.
Keep legacy descriptor reference semantics behind an explicit capability.
Fold on CPU reference descriptors without autograd, restoring deployment
kernels and retaining fitting reference tables before fake tracing.

Add allocation, conditioning, dispatch, export-order and restoration tests.
---
 .../dpmodel/atomic_model/dp_atomic_model.py   |  40 ++-
 deepmd/dpmodel/descriptor/dpa4.py             |   4 +
 deepmd/dpmodel/descriptor/dpa4c.py            |   4 +
 .../descriptor/make_base_descriptor.py        |  12 +
 deepmd/pt_expt/utils/serialization.py         |  67 +++-
 .../dpmodel/test_vacuum_descriptor_graph.py   | 289 ++++++++++++++++++
 .../utils/test_serialization_vacuum_ref.py    | 253 +++++++++++++++
 7 files changed, 661 insertions(+), 8 deletions(-)
 create mode 100644 source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
 create mode 100644 source/tests/pt_expt/utils/test_serialization_vacuum_ref.py

diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py
index ffd33b5bd2..e92f78411d 100644
--- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py
@@ -403,6 +403,10 @@ def vacuum_descriptor(self) -> Array:
         condition, and a spin vector of one Bohr magneton per unpaired
         electron when it takes the native spin.
 
+        Descriptors that opt into graph vacuum evaluation use isolated nodes,
+        independently of their advertised neighbor capacity; other descriptors
+        retain their configured all-masked dense neighbor list.
+
         Returns
         -------
         Array
@@ -413,10 +417,36 @@ def vacuum_descriptor(self) -> Array:
         xp = array_api_compat.array_namespace(bias)
         device = array_api_compat.device(bias)
         ntypes = self.get_ntypes()
+        atype = xp.arange(ntypes, dtype=xp.int64, device=device)
+        conditions = {
+            name: xp.asarray(table, dtype=bias.dtype, device=device)
+            for name, table in self.vacuum_conditions().items()
+        }
+        if self.uses_graph_lower() and self.descriptor.supports_graph_vacuum():
+            from deepmd.dpmodel.utils.neighbor_graph import (
+                NeighborGraph,
+            )
+
+            # One isolated node per frame, with the same two masked guard
+            # edges as the graph builders. Never size this graph from sel:
+            # graph-only descriptors may report an unbounded sentinel there.
+            # The graph interface also carries native spin, which the dense
+            # compatibility adapters need not support.
+            graph = NeighborGraph(
+                n_node=xp.ones((ntypes,), dtype=xp.int64, device=device),
+                edge_index=xp.zeros((2, 2), dtype=xp.int64, device=device),
+                edge_vec=xp.zeros((2, 3), dtype=bias.dtype, device=device),
+                edge_mask=xp.zeros((2,), dtype=xp.bool, device=device),
+            )
+            return self.descriptor.call_graph(
+                graph,
+                atype,
+                type_embedding=self.descriptor.graph_type_embedding_table(),
+                **conditions,
+            )[0]
+
         coord = xp.zeros((ntypes, 1, 3), dtype=bias.dtype, device=device)
-        atype = xp.reshape(
-            xp.arange(ntypes, dtype=xp.int64, device=device), (ntypes, 1)
-        )
+        atype = xp.reshape(atype, (ntypes, 1))
         nlist = xp.full(
             (ntypes, 1, self.descriptor.get_nnei()),
             -1,
@@ -424,10 +454,6 @@ def vacuum_descriptor(self) -> Array:
             device=device,
         )
         mapping = xp.zeros((ntypes, 1), dtype=xp.int64, device=device)
-        conditions = {
-            name: xp.asarray(table, dtype=bias.dtype, device=device)
-            for name, table in self.vacuum_conditions().items()
-        }
         if "spin" in conditions:
             conditions["spin"] = conditions["spin"][:, None, :]
         descriptor = self.descriptor(
diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py
index 59c027bbf9..53547fd9d8 100644
--- a/deepmd/dpmodel/descriptor/dpa4.py
+++ b/deepmd/dpmodel/descriptor/dpa4.py
@@ -2426,6 +2426,10 @@ def dense_lower_supports_comm(self) -> bool:
         """
         return False
 
+    def supports_graph_vacuum(self) -> bool:
+        """The dense adapter uses the same graph core; native spin needs it."""
+        return True
+
     def uses_graph_lower(self) -> bool:
         """Whether this descriptor supports the sel-free NeighborGraph lower.
 
diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py
index e38d8374da..a367fe523c 100644
--- a/deepmd/dpmodel/descriptor/dpa4c.py
+++ b/deepmd/dpmodel/descriptor/dpa4c.py
@@ -2026,6 +2026,10 @@ def get_env_protection(self) -> float:
         """Return the direction regularization scale in Å."""
         return self._EPS
 
+    def supports_graph_vacuum(self) -> bool:
+        """Isolated references use the graph core without a dense sel capacity."""
+        return True
+
     def uses_graph_lower(self) -> bool:
         """Return whether graph-native lowering is supported."""
         return True
diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py
index b0598e9735..9ea052c368 100644
--- a/deepmd/dpmodel/descriptor/make_base_descriptor.py
+++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py
@@ -287,6 +287,18 @@ def graph_type_embedding_table(self) -> Any | None:
             """
             return None
 
+        def supports_graph_vacuum(self) -> bool:
+            """Whether isolated-atom references can use the graph interface.
+
+            Opting in requires ``call_graph`` to produce the same isolated
+            reference as the dense compatibility path, when one exists;
+            graph-only descriptors may also opt in. This is distinct from
+            ``uses_graph_lower``: some legacy descriptors retain different
+            padding semantics on their dense and graph paths, so switching
+            their reference would change the model's energy zero.
+            """
+            return False
+
         @abstractmethod
         def need_sorted_nlist_for_lower(self) -> bool:
             """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py
index e1afd50d50..2f794f35ec 100644
--- a/deepmd/pt_expt/utils/serialization.py
+++ b/deepmd/pt_expt/utils/serialization.py
@@ -1750,6 +1750,70 @@ def _trace_and_export(
         )
 
 
+@contextlib.contextmanager
+def _cpu_reference_descriptors(model: torch.nn.Module) -> Iterator[None]:
+    """Temporarily use CPU reference descriptors for the eager vacuum fold.
+
+    Deployment descriptors bind accelerator operators at construction, before
+    they move to CPU for fake tracing. Changing their device or the kernel
+    environment cannot unbind those operators. Rebuild only descriptors whose
+    fittings still need a reference, leaving the folded bias and non-persistent
+    vacuum table on the original fitting. Restore the deployment descriptors
+    before tracing so their optimized bindings remain intact.
+    """
+    from deepmd.dpmodel.atomic_model.dp_atomic_model import (
+        DPAtomicModel,
+    )
+    from deepmd.pt_expt.utils import (
+        env,
+    )
+
+    atomic_models = [
+        module
+        for module in model.modules()
+        if isinstance(module, DPAtomicModel)
+        and module.fitting_net.needs_vacuum_descriptor()
+    ]
+    if not atomic_models:
+        yield
+        return
+
+    levels = (
+        "DP_CUDA_INFER",
+        "DP_TRITON_INFER",
+        "DP_CUTILE_INFER",
+        "DP_CUTE_INFER",
+        "DP_CUDA_TRAIN",
+        "DP_TRITON_TRAIN",
+        "DP_AMP_INFER",
+    )
+    saved_levels = {name: os.environ.get(name) for name in levels}
+    saved_device = env.DEVICE
+    descriptors = []
+    try:
+        env.DEVICE = torch.device("cpu")
+        os.environ.update(dict.fromkeys(levels, "0"))
+        for atomic_model in atomic_models:
+            descriptor = atomic_model.descriptor
+            reference = type(descriptor).deserialize(descriptor.serialize()).eval()
+            # Virtual-spin backbones can disable this route at runtime; that
+            # choice is deliberately absent from the portable descriptor.
+            if not descriptor.uses_graph_lower() and reference.uses_graph_lower():
+                reference.disable_graph_lower()
+            descriptors.append((atomic_model, descriptor))
+            atomic_model.descriptor = reference
+        yield
+    finally:
+        for atomic_model, descriptor in descriptors:
+            atomic_model.descriptor = descriptor
+        env.DEVICE = saved_device
+        for name, value in saved_levels.items():
+            if value is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = value
+
+
 def _trace_and_export_impl(
     data: dict,
     model_json_override: dict | None = None,
@@ -1840,7 +1904,8 @@ def _trace_and_export_impl(
     # folded bias or vacuum table stays on that device with the rest of the
     # model.
     model.to("cpu")
-    model.fold_vacuum_reference()
+    with torch.no_grad(), _cpu_reference_descriptors(model):
+        model.fold_vacuum_reference()
 
     # Device-dependent Python branches resolve on the CPU tracing inputs, so
     # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2)
diff --git a/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py b/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
new file mode 100644
index 0000000000..0f4d1b48d3
--- /dev/null
+++ b/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Graph-native isolated references must not allocate a dense sel-sized list."""
+
+import dataclasses
+from unittest import (
+    mock,
+)
+
+import numpy as np
+import pytest
+
+from deepmd.dpmodel.atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.dpmodel.descriptor.dpa1 import (
+    DescrptDPA1,
+)
+from deepmd.dpmodel.descriptor.dpa4 import (
+    DescrptDPA4,
+)
+from deepmd.dpmodel.descriptor.dpa4c import (
+    DescrptDPA4C,
+)
+from deepmd.dpmodel.fitting import (
+    InvarFitting,
+)
+from deepmd.dpmodel.utils.neighbor_graph import (
+    graph_from_dense_quartet,
+)
+
+from ...dpa4_fixtures import (
+    jitter_zero_arrays,
+)
+
+TYPE_MAP = ["O", "H"]
+BIAS = [[-3.0], [0.5]]
+
+
+def make_model(
+    kind: str,
+    charge: bool = False,
+    spin: bool = False,
+    precision: str = "float64",
+) -> DPAtomicModel:
+    """Activate learned condition heads so missing inputs cannot pass by chance."""
+    if kind == "dpa1":
+        descriptor = DescrptDPA1(
+            rcut=4.0,
+            rcut_smth=0.5,
+            sel=[2],
+            ntypes=2,
+            attn_layer=0,
+            axis_neuron=2,
+            neuron=[4, 8],
+            precision=precision,
+            seed=3,
+        )
+    else:
+        cls = DescrptDPA4 if kind == "dpa4" else DescrptDPA4C
+        kwargs = {
+            "ntypes": 2,
+            "rcut": 4.0,
+            "channels": 8,
+            "n_radial": 4,
+            "lmax": 2,
+            "precision": precision,
+            "seed": 3,
+            "add_chg_spin_ebd": charge,
+            "default_chg_spin": [0, 1] if charge else None,
+            "use_spin": [True, False] if spin else None,
+        }
+        if kind == "dpa4":
+            kwargs.update(sel=8, mmax=1, n_blocks=1)
+        descriptor = cls(**kwargs)
+        descriptor = cls.deserialize(
+            jitter_zero_arrays(descriptor.serialize(), np.random.default_rng(7))
+        )
+        if kind == "dpa4c" and spin:
+            descriptor.spin.spin_gate[...] = 1.0
+    fitting = InvarFitting(
+        "energy",
+        2,
+        descriptor.get_dim_out(),
+        1,
+        neuron=[8],
+        mixed_types=descriptor.mixed_types(),
+        precision=precision,
+        vacuum_ref=True,
+        seed=1,
+    )
+    return DPAtomicModel(
+        descriptor,
+        fitting,
+        type_map=TYPE_MAP,
+        preset_out_bias={"energy": BIAS},
+    )
+
+
+def bounded_reference(model: DPAtomicModel):
+    """Use two explicitly bounded dense slots, never the descriptor's sel.
+
+    The native-spin interface is graph-only, so adapt the dense quartet and
+    pass spin on its flat node axis. Spin-less models also exercise the
+    descriptor's dense compatibility adapter directly.
+    """
+    dtype = model.fitting_net.bias_atom_e.dtype
+    coord = np.zeros((2, 1, 3), dtype=dtype)
+    atype = np.arange(2, dtype=np.int64).reshape(2, 1)
+    nlist = np.full((2, 1, 2), -1, dtype=np.int64)
+    mapping = np.zeros((2, 1), dtype=np.int64)
+    graph, flat_atype = graph_from_dense_quartet(coord, atype, nlist, mapping)
+    conditions = {
+        key: np.asarray(value, dtype=dtype)
+        for key, value in model.vacuum_conditions().items()
+    }
+    descriptor = model.descriptor
+    expected = descriptor.call_graph(
+        graph,
+        flat_atype,
+        type_embedding=descriptor.graph_type_embedding_table(),
+        **conditions,
+    )[0]
+    if not model.add_spin_ebd:
+        dense = descriptor(coord, atype, nlist, mapping=mapping, **conditions)[0]
+        np.testing.assert_allclose(dense.reshape(expected.shape), expected, atol=2e-6)
+    return expected, graph, flat_atype, conditions
+
+
+@pytest.mark.parametrize("kind", ["dpa4", "dpa4c"])
+@pytest.mark.parametrize(
+    "charge,spin", [(False, False), (True, False), (False, True), (True, True)]
+)
+@pytest.mark.parametrize("precision", ["float32", "float64"])
+def test_vacuum_uses_bounded_graph(kind, charge, spin, precision) -> None:
+    model = make_model(kind, charge, spin, precision)
+    expected, _, _, conditions = bounded_reference(model)
+    descriptor = model.descriptor
+    assert descriptor.supports_graph_vacuum()
+    if kind == "dpa4c":
+        assert descriptor.get_sel() == [999999]
+    call_graph = descriptor.call_graph
+
+    def checked_graph(graph, atype, **kwargs):
+        np.testing.assert_array_equal(graph.n_node, [1, 1])
+        np.testing.assert_array_equal(atype, [0, 1])
+        assert graph.edge_index.shape == (2, 2)
+        assert graph.edge_vec.shape == (2, 3)
+        assert graph.edge_mask.shape == (2,)
+        assert not graph.edge_mask.any()
+        assert graph.edge_vec.dtype == model.fitting_net.bias_atom_e.dtype
+        assert set(kwargs) == {"type_embedding", *conditions}
+        for name, value in conditions.items():
+            np.testing.assert_array_equal(kwargs[name], value)
+        return call_graph(graph, atype, **kwargs)
+
+    # Fail before allocation if the implementation attempts the dense route.
+    # Repeat with a larger sentinel to prove that the graph is sel-independent.
+    with (
+        mock.patch.object(
+            descriptor, "get_nnei", side_effect=AssertionError("dense capacity queried")
+        ),
+        mock.patch.object(
+            descriptor, "call", side_effect=AssertionError("dense adapter called")
+        ),
+        mock.patch.object(
+            descriptor, "call_graph", side_effect=checked_graph
+        ) as graph_call,
+    ):
+        for sentinel in (999999, 1999999):
+            with mock.patch.object(descriptor, "get_sel", return_value=[sentinel]):
+                actual = model.vacuum_descriptor()
+            assert actual.shape == (2, descriptor.get_dim_out())
+            assert actual.dtype == model.fitting_net.bias_atom_e.dtype
+            np.testing.assert_allclose(actual, expected, rtol=2e-6, atol=2e-6)
+        assert graph_call.call_count == 2
+    assert np.max(np.abs(actual)) > 1e-5
+    assert not np.allclose(actual[0], actual[1])
+
+
+@pytest.mark.parametrize("kind", ["dpa4", "dpa4c"])
+def test_empty_and_masked_edges_preserve_conditions(kind) -> None:
+    model = make_model(kind, charge=True, spin=True)
+    expected, dense_graph, atype, conditions = bounded_reference(model)
+    np.testing.assert_allclose(model.vacuum_descriptor(), expected, atol=1e-12)
+    for nedges in (0, 1, 7):
+        graph = dataclasses.replace(
+            dense_graph,
+            edge_index=np.tile(np.array([[1], [0]], dtype=np.int64), (1, nedges)),
+            edge_vec=np.ones((nedges, 3)),
+            edge_mask=np.zeros(nedges, dtype=bool),
+            destination_order=None,
+            destination_row_ptr=None,
+            source_order=None,
+            source_row_ptr=None,
+        )
+        actual = model.descriptor.call_graph(graph, atype, **conditions)[0]
+        np.testing.assert_allclose(actual, expected, atol=1e-12)
+    # The reference is the learned, ground-state-conditioned value, not zero
+    # or a default-condition shortcut, even when no real edges exist.
+    for name in conditions:
+        other = {**conditions, name: np.zeros_like(conditions[name])}
+        if name == "charge_spin":
+            other[name][:, 1] = 1.0
+        actual = model.descriptor.call_graph(dense_graph, atype, **other)[0]
+        assert np.max(np.abs(actual - expected)) > 1e-7, name
+
+
+def test_graph_descriptor_receives_its_type_embedding() -> None:
+    """An opted-in descriptor receives its external type-embedding table."""
+    model = make_model("dpa1")
+    assert not model.descriptor.supports_graph_vacuum()
+    expected, _, _, _ = bounded_reference(model)
+    with (
+        mock.patch.object(model.descriptor, "supports_graph_vacuum", return_value=True),
+        mock.patch.object(
+            model.descriptor, "call", side_effect=AssertionError("dense adapter called")
+        ),
+    ):
+        actual = model.vacuum_descriptor()
+    np.testing.assert_allclose(actual, expected, atol=1e-12)
+
+
+def test_dense_forward_keeps_its_own_reference_semantics() -> None:
+    """DPA1's legacy padding residual must cancel on its dense path too."""
+    model = make_model("dpa1")
+    model.descriptor.se_atten.mean[...] = 0.3
+    coord = np.zeros((2, 1, 3))
+    atype = np.arange(2, dtype=np.int64).reshape(2, 1)
+    nlist = np.full((2, 1, 2), -1, dtype=np.int64)
+    mapping = np.zeros((2, 1), dtype=np.int64)
+    assert model.descriptor.uses_graph_lower()
+    assert not model.descriptor.supports_graph_vacuum()
+    graph, flat_atype = graph_from_dense_quartet(coord, atype, nlist, mapping)
+    graph_reference = model.descriptor.call_graph(
+        graph,
+        flat_atype,
+        type_embedding=model.descriptor.graph_type_embedding_table(),
+    )[0]
+    dense_reference = model.vacuum_descriptor()
+    assert np.max(np.abs(graph_reference - dense_reference)) > 1e-7
+    before = model.forward_atomic(coord, atype, nlist, mapping=mapping)["energy"]
+    np.testing.assert_allclose(
+        before[:, 0], model.fitting_net.bias_atom_e, rtol=1e-12, atol=1e-12
+    )
+    # Graph support alone must not silently change the legacy reference.
+    model.fold_vacuum_reference()
+    after = model.forward_atomic(coord, atype, nlist, mapping=mapping)["energy"]
+    np.testing.assert_allclose(after, before, rtol=1e-12, atol=1e-12)
+
+
+def test_disabled_graph_lower_keeps_dense_reference() -> None:
+    """An explicit lower opt-out still selects the dense compatibility path."""
+    model = make_model("dpa4", charge=True)
+    expected = model.vacuum_descriptor()
+    model.descriptor.disable_graph_lower()
+    with mock.patch.object(
+        model.descriptor, "call", wraps=model.descriptor.call
+    ) as call:
+        actual = model.vacuum_descriptor()
+    call.assert_called_once()
+    np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize(
+    "kind,charge,spin",
+    [
+        ("dpa4", False, False),
+        ("dpa4", True, True),
+        ("dpa4c", False, False),
+        ("dpa4c", True, True),
+    ],
+)
+def test_folding_matches_appended_reference_frames(kind, charge, spin) -> None:
+    """Folding must match reference nodes carried alongside a real cluster."""
+    model = make_model(kind, charge, spin)
+    coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.2, 0.0]]])
+    atype = np.array([[0, 1]], dtype=np.int64)
+    nlist = np.array([[[1, -1], [0, -1]]], dtype=np.int64)
+    graph, atype = graph_from_dense_quartet(coord, atype, nlist, None)
+    conditions = {}
+    if charge:
+        conditions["charge_spin"] = np.array([[1.0, 2.0]])
+    if spin:
+        conditions["spin"] = np.array([[0.2, 0.4, 1.0], [0.0, 0.0, 0.0]])
+    before = model.forward_atomic_graph(graph, atype, **conditions)["energy"]
+    model.fold_vacuum_reference()
+    assert not model.fitting_net.needs_vacuum_descriptor()
+    after = model.forward_atomic_graph(graph, atype, **conditions)["energy"]
+    np.testing.assert_allclose(after, before, rtol=1e-12, atol=1e-12)
diff --git a/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py b/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py
new file mode 100644
index 0000000000..c87f890924
--- /dev/null
+++ b/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py
@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""The eager CPU vacuum fold must not execute deployment-only operators."""
+
+import copy
+import os
+
+import pytest
+import torch
+from torch.fx.experimental.proxy_tensor import (
+    make_fx,
+)
+
+from deepmd.pt_expt.kernels import (
+    autotune,
+)
+from deepmd.pt_expt.kernels.cuda.dpa4 import (
+    edge_radial,
+)
+from deepmd.pt_expt.model.get_model import (
+    get_model,
+)
+from deepmd.pt_expt.model.model import (
+    BaseModel,
+)
+from deepmd.pt_expt.utils import (
+    env,
+    serialization,
+)
+
+from ..model.test_dpa4_vacuum_ref import (
+    CONFIG,
+    PRESET,
+)
+
+LEVELS = (
+    "DP_CUDA_INFER",
+    "DP_TRITON_INFER",
+    "DP_CUTILE_INFER",
+    "DP_CUTE_INFER",
+    "DP_CUDA_TRAIN",
+    "DP_TRITON_TRAIN",
+    "DP_AMP_INFER",
+)
+
+
+@torch.library.custom_op(
+    "deepmd_test::vacuum_cuda_radial", mutates_args=(), device_types="cuda"
+)
+def cuda_radial(
+    edge_len: torch.Tensor, keep: torch.Tensor, freqs: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """A CUDA-only operator, including on test hosts without a CUDA build."""
+    raise AssertionError("the eager fold must never call the accelerator")
+
+
+@cuda_radial.register_fake
+def cuda_radial_fake(edge_len, keep, freqs):
+    return edge_len.new_empty(edge_len.shape), edge_len.new_empty(
+        (edge_len.shape[0], freqs.numel())
+    )
+
+
+def reference_model(monkeypatch, numb_fparam=0, numb_aparam=0, vacuum_ref=True):
+    monkeypatch.setattr(env, "DEVICE", torch.device("cpu"))
+    for name in LEVELS:
+        monkeypatch.setenv(name, "0")
+    config = copy.deepcopy(CONFIG)
+    config["descriptor"].update(precision="float32", compute_precision="float32")
+    config["fitting_net"].update(
+        precision="float32",
+        vacuum_ref=vacuum_ref,
+        numb_fparam=numb_fparam,
+        numb_aparam=numb_aparam,
+    )
+    config["preset_out_bias"] = PRESET
+    return get_model(config).eval()
+
+
+@pytest.mark.parametrize(("numb_fparam", "numb_aparam"), [(0, 0), (1, 0), (0, 1)])
+def test_cuda_export_folds_on_cpu_before_restoring_deployment_dispatch(
+    monkeypatch, numb_fparam, numb_aparam
+) -> None:
+    """Construction-time CUDA bindings survive an eager, no-grad reference fold."""
+    reference = reference_model(monkeypatch, numb_fparam, numb_aparam)
+    data = {"model": reference.serialize()}
+    reference.fold_vacuum_reference()
+    expected = reference.atomic_model.fitting_net
+
+    # Exercise the production constructor gate without requiring CUDA hardware
+    # or its extension. This op has a fake implementation but no CPU kernel,
+    # exactly the distinction between tracing and the eager vacuum evaluation.
+    monkeypatch.setattr(
+        edge_radial,
+        "make_cuda_edge_radial",
+        lambda envelope, basis: (
+            lambda distance, keep: cuda_radial(distance, keep, basis.adam_freqs)
+        ),
+    )
+    monkeypatch.setenv("DP_CUDA_INFER", "1")
+    model = BaseModel.deserialize(data["model"]).eval()
+    am = model.atomic_model
+    descriptor = am.descriptor
+    assert descriptor._cuda_radial_fn is not None
+    with pytest.raises(NotImplementedError, match="CPU"):
+        model.fold_vacuum_reference()
+
+    # A CUDA target is restored before the exporter tunes and traces. Setting
+    # it only after construction also makes this test work on CPU-only hosts.
+    monkeypatch.setattr(env, "DEVICE", torch.device("cuda"))
+    monkeypatch.setattr(BaseModel, "deserialize", lambda data: model)
+    events = []
+    move = model.to
+    fold = model.fold_vacuum_reference
+
+    def move_to_cpu(device):
+        events.append("move")
+        assert device == "cpu"
+        return move(device)
+
+    def fold_on_cpu():
+        events.append("fold")
+        assert env.DEVICE.type == "cpu"
+        assert not torch.is_grad_enabled()
+        assert am.descriptor is not descriptor
+        assert not am.descriptor.training
+        assert am.descriptor._cuda_radial_fn is None
+        assert all(t.device.type == "cpu" for t in model.parameters())
+        assert all(t.device.type == "cpu" for t in model.buffers())
+        fold()
+
+    class StopBeforeTrace(Exception):
+        pass
+
+    def check_before_trace(model, target):
+        events.append("trace")
+        assert target.type == env.DEVICE.type == "cuda"
+        assert torch.is_grad_enabled(), "force tracing still needs autograd"
+        assert am.descriptor is descriptor
+        assert os.environ["DP_CUDA_INFER"] == "1"
+        raise StopBeforeTrace
+
+    monkeypatch.setattr(model, "to", move_to_cpu)
+    monkeypatch.setattr(model, "fold_vacuum_reference", fold_on_cpu)
+    monkeypatch.setattr(autotune, "run_autotune", check_before_trace)
+    with pytest.raises(StopBeforeTrace):
+        serialization._trace_and_export_impl(data, target_device=env.DEVICE)
+    assert events == ["move", "fold", "trace"]
+
+    fitting = am.fitting_net
+    assert fitting.vacuum_ref == expected.vacuum_ref
+    torch.testing.assert_close(fitting.bias_atom_e, expected.bias_atom_e)
+    torch.testing.assert_close(fitting.vacuum_table, expected.vacuum_table)
+    if numb_fparam or numb_aparam:
+        assert fitting.vacuum_table.device.type == "cpu"
+        assert "vacuum_table" not in fitting.state_dict()
+    assert all(t.device.type == "cpu" for t in model.buffers())
+
+    # Restoring the descriptor must preserve optimized fake tracing, rather
+    # than silently exporting the reference implementation for the real atoms.
+    with torch.no_grad():
+        traced = make_fx(
+            lambda: am.vacuum_descriptor(),
+            tracing_mode="fake",
+            _allow_non_fake_inputs=True,
+        )()
+    assert any(node.target == cuda_radial._opoverload for node in traced.graph.nodes)
+
+
+@pytest.mark.parametrize("failure", ["construction", "fold"])
+def test_reference_context_restores_state_after_failure(monkeypatch, failure) -> None:
+    model = reference_model(monkeypatch)
+    am = model.atomic_model
+    descriptor = am.descriptor
+    monkeypatch.setattr(env, "DEVICE", torch.device("meta"))
+    for i, name in enumerate(LEVELS):
+        if i % 2:
+            monkeypatch.delenv(name)
+        else:
+            monkeypatch.setenv(name, "1")
+    saved = {name: os.environ.get(name) for name in LEVELS}
+
+    def fail(data):
+        raise RuntimeError("reference construction failed")
+
+    if failure == "construction":
+        monkeypatch.setattr(type(descriptor), "deserialize", fail)
+    with pytest.raises(RuntimeError, match="failed"):
+        with serialization._cpu_reference_descriptors(model):
+            assert am.descriptor is not descriptor
+            assert env.DEVICE.type == "cpu"
+            assert all(os.environ[name] == "0" for name in LEVELS)
+            raise RuntimeError("reference fold failed")
+    assert am.descriptor is descriptor
+    assert env.DEVICE.type == "meta"
+    assert {name: os.environ.get(name) for name in LEVELS} == saved
+
+
+@pytest.mark.parametrize("folded_table", [False, True])
+def test_unneeded_reference_does_not_rebuild_descriptor(
+    monkeypatch, folded_table
+) -> None:
+    model = reference_model(monkeypatch, numb_fparam=1, vacuum_ref=folded_table)
+    if folded_table:
+        model.fold_vacuum_reference()
+    descriptor = model.atomic_model.descriptor
+
+    def fail(data):
+        raise AssertionError("an unreferenced descriptor should not be rebuilt")
+
+    monkeypatch.setattr(type(descriptor), "deserialize", fail)
+    monkeypatch.setattr(env, "DEVICE", torch.device("meta"))
+    with serialization._cpu_reference_descriptors(model):
+        assert model.atomic_model.descriptor is descriptor
+        assert env.DEVICE.type == "meta"
+
+
+def test_reference_descriptor_preserves_disabled_graph_route(monkeypatch) -> None:
+    """Virtual-spin backbones disable graph lowering outside serialization."""
+    monkeypatch.setattr(env, "DEVICE", torch.device("cpu"))
+    for name in LEVELS:
+        monkeypatch.setenv(name, "0")
+    model = get_model(
+        {
+            "type_map": ["O", "H"],
+            "descriptor": {
+                "type": "dpa1",
+                "rcut": 4.0,
+                "rcut_smth": 0.5,
+                "sel": [8],
+                "attn_layer": 0,
+                "axis_neuron": 2,
+                "neuron": [4, 8],
+                "seed": 1,
+            },
+            "fitting_net": {
+                "type": "ener",
+                "neuron": [8],
+                "vacuum_ref": True,
+                "seed": 1,
+            },
+            "preset_out_bias": PRESET,
+        }
+    ).eval()
+    am = model.atomic_model
+    descriptor = am.descriptor
+    descriptor.disable_graph_lower()
+    with torch.no_grad():
+        descriptor.se_atten.mean.fill_(0.3)
+        expected = am.vacuum_descriptor()
+        with serialization._cpu_reference_descriptors(model):
+            assert not am.descriptor.uses_graph_lower()
+            torch.testing.assert_close(am.vacuum_descriptor(), expected)
+    assert am.descriptor is descriptor

From 1cfa25afc23ee4a1b46573ce2cea50cdde400703 Mon Sep 17 00:00:00 2001
From: OutisLi <137472077+OutisLi@users.noreply.github.com>
Date: Thu, 15 Oct 2026 17:59:06 +0800
Subject: [PATCH 3/5] style(tests): use canonical NumPy capitalization

---
 source/tests/pt_expt/utils/test_common.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/source/tests/pt_expt/utils/test_common.py b/source/tests/pt_expt/utils/test_common.py
index f0fa2c524f..856144d413 100644
--- a/source/tests/pt_expt/utils/test_common.py
+++ b/source/tests/pt_expt/utils/test_common.py
@@ -130,7 +130,7 @@ class MockModule(MockNativeOP):
 
 
 def test_dpmodel_setattr_keeps_module_device(monkeypatch) -> None:
-    """Numpy buffer assignment after Module.to keeps the module device.
+    """NumPy buffer assignment after Module.to keeps the module device.
 
     ``env.DEVICE`` is patched to ``meta`` so a regression that always places
     new buffers on the process default would land off-CPU and fail the assert.

From fd0a9cdf3951ba39efde9b6142dc7b29fa6920a0 Mon Sep 17 00:00:00 2001
From: OutisLi <137472077+OutisLi@users.noreply.github.com>
Date: Sat, 17 Oct 2026 11:18:38 +0800
Subject: [PATCH 4/5] test(vacuum): tighten precision and cover buffer device
 fallbacks

Use precision-specific dense/graph comparison tolerances without loosening
float32 checks. Cover new buffers inheriting parameter or buffer devices
with CPU/meta distinctions, and document export test parametrizations.
---
 .../dpmodel/test_vacuum_descriptor_graph.py   |  8 ++++--
 source/tests/pt_expt/utils/test_common.py     | 27 +++++++++++++++++++
 .../utils/test_serialization_vacuum_ref.py    |  8 +++---
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py b/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
index 0f4d1b48d3..0b4a2ba6c7 100644
--- a/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
+++ b/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
@@ -122,7 +122,10 @@ def bounded_reference(model: DPAtomicModel):
     )[0]
     if not model.add_spin_ebd:
         dense = descriptor(coord, atype, nlist, mapping=mapping, **conditions)[0]
-        np.testing.assert_allclose(dense.reshape(expected.shape), expected, atol=2e-6)
+        rtol, atol = (1e-7, 2e-6) if dtype == np.float32 else (1e-12, 1e-12)
+        np.testing.assert_allclose(
+            dense.reshape(expected.shape), expected, rtol=rtol, atol=atol
+        )
     return expected, graph, flat_atype, conditions
 
 
@@ -132,6 +135,7 @@ def bounded_reference(model: DPAtomicModel):
 )
 @pytest.mark.parametrize("precision", ["float32", "float64"])
 def test_vacuum_uses_bounded_graph(kind, charge, spin, precision) -> None:
+    tol = 2e-6 if precision == "float32" else 1e-12
     model = make_model(kind, charge, spin, precision)
     expected, _, _, conditions = bounded_reference(model)
     descriptor = model.descriptor
@@ -171,7 +175,7 @@ def checked_graph(graph, atype, **kwargs):
                 actual = model.vacuum_descriptor()
             assert actual.shape == (2, descriptor.get_dim_out())
             assert actual.dtype == model.fitting_net.bias_atom_e.dtype
-            np.testing.assert_allclose(actual, expected, rtol=2e-6, atol=2e-6)
+            np.testing.assert_allclose(actual, expected, rtol=tol, atol=tol)
         assert graph_call.call_count == 2
     assert np.max(np.abs(actual)) > 1e-5
     assert not np.allclose(actual[0], actual[1])
diff --git a/source/tests/pt_expt/utils/test_common.py b/source/tests/pt_expt/utils/test_common.py
index 856144d413..179426c929 100644
--- a/source/tests/pt_expt/utils/test_common.py
+++ b/source/tests/pt_expt/utils/test_common.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
 import numpy as np
+import pytest
 import torch
 
 from deepmd.dpmodel.common import (
@@ -159,3 +160,29 @@ class Holder(HolderDP):
     module.stat = np.ones(3, dtype=np.float64)
     assert module.stat.device.type == "cpu"
     np.testing.assert_array_equal(module.stat.detach().cpu().numpy(), np.ones(3))
+
+
+@pytest.mark.parametrize("resident", ["parameter", "buffer"])  # device source
+def test_dpmodel_setattr_uses_other_module_tensors(monkeypatch, resident) -> None:
+    """A new NumPy buffer follows an existing parameter or other buffer."""
+
+    @torch_module
+    class Holder(NativeOP):
+        def call(self, *args, **kwargs) -> None:
+            pass
+
+    monkeypatch.setattr(env, "DEVICE", torch.device("cpu"))
+    module = Holder()
+    tensor = torch.zeros(3, device="cpu")
+    if resident == "parameter":
+        module.weight = torch.nn.Parameter(tensor)
+    else:
+        module.register_buffer("other", tensor)
+        assert not list(module.parameters())
+    module.to("meta")
+
+    assert "stat" not in module._buffers
+    module.stat = np.ones(3, dtype=np.float64)
+    assert module.stat.device.type == "meta"
+    assert module.stat.dtype == torch.float64
+    assert module.stat.shape == (3,)
diff --git a/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py b/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py
index c87f890924..37a42aa0ec 100644
--- a/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py
+++ b/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py
@@ -76,7 +76,9 @@ def reference_model(monkeypatch, numb_fparam=0, numb_aparam=0, vacuum_ref=True):
     return get_model(config).eval()
 
 
-@pytest.mark.parametrize(("numb_fparam", "numb_aparam"), [(0, 0), (1, 0), (0, 1)])
+@pytest.mark.parametrize(
+    ("numb_fparam", "numb_aparam"), [(0, 0), (1, 0), (0, 1)]
+)  # frame/atomic parameter dimensions
 def test_cuda_export_folds_on_cpu_before_restoring_deployment_dispatch(
     monkeypatch, numb_fparam, numb_aparam
 ) -> None:
@@ -166,7 +168,7 @@ def check_before_trace(model, target):
     assert any(node.target == cuda_radial._opoverload for node in traced.graph.nodes)
 
 
-@pytest.mark.parametrize("failure", ["construction", "fold"])
+@pytest.mark.parametrize("failure", ["construction", "fold"])  # failing stage
 def test_reference_context_restores_state_after_failure(monkeypatch, failure) -> None:
     model = reference_model(monkeypatch)
     am = model.atomic_model
@@ -195,7 +197,7 @@ def fail(data):
     assert {name: os.environ.get(name) for name in LEVELS} == saved
 
 
-@pytest.mark.parametrize("folded_table", [False, True])
+@pytest.mark.parametrize("folded_table", [False, True])  # disabled vs already folded
 def test_unneeded_reference_does_not_rebuild_descriptor(
     monkeypatch, folded_table
 ) -> None:

From 8ac623c4a94e39df270ceaec2141f6dca6d98d94 Mon Sep 17 00:00:00 2001
From: OutisLi <137472077+OutisLi@users.noreply.github.com>
Date: Sun, 18 Oct 2026 14:39:31 +0800
Subject: [PATCH 5/5] fix(vacuum): use dense references for legacy graph
 descriptors

Append isolated reference frames only for descriptors that opt into graph
vacuum evaluation. DPA1 and DPA2 now use the same differentiable dense
reference in graph inference and folding; DPA4/C keep the shared graph call.

Cover nonzero normalization, unconditioned and fparam/aparam folding, and
parameter-gradient parity. Eight regressions fail numerically before the
fix; 80 related tests pass afterward, with two GPU tests skipped.
---
 .../dpmodel/atomic_model/dp_atomic_model.py   | 11 ++-
 .../dpmodel/test_vacuum_descriptor_graph.py   | 99 +++++++++++++++++++
 .../pt_expt/model/test_legacy_vacuum_ref.py   | 61 ++++++++++++
 3 files changed, 167 insertions(+), 4 deletions(-)
 create mode 100644 source/tests/pt_expt/model/test_legacy_vacuum_ref.py

diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py
index b78ec05e24..b1e0e98af9 100644
--- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py
@@ -636,10 +636,11 @@ def forward_atomic_graph(
                 atype.shape[0],
             )
         # === Step 2. Vacuum reference nodes ===
-        # The reference atoms are appended as single-node frames, so the same
-        # descriptor call yields the vacuum descriptor of every type.
+        # Opted-in descriptors evaluate the reference in the same graph call.
+        # Legacy descriptors must retain the dense reference used by folding.
         vacuum_ref = self.fitting_net.needs_vacuum_descriptor()
-        if vacuum_ref:
+        graph_vacuum = vacuum_ref and self.descriptor.supports_graph_vacuum()
+        if graph_vacuum:
             graph, atype, charge_spin, spin = self.append_vacuum_frames(
                 graph, atype, charge_spin, spin
             )
@@ -668,12 +669,14 @@ def forward_atomic_graph(
         gg, rot_mat = outputs[0], outputs[1]
         node_gate = outputs[2] if len(outputs) > 2 else None
         vacuum_kwargs = {}
-        if vacuum_ref:
+        if graph_vacuum:
             vacuum_kwargs["vacuum_descriptor"] = gg[n_real:]
             gg = gg[:n_real]
             atype = atype[:n_real]
             rot_mat = None if rot_mat is None else rot_mat[:n_real]
             node_gate = None if node_gate is None else node_gate[:n_real]
+        elif vacuum_ref:
+            vacuum_kwargs["vacuum_descriptor"] = self.vacuum_descriptor()
         # === Step 4. Fitting ===
         # The source gate of a bridged descriptor fades the learned part of
         # the fitting output: an atom of a frozen pair keeps only the bias of
diff --git a/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py b/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
index 0b4a2ba6c7..365185f8e7 100644
--- a/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
+++ b/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py
@@ -15,6 +15,9 @@
 from deepmd.dpmodel.descriptor.dpa1 import (
     DescrptDPA1,
 )
+from deepmd.dpmodel.descriptor.dpa2 import (
+    DescrptDPA2,
+)
 from deepmd.dpmodel.descriptor.dpa4 import (
     DescrptDPA4,
 )
@@ -41,6 +44,8 @@ def make_model(
     charge: bool = False,
     spin: bool = False,
     precision: str = "float64",
+    numb_fparam: int = 0,
+    numb_aparam: int = 0,
 ) -> DPAtomicModel:
     """Activate learned condition heads so missing inputs cannot pass by chance."""
     if kind == "dpa1":
@@ -55,6 +60,29 @@ def make_model(
             precision=precision,
             seed=3,
         )
+    elif kind == "dpa2":
+        descriptor = DescrptDPA2(
+            ntypes=2,
+            repinit={
+                "rcut": 4.0,
+                "rcut_smth": 0.5,
+                "nsel": 2,
+                "neuron": [4, 8],
+                "axis_neuron": 2,
+            },
+            repformer={
+                "rcut": 3.0,
+                "rcut_smth": 0.5,
+                "nsel": 1,
+                "nlayers": 1,
+                "g1_dim": 8,
+                "g2_dim": 4,
+                "axis_neuron": 2,
+                "update_g1_has_attn": False,
+            },
+            precision=precision,
+            seed=3,
+        )
     else:
         cls = DescrptDPA4 if kind == "dpa4" else DescrptDPA4C
         kwargs = {
@@ -86,6 +114,8 @@ def make_model(
         mixed_types=descriptor.mixed_types(),
         precision=precision,
         vacuum_ref=True,
+        numb_fparam=numb_fparam,
+        numb_aparam=numb_aparam,
         seed=1,
     )
     return DPAtomicModel(
@@ -291,3 +321,72 @@ def test_folding_matches_appended_reference_frames(kind, charge, spin) -> None:
     assert not model.fitting_net.needs_vacuum_descriptor()
     after = model.forward_atomic_graph(graph, atype, **conditions)["energy"]
     np.testing.assert_allclose(after, before, rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize("kind", ["dpa1", "dpa2"])  # legacy descriptor
+@pytest.mark.parametrize(
+    "numb_fparam,numb_aparam", [(0, 0), (1, 0), (0, 1)]
+)  # bias folding or conditioned reference table
+def test_legacy_graph_uses_dense_vacuum_reference(
+    kind, numb_fparam, numb_aparam
+) -> None:
+    """Graph inference and folding must subtract the same legacy reference."""
+    model = make_model(kind, numb_fparam=numb_fparam, numb_aparam=numb_aparam)
+    descriptor = model.descriptor
+    blocks = (
+        [descriptor.se_atten]
+        if kind == "dpa1"
+        else [descriptor.repinit, descriptor.repformers]
+    )
+    for block in blocks:
+        block.mean[...] = 0.3
+    assert descriptor.uses_graph_lower()
+    assert not descriptor.supports_graph_vacuum()
+    iso_coord = np.zeros((2, 1, 3))
+    iso_atype = np.arange(2, dtype=np.int64).reshape(2, 1)
+    iso_nlist = np.full((2, 1, 2), -1, dtype=np.int64)
+    iso_graph, iso_atype = graph_from_dense_quartet(
+        iso_coord, iso_atype, iso_nlist, None
+    )
+    graph_reference = descriptor.call_graph(
+        iso_graph, iso_atype, type_embedding=descriptor.graph_type_embedding_table()
+    )[0]
+    dense_reference = model.vacuum_descriptor()
+    assert np.max(np.abs(graph_reference - dense_reference)) > 1e-7
+
+    coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.2, 0.0]]])
+    atype = np.array([[0, 1]], dtype=np.int64)
+    nlist = np.array([[[1, -1], [0, -1]]], dtype=np.int64)
+    graph, atype = graph_from_dense_quartet(coord, atype, nlist, None)
+    fparam = np.array([[0.7]]) if numb_fparam else None
+    aparam = np.array([[0.2], [0.6]]) if numb_aparam else None
+    gg = descriptor.call_graph(
+        graph, atype, type_embedding=descriptor.graph_type_embedding_table()
+    )[0]
+    expected = model.fitting_net.call_graph(
+        gg,
+        atype,
+        fparam=np.repeat(fparam, 2, axis=0) if numb_fparam else None,
+        aparam=aparam,
+        vacuum_descriptor=dense_reference,
+    )["energy"]
+    with (
+        mock.patch.object(
+            model, "append_vacuum_frames", wraps=model.append_vacuum_frames
+        ) as append,
+        mock.patch.object(
+            model, "vacuum_descriptor", wraps=model.vacuum_descriptor
+        ) as reference,
+    ):
+        before = model.forward_atomic_graph(graph, atype, fparam=fparam, aparam=aparam)[
+            "energy"
+        ]
+    np.testing.assert_allclose(before, expected, rtol=1e-12, atol=1e-12)
+    append.assert_not_called()
+    reference.assert_called_once()
+    model.fold_vacuum_reference()
+    assert not model.fitting_net.needs_vacuum_descriptor()
+    after = model.forward_atomic_graph(graph, atype, fparam=fparam, aparam=aparam)[
+        "energy"
+    ]
+    np.testing.assert_allclose(after, before, rtol=1e-12, atol=1e-12)
diff --git a/source/tests/pt_expt/model/test_legacy_vacuum_ref.py b/source/tests/pt_expt/model/test_legacy_vacuum_ref.py
new file mode 100644
index 0000000000..e6aea5c55b
--- /dev/null
+++ b/source/tests/pt_expt/model/test_legacy_vacuum_ref.py
@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Legacy vacuum references remain differentiable on the graph route."""
+
+import pytest
+import torch
+
+from deepmd.dpmodel.atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.dpmodel.utils.neighbor_graph import (
+    graph_from_dense_quartet,
+)
+from deepmd.pt_expt.common import (
+    auto_wrapped_class,
+)
+from deepmd.pt_expt.utils import (
+    env,
+)
+
+from ...common.dpmodel.test_vacuum_descriptor_graph import (
+    make_model,
+)
+
+
+@pytest.mark.parametrize("kind", ["dpa1", "dpa2"])  # legacy descriptor
+def test_graph_vacuum_reference_preserves_parameter_gradients(kind) -> None:
+    model = auto_wrapped_class(DPAtomicModel).deserialize(make_model(kind).serialize())
+    descriptor = model.descriptor
+    blocks = (
+        [descriptor.se_atten]
+        if kind == "dpa1"
+        else [descriptor.repinit, descriptor.repformers]
+    )
+    with torch.no_grad():
+        for block in blocks:
+            block.mean.fill_(0.3)
+    coord = torch.tensor(
+        [[[0.0, 0.0, 0.0], [1.0, 0.2, 0.0]]],
+        dtype=torch.float64,
+        device=env.DEVICE,
+    )
+    atype = torch.tensor([[0, 1]], dtype=torch.int64, device=env.DEVICE)
+    nlist = torch.tensor([[[1, -1], [0, -1]]], dtype=torch.int64, device=env.DEVICE)
+    graph, atype = graph_from_dense_quartet(coord, atype, nlist, None)
+    gg = descriptor.call_graph(
+        graph, atype, type_embedding=descriptor.graph_type_embedding_table()
+    )[0]
+    reference = model.vacuum_descriptor()
+    parameters = tuple(p for p in model.parameters() if p.requires_grad)
+    reference_grad = torch.autograd.grad(
+        reference.sum(), parameters, allow_unused=True, retain_graph=True
+    )
+    assert any(g is not None and torch.any(g.abs() > 1e-7) for g in reference_grad)
+    expected = model.fitting_net.call_graph(gg, atype, vacuum_descriptor=reference)[
+        "energy"
+    ]
+    actual = model.forward_atomic_graph(graph, atype)["energy"]
+    torch.testing.assert_close(actual, expected, rtol=1e-12, atol=1e-12)
+    expected_grad = torch.autograd.grad(expected.sum(), parameters, allow_unused=True)
+    actual_grad = torch.autograd.grad(actual.sum(), parameters, allow_unused=True)
+    torch.testing.assert_close(actual_grad, expected_grad, rtol=1e-12, atol=1e-12)