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

deepmd/dpmodel/descriptor/dpa4.py

@@ -2748,6 +2748,11 @@ def serialize(self) -> dict[str, Any]:
                 "mlp_bias": self.mlp_bias,
                 "exclude_types": self.exclude_types,
                 "eps": self.eps,
+                # Must round-trip: pt_expt rebuilds the descriptor from this
+                # dict, so omitting the key silently reset a configured
+                # ``use_amp: false`` to True and kept training in bfloat16.
+                # Older records without it still load (__init__ defaults it).
+                "use_amp": self.use_amp,
                 "trainable": self.trainable,
                 "seed": self.seed,
                 "inner_clamp_r_inner": self.inner_clamp_r_inner,

deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py

@@ -99,6 +99,38 @@ 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 degree axis, not over ``N``: the latter would broadcast
+    ``weight`` to ``(N, D, i, o)`` and make autograd reduce that expansion on
+    every backward.  The transposes touch only ``coeff``, which is smaller.
+    """
+    n_batch, coeff_dim, n_focus, _ = coeff.shape
+    coeff_d = xp.reshape(
+        xp.permute_dims(coeff, (1, 0, 2, 3)), (coeff_dim, n_batch * n_focus, -1)
+    )  # (D, N*F, i)
+    out = xp.matmul(coeff_d, weight)  # (D, N*F, o)
+    out = xp.reshape(out, (coeff_dim, n_batch, n_focus, -1))
+    return xp.permute_dims(out, (1, 0, 2, 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 +525,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 degree axis, 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 +606,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 degree axis, 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,16 @@ 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", batched over the small (D, F) axes.
+        # Batching over the node axis N instead would broadcast the weight to
+        # (N, D, F, Cin, Cout) -- for the water example a 165K-element parameter
+        # blown up to 191M elements per call -- and autograd would then reduce
+        # that expansion back down.  It was the costliest kernel of a step.
         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/pt/model/descriptor/sezm.py

@@ -2561,6 +2561,9 @@ def serialize(self) -> dict[str, Any]:
                 "mlp_bias": self.mlp_bias,
                 "exclude_types": self.exclude_types,
                 "eps": self.eps,
+                # Kept in step with the dpmodel serialize contract so both
+                # backends' records carry the same keys.
+                "use_amp": self.use_amp,
                 "trainable": self.trainable,
                 "seed": self.seed,
                 "inner_clamp_r_inner": self.inner_clamp_r_inner,

source/tests/common/dpmodel/test_descrpt_dpa4.py

@@ -265,6 +265,27 @@ 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)
 
+    @pytest.mark.parametrize(
+        "use_amp",
+        [
+            True,  # the constructor default; must not be clobbered either
+            False,  # the value that was silently lost, re-enabling autocast
+        ],
+    )
+    def test_use_amp_survives_roundtrip(self, use_amp) -> None:
+        """``use_amp`` must round-trip through serialize/deserialize.
+
+        The key was missing from the config, so a backend that rebuilds from
+        it (pt_expt does) reset ``use_amp: false`` to True and kept training in
+        bfloat16.  The forward-output round-trip test can't catch this --
+        dpmodel never autocasts, so outputs match either way.
+        """
+        dd = make_descriptor(use_amp=use_amp)
+        assert dd.use_amp is use_amp
+        assert dd.serialize()["config"]["use_amp"] is use_amp
+        dd2 = DescrptDPA4.deserialize(dd.serialize())
+        assert dd2.use_amp is use_amp
+
     def test_value_errors(self) -> None:
         with pytest.raises(ValueError):  # kmax must be <= lmax
             make_descriptor(kmax=4, lmax=3)