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.
@@ -401,8 +433,12 @@ def call(
         router = self.router(scalar_pair)
         router = xp.exp(router - xp.max(router, axis=-1, keepdims=True))
         router = router / xp.sum(router, axis=-1, keepdims=True)
-        # einsum "ngfhc,nfh->ngfc" as a broadcast sum over the branch axis
-        out = xp.sum(value * router[:, None, :, :, None], axis=3)  # (N, G, F, C)
+        # einsum "ngfhc,nfh->ngfc" as a broadcast multiply and reduction.
+        # Spelling it as a matmul over H gives a batched GEMM with M=1, K=H,
+        # which cuBLAS serves from its slow small-N kernels: 7.5 ms vs 1.8 ms
+        # here at H=1, and no better at H=3.  The intermediate this form
+        # materialises is only H (a handful) times the result.
+        out = xp.sum(value * router[:, None, :, :, None], axis=3)
 
         # === Step 3. Project back to coefficients and mix output channels ===
         return _project_frames(from_grid(out), self.out_proj, self.n_frames)
@@ -493,9 +529,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 +610,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,12 @@ 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", batched over the small focus axis F.
+        # Batching over B instead would broadcast the weight to (B, F, Cin, Cout)
+        # and force autograd to reduce that expansion on every backward.
         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 +441,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,

deepmd/pt_expt/model/get_model.py

@@ -8,6 +8,7 @@
 
 import copy
 import logging
+import os
 from typing import (
     TYPE_CHECKING,
 )
@@ -57,8 +58,46 @@
 
 log = logging.getLogger(__name__)
 
-# Warn at most once per process for backend-ignored switches (keyed by name).
-_WARNED_ONCE: set[str] = set()
+#: ``DP_TF32_INFER`` -> eval-time matmul precision.  Same table as the pt
+#: backend's ``sezm_model._TF32_INFER_PRECISION_CHOICES``.
+_TF32_INFER_PRECISION_CHOICES = {
+    "0": "highest",
+    "1": "high",
+    "2": "medium",
+}
+
+
+def _apply_tf32_policy(model: BaseModel, data: dict) -> BaseModel:
+    """Attach the DPA4/SeZM TF32 matmul-precision policy to a built model.
+
+    As in pt: training forwards follow ``enable_tf32`` (default ``True``),
+    eval forwards follow ``DP_TF32_INFER``.  The policy is applied in
+    ``call_common``, and in ``_CompiledModel.forward`` when compiled.
+
+    Parameters
+    ----------
+    model : BaseModel
+        The freshly built model to configure.
+    data : dict
+        The model config section, read for ``enable_tf32``.
+
+    Returns
+    -------
+    BaseModel
+        The same model, with the precision policy attached.
+
+    Raises
+    ------
+    ValueError
+        If ``DP_TF32_INFER`` is set to anything other than ``0``, ``1``, or
+        ``2``.
+    """
+    model.enable_tf32 = bool(data.get("enable_tf32", True))
+    tf32_infer_env = os.environ.get("DP_TF32_INFER", "0").strip().lower()
+    if tf32_infer_env not in _TF32_INFER_PRECISION_CHOICES:
+        raise ValueError(f"DP_TF32_INFER must be one of 0/1/2, got {tf32_infer_env!r}")
+    model.tf32_infer_precision = _TF32_INFER_PRECISION_CHOICES[tf32_infer_env]
+    return model
 
 
 _model_factory = BackendModelFactory(
@@ -90,17 +129,11 @@ def get_sezm_model(data: dict) -> BaseModel:
 
     Notes
     -----
-    ``enable_tf32`` is accepted but ignored: the pt backend uses it to toggle
-    TF32 matmul precision, while the pt_expt backend always runs at full
-    ("highest") matmul precision, which is numerically conservative.
+    ``enable_tf32`` behaves as in pt: training forwards run at TF32 ("high")
+    precision when set (the default), eval forwards follow ``DP_TF32_INFER``.
+    See :func:`_apply_tf32_policy`.
     """
     data = copy.deepcopy(data)
-    if bool(data.get("enable_tf32", True)) and "enable_tf32" not in _WARNED_ONCE:
-        log.warning(
-            "`enable_tf32` has no effect on the pt_expt backend, which "
-            "always runs at full ('highest') matmul precision; ignoring it."
-        )
-        _WARNED_ONCE.add("enable_tf32")
     if "spin" in data:
         if str(data["spin"].get("scheme", "deepspin")) != "native":
             raise NotImplementedError(
@@ -172,8 +205,9 @@ def get_sezm_model(data: dict) -> BaseModel:
         pair_exclude_types=pair_exclude_types,
     )
     if bridging_enabled:
-        return _compose_bridging(model, data, bridging_method)
-    return model
+        # The TF32 policy attaches to whichever model is returned.
+        return _apply_tf32_policy(_compose_bridging(model, data, bridging_method), data)
+    return _apply_tf32_policy(model, data)
 
 
 def _compose_bridging(
@@ -338,7 +372,10 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
             "spin scheme 'native' requires an atomic model declaring "
             "supports_native_spin()"
         )
-    return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin)
+    return _apply_tf32_policy(
+        NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin),
+        data,
+    )
 
 
 def get_linear_model(model_params: dict) -> BaseModel:

deepmd/pt_expt/model/make_model.py

@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import contextlib
 import math
 import types
+from collections.abc import (
+    Generator,
+)
 from typing import (
     Any,
 )
@@ -470,6 +474,55 @@ def get_min_nbor_dist(self) -> float | None:
             """Get the minimum distance between two atoms."""
             return self.min_nbor_dist
 
+        # === TF32 matmul precision ===
+        # Same policy as pt's SeZMModel: training follows ``enable_tf32``,
+        # eval follows ``DP_TF32_INFER``.  The DPA4/SeZM builders in
+        # ``get_model`` set both; every other model keeps these defaults,
+        # which mean full fp32 either way.
+        enable_tf32: bool = False
+        tf32_infer_precision: str = "highest"
+
+        @contextlib.contextmanager
+        def tf32_precision_ctx(self) -> Generator[None, None, None]:
+            """Select the matmul precision for one forward, then restore it.
+
+            Yields
+            ------
+            None
+                With ``torch.set_float32_matmul_precision`` set for the
+                duration of the block.
+            """
+            if not torch.cuda.is_available():
+                yield
+                return
+            prev_precision = torch.get_float32_matmul_precision()
+            try:
+                if self.training:
+                    precision = "high" if self.enable_tf32 else "highest"
+                else:
+                    precision = self.tf32_infer_precision
+                torch.set_float32_matmul_precision(precision)
+                yield
+            finally:
+                torch.set_float32_matmul_precision(prev_precision)
+
+        def call_common(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
+            """Run the shared dense/graph forward under the TF32 policy.
+
+            Every model's ``forward`` reaches the backbone through here, so
+            this is where eager forwards pick their matmul precision.  Export
+            traces root at ``call_common_lower``, so the switch stays out of
+            exported graphs.  Compiled training skips this method and applies
+            the policy in ``_CompiledModel.forward`` instead.
+
+            Returns
+            -------
+            dict[str, torch.Tensor]
+                The backbone's output dict, unchanged.
+            """
+            with self.tf32_precision_ctx():
+                return super().call_common(*args, **kwargs)
+
         def forward(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
             """Default forward delegates to call().
 

deepmd/pt_expt/train/training.py

@@ -961,7 +961,23 @@ def __getattr__(self, name: str) -> Any:
         except AttributeError:
             return getattr(self.original_model, name)
 
-    def forward(
+    def forward(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
+        """Run the compiled forward under the wrapped model's TF32 policy.
+
+        This path never reaches ``call_common``, where eager forwards set their
+        precision, so it applies the same policy here.  The context also covers
+        the lazy compile below: Inductor picks its GEMM backend while lowering,
+        so setting precision only around the call would miss the kernels.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            The model prediction dict.
+        """
+        with self.original_model.tf32_precision_ctx():
+            return self._forward_dispatch(*args, **kwargs)
+
+    def _forward_dispatch(
         self,
         coord: torch.Tensor,
         atype: torch.Tensor,

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)