fix(pt_expt): fold vacuum reference on CPU tracing device (#6076)

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

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)

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)

source/tests/pt_expt/utils/test_common.py

@@ -127,3 +127,35 @@
     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))