perf(dpa4): batch the SO3/grid contractions over (D,F); keep use_amp through pt_expt assembly

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."""

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

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:

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

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

deepmd/pt_expt/common.py

@@ -136,32 +136,54 @@ 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__,
             (cls,),
             {"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__}: "

deepmd/pt_expt/model/get_model.py

@@ -9,18 +9,22 @@
 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,
 )
 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
 

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

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.