diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py index 4c63d1a492..dca7d7a36f 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py @@ -99,6 +99,37 @@ def _build_frame_degree_index( raise ValueError("`coefficient_layout` must be either 'packed' or 'm_major'") +def _degree_batched_matmul(xp: Any, coeff: Any, weight: Any) -> Any: + """Contract ``einsum("ndfi,dio->ndfo")`` batched over the degree axis. + + Parameters + ---------- + xp : Any + The array namespace of ``coeff``. + coeff : Array + Coefficients with shape ``(N, D, F, i)``. + weight : Array + Per-degree weights with shape ``(D, i, o)``. + + Returns + ------- + Array + Contracted coefficients with shape ``(N, D, F, o)``. + + Notes + ----- + Batching over the ``(D, F)`` axes, not over ``N``: expanding ``weight`` + across ``F`` costs ``D*F*i*o`` elements, whereas batching over ``N`` + (or collapsing ``N*F``, which needs a materialized permuted copy of + ``coeff``) touches ``N*D*F*i`` elements — a factor ``N/o`` more. No + reshape is involved, so an empty ``N`` batch (empty graph/edge set, or + a distributed rank owning no nodes) flows through naturally. + """ + coeff_df = xp.permute_dims(coeff, (1, 2, 0, 3)) # (D, F, N, i) + out = xp.matmul(coeff_df, weight[:, None, :, :]) # (D, F, N, o) + return xp.permute_dims(out, (2, 0, 1, 3)) # (N, D, F, o) + + def _project_frames(coeff: Any, proj: ChannelLinear, n_frames: int) -> Any: """ Apply a channel-only linear map to each Wigner-D frame independently. @@ -493,9 +524,8 @@ def call(self, coeff: Any) -> Any: weight = xp_asarray_nodetach(xp, self.weight[...], device=device) degree_index = xp_asarray_nodetach(xp, self.degree_index, device=device) weight = xp.take(weight, degree_index, axis=0) - # einsum "ndfi,dio->ndfo" as a broadcast batched matmul: - # (N, D, F, i) @ (1, D, i, o) -> (N, D, F, o) - return xp.matmul(coeff, weight[None, ...]) + # Batched over the (D, F) axes, never over N -- see the helper's note. + return _degree_batched_matmul(xp, coeff, weight) def serialize(self) -> dict[str, Any]: """Serialize the FrameContract to a dict.""" @@ -575,9 +605,8 @@ def call(self, coeff: Any) -> Any: weight = xp_asarray_nodetach(xp, self.weight[...], device=device) degree_index = xp_asarray_nodetach(xp, self.degree_index, device=device) weight = xp.take(weight, degree_index, axis=0) - # einsum "ndfi,dio->ndfo" as a broadcast batched matmul: - # (N, D, F, i) @ (1, D, i, o) -> (N, D, F, o) - return xp.matmul(coeff, weight[None, ...]) + # Batched over the (D, F) axes, never over N -- see the helper's note. + return _degree_batched_matmul(xp, coeff, weight) def serialize(self) -> dict[str, Any]: """Serialize the FrameExpand to a dict.""" diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/lora.py b/deepmd/dpmodel/descriptor/dpa4_nn/lora.py index a4600b7dbe..ba8dd1ab1a 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/lora.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/lora.py @@ -189,10 +189,12 @@ def call(self, x: Array) -> Array: ) expand_index = xp_asarray_nodetach(xp, self.expand_index, device=device) weight_expanded = xp.take(weight, expand_index, axis=0) - # einsum "ndfi,difo->ndfo" as a broadcast batched matmul: - # (N, D, F, 1, Cin) @ (1, D, F, Cin, Cout) -> (N, D, F, 1, Cout) + # einsum "ndfi,difo->ndfo", batched over the small (D, F) axes rather + # than over N, which would broadcast the weight and make autograd + # reduce the expansion. LoRA twin of the so3.py contraction. weight_expanded = xp.permute_dims(weight_expanded, (0, 2, 1, 3)) - out = xp.matmul(x[:, :, :, None, :], weight_expanded[None, ...])[..., 0, :] + out = xp.matmul(xp.permute_dims(x, (1, 2, 0, 3)), weight_expanded) + out = xp.permute_dims(out, (2, 0, 1, 3)) # (N, D, F, Cout) if self.mlp_bias: bias = xp.reshape( xp_asarray_nodetach(xp, self.bias[...], device=device), diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py index 8ca2dbc855..a5f1694942 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py @@ -131,10 +131,13 @@ def call(self, x: Any) -> Any: xp, self.weight[...], device=array_api_compat.device(x) ) weight = xp.reshape(weight, (self.in_channels, self.n_focus, self.out_channels)) - # einsum "bfi,ifo->bfo" as a broadcast batched matmul: - # (B, F, 1, Cin) @ (1, F, Cin, Cout) -> (B, F, 1, Cout) + # einsum "bfi,ifo->bfo" as F independent (B, Cin) x (Cin, Cout) GEMMs. + # B stays the GEMM rows so the weight is used in place; making B the + # batch axis would broadcast it to (B, F, Cin, Cout) and leave autograd + # reducing that expansion. At n_focus=1 both permutes are free views. weight = xp.permute_dims(weight, (1, 0, 2)) # (F, Cin, Cout) - out = xp.matmul(x[:, :, None, :], weight[None, ...])[..., 0, :] + out = xp.matmul(xp.permute_dims(x, (1, 0, 2)), weight) # (F, B, Cout) + out = xp.permute_dims(out, (1, 0, 2)) # (B, F, Cout) if self.use_bias: bias = xp_asarray_nodetach( xp, self.bias[...], device=array_api_compat.device(x) @@ -439,12 +442,14 @@ def call(self, x: Any) -> Any: weight_expanded = xp.take(weight, expand_index, axis=0) # (D, Cin, F, Cout) # === Step 2. Per-focus, per-degree channel mixing === - # einsum "ndfi,difo->ndfo" as a broadcast batched matmul: - # (N, D, F, 1, Cin) @ (1, D, F, Cin, Cout) -> (N, D, F, 1, Cout) + # einsum "ndfi,difo->ndfo". Batch over (D, F) so N remains the GEMM + # row dimension: this avoids materializing N copies of the weight and + # the corresponding gradient reduction on every backward. weight_expanded = xp.permute_dims( weight_expanded, (0, 2, 1, 3) ) # (D, F, Cin, Cout) - out = xp.matmul(x[:, :, :, None, :], weight_expanded[None, ...])[..., 0, :] + out = xp.matmul(xp.permute_dims(x, (1, 2, 0, 3)), weight_expanded) + out = xp.permute_dims(out, (2, 0, 1, 3)) # (N, D, F, Cout) # === Step 3. Add l=0 bias === if self.mlp_bias: diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index f0442828ae..d570e12eb5 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -135,6 +135,7 @@ def get_linear_atomic_model( backend_name: str, atomic_model: type, pairtab_model: type, + linear_atomic_model: type | None = None, descriptor_child_builder: "Callable[[dict], Any | None] | None" = None, ) -> Any: """Build the ``LinearEnergyAtomicModel`` composition from a config. @@ -162,6 +163,13 @@ def get_linear_atomic_model( Backend learned atomic-model class. pairtab_model : type Backend pair-tabulation atomic-model class. + linear_atomic_model : type, optional + Backend linear composition atomic-model class. Defaults to the + dpmodel class. A backend that wraps dpmodel classes must pass its + own wrapper: otherwise the composition it gets back is a dpmodel + instance that its model wrapper has to convert, and conversion + keeps only what the portable record carries -- dropping any + runtime state the children hold (e.g. ``use_amp``). descriptor_child_builder : callable, optional Backend hook for descriptor-bearing children: called with the child config (``type_map`` and derived clamp radii already @@ -183,9 +191,11 @@ def get_linear_atomic_model( InnerPotentialAtomicModel, ) from deepmd.dpmodel.atomic_model.linear_atomic_model import ( - LinearEnergyAtomicModel, + LinearEnergyAtomicModel as LinearEnergyAtomicModelDP, ) + LinearEnergyAtomicModel = linear_atomic_model or LinearEnergyAtomicModelDP + data = copy.deepcopy(data) type_map = data["type_map"] children = data["models"] @@ -433,6 +443,7 @@ def __init__( atomic_model: type | None = None, pairtab_model: type | None = None, zbl_model: type | None = None, + linear_atomic_model: type | None = None, ) -> None: """Store backend-native classes used by all model construction paths.""" self.descriptor_base = descriptor_base @@ -442,6 +453,7 @@ def __init__( self.atomic_model = atomic_model self.pairtab_model = pairtab_model self.zbl_model = zbl_model + self.linear_atomic_model = linear_atomic_model def get_model_components(self, data: dict) -> tuple[Any, Any, str]: """Construct descriptor and fitting objects for this backend.""" @@ -478,6 +490,7 @@ def get_linear_atomic_model( backend_name=self.backend_name, atomic_model=self.atomic_model, pairtab_model=self.pairtab_model, + linear_atomic_model=self.linear_atomic_model, descriptor_child_builder=descriptor_child_builder, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py index 7c36e99e7b..a2e5efdd4f 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py +++ b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py @@ -343,6 +343,20 @@ def forward( return _project_frames(from_grid(out), self.out_proj, self.n_frames) +def _degree_batched_matmul(coeff: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Contract ``einsum("ndfi,dio->ndfo", coeff, weight)``. + + Batched over the ``(D, F)`` axes, not over ``N`` (and not by collapsing + ``N*F``, which would materialize a permuted copy of ``coeff``): + expanding ``weight`` across ``F`` costs ``D*F*i*o`` elements versus + ``N*D*F*i`` for the coefficient copy -- a factor ``N/o`` more. No + reshape is involved, so an empty ``N`` batch flows through naturally. + """ + coeff_df = coeff.permute(1, 2, 0, 3) # (D, F, N, i) + out = torch.matmul(coeff_df, weight.unsqueeze(1)) # (D, F, N, o) + return out.permute(2, 0, 1, 3) # (N, D, F, o) + + class FrameContract(nn.Module): """Per-degree frame/channel contraction that preserves the order index.""" @@ -387,7 +401,7 @@ def __init__( def forward(self, coeff: torch.Tensor) -> torch.Tensor: """Contract ``(N, D, F, K*C)`` frame coefficients to ``(N, D, F, C)``.""" weight = self.weight.index_select(0, self.degree_index) - return torch.einsum("ndfi,dio->ndfo", coeff, weight) + return _degree_batched_matmul(coeff, weight) class FrameExpand(nn.Module): @@ -434,7 +448,7 @@ def __init__( def forward(self, coeff: torch.Tensor) -> torch.Tensor: """Expand ``(N, D, F, C)`` coefficients to ``(N, D, F, K*C)``.""" weight = self.weight.index_select(0, self.degree_index) - return torch.einsum("ndfi,dio->ndfo", coeff, weight) + return _degree_batched_matmul(coeff, weight) class BaseGridNet(nn.Module): diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py index 93092ecd9c..076b9b6a04 100644 --- a/deepmd/pt_expt/common.py +++ b/deepmd/pt_expt/common.py @@ -136,24 +136,29 @@ def try_convert_module(value: Any) -> torch.nn.Module | None: _AUTO_WRAPPED_CLASSES: dict[type, type] = {} -def _auto_wrap_native_op(value: NativeOP) -> torch.nn.Module: - """Auto-wrap any NativeOP as a torch.nn.Module via ``torch_module``. +def auto_wrapped_class(cls: type) -> type: + """Return the cached ``torch_module`` auto-wrap of a dpmodel class. Creates a subclass with a generic ``forward`` that delegates to ``call``, then applies ``torch_module`` to get full ``__setattr__`` / post-init list conversion. The wrapped class is cached per dpmodel type. + Invariant: construct this wrapped class directly whenever live, + constructor-supplied components must retain non-serialized runtime + state. Converting a populated raw dpmodel instance instead goes + through the ``serialize()``/``deserialize()`` round-trip of + ``_auto_wrap_native_op``, which preserves only the portable record. + Parameters ---------- - value : NativeOP - The dpmodel object to wrap. + cls : type + The dpmodel NativeOP class to wrap. Returns ------- - torch.nn.Module - The wrapped pt_expt module, deserialized from value's serialized state. + type + The ``torch_module``-wrapped subclass. """ - cls = type(value) if cls not in _AUTO_WRAPPED_CLASSES: wrapped = type( cls.__name__, @@ -161,7 +166,24 @@ def _auto_wrap_native_op(value: NativeOP) -> torch.nn.Module: {"forward": lambda self, *args, **kwargs: self.call(*args, **kwargs)}, ) _AUTO_WRAPPED_CLASSES[cls] = torch_module(wrapped) - wrapped_cls = _AUTO_WRAPPED_CLASSES[cls] + return _AUTO_WRAPPED_CLASSES[cls] + + +def _auto_wrap_native_op(value: NativeOP) -> torch.nn.Module: + """Auto-wrap any NativeOP as a torch.nn.Module via ``torch_module``. + + Parameters + ---------- + value : NativeOP + The dpmodel object to wrap. + + Returns + ------- + torch.nn.Module + The wrapped pt_expt module, deserialized from value's serialized state. + """ + cls = type(value) + wrapped_cls = auto_wrapped_class(cls) if not (hasattr(value, "serialize") and hasattr(wrapped_cls, "deserialize")): raise TypeError( f"Cannot auto-wrap {cls.__name__}: " diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index a994958df4..4b8230f2e7 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -9,11 +9,12 @@ import copy import logging -from deepmd.dpmodel.atomic_model.dp_atomic_model import ( - DPAtomicModel, +from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel as LinearEnergyAtomicModelDP, ) from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( - PairTabAtomicModel, + PairTabAtomicModel as PairTabAtomicModelDP, ) from deepmd.dpmodel.model.model_factory import ( BackendModelFactory, @@ -21,6 +22,9 @@ from deepmd.dpmodel.model.model_factory import ( get_spin_model as get_spin_model_from_factory, ) +from deepmd.pt_expt.common import ( + auto_wrapped_class, +) from deepmd.pt_expt.descriptor import ( BaseDescriptor, ) @@ -56,6 +60,12 @@ _WARNED_ONCE: set[str] = set() +# wrapped atomic classes: constructed directly so live children keep their +# runtime state (see the auto_wrapped_class invariant) +DPAtomicModel = auto_wrapped_class(DPAtomicModelDP) +PairTabAtomicModel = auto_wrapped_class(PairTabAtomicModelDP) +LinearEnergyAtomicModel = auto_wrapped_class(LinearEnergyAtomicModelDP) + _model_factory = BackendModelFactory( descriptor_base=BaseDescriptor, fitting_base=BaseFitting, @@ -64,6 +74,7 @@ atomic_model=DPAtomicModel, pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, + linear_atomic_model=LinearEnergyAtomicModel, ) get_zbl_model = _model_factory.get_zbl_model diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 7aa82215b2..cee2bee6b5 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -28,6 +28,7 @@ cuda_infer_level, ) from deepmd.pt_expt.common import ( + auto_wrapped_class, torch_module, ) from deepmd.pt_expt.utils.graph_builder import ( @@ -465,7 +466,9 @@ def make_model( The model. """ - DPModel = make_model_dp(T_AtomicModel) + # wrapped atomic class: live descriptor/fitting keep their runtime + # state (see the auto_wrapped_class invariant) + DPModel = make_model_dp(auto_wrapped_class(T_AtomicModel)) @torch_module class CM(DPModel, *T_Bases): diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py index b0678b3711..cb1465ae8a 100644 --- a/source/tests/common/dpmodel/test_descrpt_dpa4.py +++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py @@ -265,6 +265,19 @@ def test_supported_feature_roundtrip(self, overrides) -> None: out2 = np.asarray(dd2.call(coord.reshape(nf, -1), atype, nlist)[0]) np.testing.assert_array_equal(out1, out2) + def test_use_amp_stays_out_of_the_portable_record(self) -> None: + """``use_amp`` is a runtime/training policy, not model state. + + The portable serialization must not carry it (a ``use_amp: true`` + record would e.g. be rejected by the JAX deserializer); a fresh + deserialize falls back to the constructor default. Construction-time + survival is pinned at the pt_expt assembly boundary instead + (``test_get_model_dpa4.py``). + """ + dd = make_descriptor(use_amp=False) + assert dd.use_amp is False + assert "use_amp" not in dd.serialize()["config"] + def test_legacy_spin_gate_is_squared_on_deserialize(self) -> None: """Version 1.2 stores the env-seed spin gate after the quadratic form. diff --git a/source/tests/common/dpmodel/test_dpa4_frame_mixers.py b/source/tests/common/dpmodel/test_dpa4_frame_mixers.py index 38498c84ec..8958b204dc 100644 --- a/source/tests/common/dpmodel/test_dpa4_frame_mixers.py +++ b/source/tests/common/dpmodel/test_dpa4_frame_mixers.py @@ -3,9 +3,12 @@ These mirror the current pt ``deepmd.pt.model.descriptor.sezm_nn.grid_net`` ``FrameContract`` / -``FrameExpand`` (and the ``_build_frame_degree_index`` helper). The pt mixers -realise a per-degree ``einsum("ndfi,dio->ndfo", coeff, weight[degree_index])``; -the dpmodel port realises the same map as a broadcast batched ``xp.matmul``. +``FrameExpand`` (and the ``_build_frame_degree_index`` helper). The +backend-independent mathematical contract of both mixers is the per-degree +``einsum("ndfi,dio->ndfo", coeff, weight[degree_index])``; both backends +realise it through the same ``(D, F)``-batched matmul lowering +(``_degree_batched_matmul``), which these tests pin against the einsum +contract for values and gradients. pt imports live inside the test functions because ruff TID253 bans module-level ``deepmd.pt`` imports under ``source/tests/common``. pt modules @@ -211,3 +214,134 @@ def test_torch_namespace(cls) -> None: rtol=1e-12, atol=1e-12, ) + + +@pytest.mark.parametrize( + "cls", + [ + DPFrameContract, # (N, D, F, K*C) -> (N, D, F, C) + DPFrameExpand, # (N, D, F, C) -> (N, D, F, K*C) + ], +) +def test_empty_batch_passes_through(cls) -> None: + """An empty node axis yields ``(0, D, F, o)`` on every namespace. + + Reachable when the cross-grid leading axis is an empty graph/edge set + or a distributed rank owns no nodes. + """ + import torch + + lmax, kmax, channels, n_focus = 2, 1, 4, 2 + n_frames = 2 * kmax + 1 + coeff_dim = (lmax + 1) ** 2 + mod = cls( + lmax=lmax, + mmax=lmax, + coefficient_layout="packed", + n_frames=n_frames, + channels=channels, + precision="float64", + trainable=True, + seed=7, + ) + in_dim = n_frames * channels if cls is DPFrameContract else channels + out_dim = channels if cls is DPFrameContract else n_frames * channels + coeff = np.zeros((0, coeff_dim, n_focus, in_dim), dtype=np.float64) + out = mod.call(coeff) + assert out.shape == (0, coeff_dim, n_focus, out_dim) + t_out = mod.call(torch.from_numpy(coeff)) + assert tuple(t_out.shape) == (0, coeff_dim, n_focus, out_dim) + + +@pytest.mark.parametrize( + "kind", + [ + "contract", # (N, D, F, K*C) -> (N, D, F, C) + "expand", # (N, D, F, C) -> (N, D, F, K*C) + ], +) +def test_focus_batched_lowering_matches_einsum_backward(kind) -> None: + """``F > 1`` forward AND backward parity of the ``(D, F)``-batched + lowering against the ``einsum("ndfi,dio->ndfo")`` contract, for the + input and the weight gradients. + """ + import torch + + lmax, kmax, channels = 2, 1, 4 + n_frames = 2 * kmax + 1 + coeff_dim = (lmax + 1) ** 2 + n_batch, n_focus = 5, 2 + rng = np.random.default_rng(2026) + + if kind == "contract": + from deepmd.pt.model.descriptor.sezm_nn.grid_net import FrameContract as PTMixer + + in_dim = n_frames * channels + else: + from deepmd.pt.model.descriptor.sezm_nn.grid_net import FrameExpand as PTMixer + + in_dim = channels + pt_mod = PTMixer( + lmax=lmax, + mmax=lmax, + coefficient_layout="packed", + n_frames=n_frames, + channels=channels, + dtype=torch.float64, + trainable=True, + seed=7, + ).to("cpu") + + coeff = torch.from_numpy( + rng.normal(size=(n_batch, coeff_dim, n_focus, in_dim)) + ).requires_grad_(True) + out = pt_mod(coeff) + grad_out = torch.from_numpy(rng.normal(size=tuple(out.shape))) + out.backward(grad_out) + grad_in_mod = coeff.grad.detach().clone() + grad_w_mod = pt_mod.weight.grad.detach().clone() + + pt_mod.weight.grad = None + coeff_ref = coeff.detach().clone().requires_grad_(True) + ref = torch.einsum( + "ndfi,dio->ndfo", + coeff_ref, + pt_mod.weight.index_select(0, pt_mod.degree_index), + ) + np.testing.assert_allclose( + out.detach().numpy(), ref.detach().numpy(), rtol=1e-12, atol=1e-12 + ) + ref.backward(grad_out) + np.testing.assert_allclose( + grad_in_mod.numpy(), coeff_ref.grad.numpy(), rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose( + grad_w_mod.numpy(), pt_mod.weight.grad.numpy(), rtol=1e-12, atol=1e-12 + ) + + # the dpmodel lowering agrees on the torch namespace, gradients included + import array_api_compat + + from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import ( + _degree_batched_matmul, + ) + + coeff_dp = coeff.detach().clone().requires_grad_(True) + # leaf copy of the per-degree parameter: its gradient pins the dpmodel + # lowering's WEIGHT backward too, not only the input backward + weight_dp = pt_mod.weight.detach().clone().requires_grad_(True) + dp_out = _degree_batched_matmul( + array_api_compat.array_namespace(coeff_dp), + coeff_dp, + weight_dp.index_select(0, pt_mod.degree_index), + ) + np.testing.assert_allclose( + dp_out.detach().numpy(), out.detach().numpy(), rtol=1e-12, atol=1e-12 + ) + dp_out.backward(grad_out) + np.testing.assert_allclose( + coeff_dp.grad.numpy(), grad_in_mod.numpy(), rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose( + weight_dp.grad.numpy(), grad_w_mod.numpy(), rtol=1e-12, atol=1e-12 + ) diff --git a/source/tests/common/dpmodel/test_dpa4_lora.py b/source/tests/common/dpmodel/test_dpa4_lora.py index 54d4257394..c46ed3ec05 100644 --- a/source/tests/common/dpmodel/test_dpa4_lora.py +++ b/source/tests/common/dpmodel/test_dpa4_lora.py @@ -1,5 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Torch-free tests for the dpmodel DPA4 (SeZM) LoRA fine-tune freeze policy.""" +"""Tests for the dpmodel DPA4 (SeZM) LoRA adapters: the fine-tune freeze +policy and the ``LoRASO3`` contraction contract (torch imported lazily). +""" + +import numpy as np +import pytest from deepmd.dpmodel.descriptor.dpa4 import ( DescrptDPA4, @@ -57,3 +62,54 @@ def test_apply_lora_marks_adapters_trainable() -> None: ] assert type_embeddings assert all(not m.trainable for m in type_embeddings) + + +@pytest.mark.parametrize( + "n_focus", + [ + 1, # the common single-focus configuration + 2, # shipped spin/property DPA4 examples + ], +) +def test_lora_so3_call_matches_einsum_contract(n_focus) -> None: + """Direct regression for the dpmodel ``LoRASO3.call`` contraction. + + ``B_by_l`` is set nonzero so the adapter delta participates; the + forward must equal ``einsum("ndfi,difo->ndfo")`` over the effective + (base + scaled ``B @ A``) per-degree weight, on both the NumPy and + the Torch array namespaces. + """ + import torch + + from deepmd.dpmodel.descriptor.dpa4_nn.lora import ( + LoRASO3, + ) + + lmax, cin, cout, rank = 2, 3, 4, 2 + mod = LoRASO3( + lmax=lmax, + in_channels=cin, + out_channels=cout, + n_focus=n_focus, + precision="float64", + trainable=True, + seed=3, + lora_rank=rank, + ) + rng = np.random.default_rng(11) + # unlock the adapter: B is zero-initialised, which would hide a delta bug + mod.B_by_l = rng.normal(size=mod.B_by_l.shape).astype(np.float64) + + coeff_dim = (lmax + 1) ** 2 + x = rng.normal(size=(5, coeff_dim, n_focus, cin)) + delta = np.matmul(mod.B_by_l, mod.A_by_l).transpose(0, 2, 1) * mod.scaling + w_eff = np.reshape(mod.weight + delta, (lmax + 1, cin, n_focus, cout)) + w_deg = w_eff[np.asarray(mod.expand_index)] # (D, Cin, F, Cout) + ref = np.einsum("ndfi,difo->ndfo", x, w_deg) + + out_np = np.asarray(mod.call(x)) + np.testing.assert_allclose(out_np, ref, rtol=1e-12, atol=1e-12) + out_torch = mod.call(torch.from_numpy(x)) + np.testing.assert_allclose( + out_torch.detach().cpu().numpy(), ref, rtol=1e-12, atol=1e-12 + ) diff --git a/source/tests/pt_expt/model/test_get_model_dpa4.py b/source/tests/pt_expt/model/test_get_model_dpa4.py index bbb59af1f0..86e180b378 100644 --- a/source/tests/pt_expt/model/test_get_model_dpa4.py +++ b/source/tests/pt_expt/model/test_get_model_dpa4.py @@ -321,5 +321,107 @@ def test_unrelated_construction_error_propagates(self) -> None: get_model(raw) +class TestUseAmpSurvivesAssembly(unittest.TestCase): + """``use_amp`` is runtime policy: it must survive model ASSEMBLY without + entering the portable serialization record (which the JAX deserializer + rejects for ``use_amp: true``). The wrapping of the atomic model must + therefore not round-trip the constructed descriptor through + ``serialize()``/``deserialize()``. + """ + + def test_get_model_keeps_use_amp_false(self) -> None: + model = get_model( + _make_raw_model_config( + descriptor={ + "sel": 20, + "rcut": 4.0, + "channels": 8, + "n_radial": 4, + "lmax": 1, + "mmax": 1, + "n_blocks": 1, + "precision": "float64", + "seed": 1, + "use_amp": False, + } + ) + ) + assert model.atomic_model.descriptor.use_amp is False + + def test_get_model_keeps_the_use_amp_default(self) -> None: + model = get_model(_make_raw_model_config()) + assert model.atomic_model.descriptor.use_amp is True + + def test_standard_type_keeps_use_amp_false(self) -> None: + """The plain `standard` route wraps through the same boundary.""" + cfg = _make_raw_model_config( + descriptor={ + "type": "dpa4", + "sel": 20, + "rcut": 4.0, + "channels": 8, + "n_radial": 4, + "lmax": 1, + "mmax": 1, + "n_blocks": 1, + "precision": "float64", + "seed": 1, + "use_amp": False, + }, + fitting_net={ + "type": "dpa4_ener", + "precision": "float64", + "seed": 1, + }, + ) + del cfg["type"] + model = get_model(cfg) + assert model.atomic_model.descriptor.use_amp is False + + def test_bridged_composition_keeps_use_amp_false(self) -> None: + """The ZBL composition path must be lossless too: the learned child + is a live module, not a serialize round-trip rebuild. + """ + model = get_model( + _make_raw_model_config( + descriptor={ + "sel": 20, + "rcut": 4.0, + "channels": 8, + "n_radial": 4, + "lmax": 1, + "mmax": 1, + "n_blocks": 1, + "precision": "float64", + "seed": 1, + "use_amp": False, + }, + bridging_method="ZBL", + ) + ) + assert model.atomic_model.models[0].descriptor.use_amp is False + + def test_linear_ener_child_keeps_use_amp_false(self) -> None: + """The explicit `linear_ener` route builds its children as wrapped + modules directly -- same lossless rule as the bridged path. + """ + base = _make_raw_model_config() + child_descriptor = dict(base["descriptor"], type="dpa4", use_amp=False) + child_fitting = dict(base["fitting_net"], type="dpa4_ener") + model = get_model( + { + "type": "linear_ener", + "type_map": base["type_map"], + "models": [ + { + "descriptor": child_descriptor, + "fitting_net": child_fitting, + } + ], + } + ) + assert model.atomic_model.models[0].descriptor.use_amp is False + + if __name__ == "__main__": unittest.main()