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] 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))