diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index b9f20ed8c7..b1e0e98af9 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -413,11 +413,11 @@ def vacuum_descriptor(self) -> Array: conditioned as the neutral ground-state atom: zero charge with the ground-state multiplicity when the descriptor takes the charge/spin condition, and a spin vector of one Bohr magneton per unpaired - electron when it takes the native spin. A descriptor with a graph - lower is evaluated on the neighbor graph of those single-atom frames, - which holds no real edge; a dense descriptor on a neighbor list of its - own capacity filled with padding, since its equations are sized by - that capacity. + 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 ------- @@ -429,39 +429,36 @@ def vacuum_descriptor(self) -> Array: xp = array_api_compat.array_namespace(bias) device = array_api_compat.device(bias) ntypes = self.get_ntypes() - 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.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.descriptor.uses_graph_lower(): + if self.uses_graph_lower() and self.descriptor.supports_graph_vacuum(): from deepmd.dpmodel.utils.neighbor_graph import ( - build_neighbor_graph, + NeighborGraph, ) - graph = build_neighbor_graph( - coord, atype, None, self.descriptor.get_rcut(), with_csr=True - ) - # Only descriptors whose ``call_graph`` declares a condition take - # the keyword; a declared condition the model does not use is None. - spin_kwargs = ( - {"spin": conditions.get("spin")} if self.supports_native_spin() else {} - ) - charge_spin_kwargs = ( - {"charge_spin": conditions.get("charge_spin")} - if self.supports_charge_spin - else {} + # 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, - xp.reshape(atype, (ntypes,)), + atype, type_embedding=self.descriptor.graph_type_embedding_table(), - **spin_kwargs, - **charge_spin_kwargs, + **conditions, )[0] + + coord = xp.zeros((ntypes, 1, 3), dtype=bias.dtype, device=device) + atype = xp.reshape(atype, (ntypes, 1)) nlist = xp.full( (ntypes, 1, self.descriptor.get_nnei()), -1, @@ -639,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 ) @@ -671,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/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 7e264fb3e9..89e47669e4 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -2512,6 +2512,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 7dbd9c7f3d..2dc880da4f 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -2118,6 +2118,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/common.py b/deepmd/pt_expt/common.py index 1c19f7b1b0..b91c24d448 100644 --- a/deepmd/pt_expt/common.py +++ b/deepmd/pt_expt/common.py @@ -273,6 +273,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. @@ -319,6 +341,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 @@ -328,10 +354,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:** @@ -354,11 +380,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 b3edf999ad..22f5a2a494 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, @@ -1833,12 +1897,15 @@ 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") + 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..365185f8e7 --- /dev/null +++ b/source/tests/common/dpmodel/test_vacuum_descriptor_graph.py @@ -0,0 +1,392 @@ +# 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.dpa2 import ( + DescrptDPA2, +) +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", + numb_fparam: int = 0, + numb_aparam: int = 0, +) -> 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, + ) + 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 = { + "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, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + 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] + 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 + + +@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: + tol = 2e-6 if precision == "float32" else 1e-12 + 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=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]) + + +@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) + + +@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_dpa4_vacuum_ref.py b/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py index 3a4ee60793..6e3923453b 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/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) diff --git a/source/tests/pt_expt/utils/test_common.py b/source/tests/pt_expt/utils/test_common.py index 57ee153e54..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 ( @@ -127,3 +128,61 @@ 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)) + + +@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 new file mode 100644 index 0000000000..98bb6e5ee9 --- /dev/null +++ b/source/tests/pt_expt/utils/test_serialization_vacuum_ref.py @@ -0,0 +1,255 @@ +# 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") + 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)] +) # frame/atomic parameter dimensions +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"]) # failing stage +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]) # disabled vs already folded +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