From 7518a417cbebf1367c262b082c2dcce818250d8b Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 5 Aug 2026 01:37:00 +0800
Subject: [PATCH 01/16] perf(dpmodel): contract the DPA4 grid-branch router
 with matmul

The GridBranch router contraction einsum("ngfhc,nfh->ngfc") was written as
a broadcast multiply followed by a reduce over the branch axis.  That
materialises the entire (N, G, F, H, C) product -- roughly 0.8 GB at the
grid resolution of examples/water/dpa4 -- writes it to memory and reads it
straight back, and the backward pays the same traffic again.

An op-level CUDA profile of a DPA4 training step measured this single
reduce at 45.6 ms per call over a [1152, 9, 1, 32, 576] operand, three
calls per step: the most expensive kernel in the run.  The pt backend
spells the same contraction as torch.einsum and never builds the
intermediate.

Use xp.matmul instead, which is array-API standard (unlike np.einsum,
which is what the broadcast form was avoiding) and contracts H in place so
only the (N, G, F, C) result is written.  matmul broadcasts its leading
batch axes, so the router reshapes to (N, 1, F, 1, H) and lines up with
value's (N, G, F, H, C) without any permute -- a permute would reintroduce
the copy this removes.
---
 deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index 4c63d1a492..6994ad1497 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -401,8 +401,22 @@ 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", expressed as a batched matmul.
+        #
+        # NOT as ``xp.sum(value * router[:, None, :, :, None], axis=3)``: that
+        # broadcast-then-reduce materialises the whole (N, G, F, H, C) product
+        # -- ~0.8 GB at this example's grid resolution -- and reads it straight
+        # back, which measured as the single most expensive kernel of a DPA4
+        # training step. ``matmul`` contracts H in place, so only the
+        # (N, G, F, C) result is written. ``matmul`` broadcasts its leading
+        # batch axes, so the router's (N, 1, F, 1, H) view lines up with
+        # value's (N, G, F, H, C) without permuting (a permute here would
+        # reintroduce the very copy this avoids).
+        router_row = xp.reshape(
+            router, (n_batch, 1, n_focus, 1, self.n_branches)
+        )  # (N, 1, F, 1, H)
+        out = xp.matmul(router_row, value)  # (N, G, F, 1, C)
+        out = xp.reshape(out, (n_batch, n_grid, n_focus, self.channels))
 
         # === Step 3. Project back to coefficients and mix output channels ===
         return _project_frames(from_grid(out), self.out_proj, self.n_frames)

From 01c58e66540672aa82dcace4923c989807570586 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 5 Aug 2026 02:13:02 +0800
Subject: [PATCH 02/16] Revert "perf(dpmodel): contract the DPA4 grid-branch
 router with matmul"

This reverts commit 7518a417c.

Measurement did not support it.  An op-level profile attributed the 45.6 ms
reduce to ExpandBackward0, not to this multiply, and re-benchmarking after
the change moved DPA4 eager training by nothing (1.514 -> 1.552 s/step, i.e.
run-to-run noise) while the offending kernel stayed byte-identical at
410.7 ms.  The GridBranch product is well under the size that would matter.

Since matmul is autocast-listed where mul/sum are not, keeping it would have
silently moved this contraction into bf16 under the autocast region for no
measured gain.  The actual site is the broadcast weight in so3.py, fixed
separately.
---
 deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py | 18 ++----------------
 1 file changed, 2 insertions(+), 16 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index 6994ad1497..4c63d1a492 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -401,22 +401,8 @@ 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", expressed as a batched matmul.
-        #
-        # NOT as ``xp.sum(value * router[:, None, :, :, None], axis=3)``: that
-        # broadcast-then-reduce materialises the whole (N, G, F, H, C) product
-        # -- ~0.8 GB at this example's grid resolution -- and reads it straight
-        # back, which measured as the single most expensive kernel of a DPA4
-        # training step. ``matmul`` contracts H in place, so only the
-        # (N, G, F, C) result is written. ``matmul`` broadcasts its leading
-        # batch axes, so the router's (N, 1, F, 1, H) view lines up with
-        # value's (N, G, F, H, C) without permuting (a permute here would
-        # reintroduce the very copy this avoids).
-        router_row = xp.reshape(
-            router, (n_batch, 1, n_focus, 1, self.n_branches)
-        )  # (N, 1, F, 1, H)
-        out = xp.matmul(router_row, value)  # (N, G, F, 1, C)
-        out = xp.reshape(out, (n_batch, n_grid, n_focus, self.channels))
+        # 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)
 
         # === Step 3. Project back to coefficients and mix output channels ===
         return _project_frames(from_grid(out), self.out_proj, self.n_frames)

From 193b49ecd336bbef436e1ce1abf736237e9ae106 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 5 Aug 2026 02:14:40 +0800
Subject: [PATCH 03/16] perf(dpmodel): stop broadcasting DPA4 so3 linear
 weights across nodes

Both so3 channel mixers spelled their einsum as a batched matmul with the
NODE/EDGE axis as the matmul BATCH and the weight carrying a dummy leading
axis:

    matmul(x[:, :, :, None, :], weight_expanded[None, ...])

matmul broadcasts batch axes, so this expands the weight to
(N, D, F, Cin, Cout).  For examples/water/dpa4 that turns a 165K-element
parameter into 191M elements -- about 0.8 GB -- on every call, and autograd
must then reduce the whole expanded gradient back to the parameter shape.

An op-level CUDA profile of a DPA4 training step attributed 45.6 ms per
call to that ExpandBackward0 reduce over a [1152, 9, 1, 32, 576] operand,
three calls per step, making it the most expensive kernel in the run; the
ChannelLinear twin cost a further ~7-9 ms per call over [102510, 1, 32, 64].
The pt backend spells the same contraction as torch.einsum and never
expands the weight.

Batch over the small (D, F) / (F,) axes instead, which keeps N as matmul
ROWS.  The weight is then used in place and its gradient is an ordinary
matmul.  The transposes this adds touch only the (N, D, F, C) operands,
which are orders of magnitude smaller than the expanded weight.
---
 deepmd/dpmodel/descriptor/dpa4_nn/so3.py | 30 +++++++++++++++++++-----
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
index 8ca2dbc855..15cd5c10a9 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
@@ -131,10 +131,19 @@ 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 a matmul batched over the FOCUS axis.
+        #
+        # NOT as ``matmul(x[:, :, None, :], weight[None, ...])``: that makes B a
+        # batch axis, so matmul broadcasts the weight to (B, F, Cin, Cout) --
+        # inflating a few-hundred-KB parameter into hundreds of millions of
+        # elements per call, whose gradient autograd must then reduce back down
+        # (an ``ExpandBackward0`` reduce that measured as the single most
+        # expensive kernel of a DPA4 training step). Batching over F instead
+        # keeps B as matmul ROWS, so the weight is used in place and its
+        # gradient is an ordinary matmul.
         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 +448,21 @@ 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" as a matmul batched over the (D, F) axes.
+        #
+        # NOT as ``matmul(x[:, :, :, None, :], weight_expanded[None, ...])``:
+        # that makes N a batch axis, so matmul broadcasts the weight to
+        # (N, D, F, Cin, Cout) -- for the water DPA4 example, a 165K-element
+        # parameter expanded to 191M elements (~0.8 GB) on every call, whose
+        # gradient autograd then reduces back down. That ``ExpandBackward0``
+        # reduce measured at 45.6 ms per call, three calls per training step:
+        # the most expensive kernel in the run. Batching over the small (D, F)
+        # axes keeps N as matmul ROWS, so the weight is never expanded.
         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:

From 75459610a34c9db99c468cd754dceb0017e2d3a1 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 5 Aug 2026 09:40:13 +0800
Subject: [PATCH 04/16] perf(dpmodel): remove the remaining DPA4
 broadcast-weight contractions

Follow-up to the so3.py fix, applying the same correction wherever a
contraction was spelled so that the NODE axis becomes the matmul BATCH and
a trainable tensor is broadcast across it:

* grid_net.GridBranch  einsum "ngfhc,nfh->ngfc" -- was a broadcast multiply
  plus a reduce, materialising an (N, G, F, H, C) product H times the size
  of its own result.
* grid_net.FrameContract / FrameExpand  einsum "ndfi,dio->ndfo" -- broadcast
  the per-degree weight to (N, D, i, o).  Both now share
  _degree_batched_matmul, which batches over the small degree axis.
* lora.call  einsum "ndfi,difo->ndfo" -- the LoRA twin of the so3.py site.

In every case autograd had to reduce the fully expanded gradient back to
the parameter shape on each step; batching over the small (D, F) axes keeps
N as matmul ROWS so the weight is used in place.

The two projection.py sites that share the [None, ...] spelling are left
alone deliberately: to_grid_mat / from_grid_mat are registered as BUFFERS
with requires_grad=False (verified on a constructed DPA4), so no gradient
is taken for them and none of the expensive half applies.

Covered by the existing pt-parity gates, which construct these classes
directly: test_dpa4_frame_mixers.py (FrameContract/FrameExpand, fp64
weight-copied vs pt), test_dpa4_gridbranch_frames.py, test_dpa4_lora.py,
and test_dpa4_dpmodel_parity.py.
---
 deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py | 65 ++++++++++++++++---
 deepmd/dpmodel/descriptor/dpa4_nn/lora.py     | 11 +++-
 2 files changed, 65 insertions(+), 11 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index 4c63d1a492..e007932d0b 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -99,6 +99,41 @@ 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
+    -----
+    The obvious spelling ``matmul(coeff, weight[None, ...])`` puts ``N`` in the
+    matmul batch, so ``weight`` is broadcast to ``(N, D, i, o)`` and autograd
+    must reduce that expanded gradient back to the parameter shape every step.
+    Batching over the small degree axis keeps ``N`` as matmul ROWS, so the
+    weight is used in place; the transposes touch only ``coeff``, which is far
+    smaller than the expanded weight would be.
+    """
+    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 +436,18 @@ 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 batched matmul.
+        #
+        # NOT as ``xp.sum(value * router[:, None, :, :, None], axis=3)``: that
+        # materialises the whole (N, G, F, H, C) product just to reduce it away,
+        # H times the size of the result.  ``matmul`` broadcasts its leading
+        # batch axes, so the router's (N, 1, F, 1, H) view contracts H in place
+        # against value's (N, G, F, H, C) with no permute and no intermediate.
+        router_row = xp.reshape(
+            router, (n_batch, 1, n_focus, 1, self.n_branches)
+        )  # (N, 1, F, 1, H)
+        out = xp.matmul(router_row, value)  # (N, G, F, 1, C)
+        out = xp.reshape(out, (n_batch, n_grid, n_focus, self.channels))
 
         # === Step 3. Project back to coefficients and mix output channels ===
         return _project_frames(from_grid(out), self.out_proj, self.n_frames)
@@ -493,9 +538,13 @@ 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, ...])
+        # einsum "ndfi,dio->ndfo" as a matmul batched over the DEGREE axis.
+        #
+        # NOT as ``matmul(coeff, weight[None, ...])``: that puts N in the matmul
+        # batch, so the weight broadcasts to (N, D, i, o) and autograd has to
+        # reduce that whole expanded gradient back to the parameter shape.  See
+        # the same fix in so3.py, where it was the costliest kernel of a step.
+        return _degree_batched_matmul(xp, coeff, weight)
 
     def serialize(self) -> dict[str, Any]:
         """Serialize the FrameContract to a dict."""
@@ -575,9 +624,9 @@ 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, ...])
+        # einsum "ndfi,dio->ndfo" as a matmul batched over the DEGREE axis; see
+        # the note in FrameContract.call for why N must not be the batch axis.
+        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..58f3297811 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/lora.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/lora.py
@@ -189,10 +189,15 @@ 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" as a matmul batched over the (D, F) axes.
+        #
+        # NOT as ``matmul(x[:, :, :, None, :], weight_expanded[None, ...])``:
+        # that makes N the matmul batch, so the weight broadcasts to
+        # (N, D, F, Cin, Cout) and autograd reduces that whole expanded
+        # gradient every step.  This is the LoRA twin of the so3.py fix.
         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),

From 3c2b9bf710324a2237b2ae0e3eb7c381c95633b7 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 5 Aug 2026 11:12:33 +0800
Subject: [PATCH 05/16] fix(dpa4): serialize use_amp so a configured false is
 not silently ignored

The descriptor's use_amp flag was never written to serialize(), and
deserialize() feeds config straight into __init__, so any rebuild fell back
to the True default.  The pt_expt backend rebuilds the descriptor from that
dict, so 'use_amp: false' in the input was silently discarded and training
stayed in bfloat16 autocast; only the pt backend, which builds once from
the config, honoured it.

Caught while benchmarking: disabling AMP made pt 23% faster on a Turing GPU
(no bf16 tensor cores) while pt-expt did not move at all, and an op-level
profile showed pt-expt still spending 45% of its device time in bf16 gemm
kernels with use_amp=false.

Add the key to both the dpmodel and pt serialize configs so the two stay
key-identical and the flag survives a cross-backend round-trip.  Records
written before this change deserialize unchanged -- the key is simply
absent and __init__ supplies the default.

The pre-existing round-trip tests compare forward OUTPUTS, which cannot
catch this: dpmodel never autocasts, so the outputs agree whatever use_amp
says.  The new test pins the attribute itself, for both boolean values, and
fails on the previous code.
---
 deepmd/dpmodel/descriptor/dpa4.py             |  7 ++++++
 deepmd/pt/model/descriptor/sezm.py            |  4 ++++
 .../tests/common/dpmodel/test_descrpt_dpa4.py | 23 +++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py
index 2687d9535d..f1eadd17ed 100644
--- a/deepmd/dpmodel/descriptor/dpa4.py
+++ b/deepmd/dpmodel/descriptor/dpa4.py
@@ -2695,6 +2695,13 @@ def serialize(self) -> dict[str, Any]:
                 "mlp_bias": self.mlp_bias,
                 "exclude_types": self.exclude_types,
                 "eps": self.eps,
+                # ``use_amp`` must round-trip: the pt_expt backend rebuilds the
+                # descriptor from this dict, so omitting it silently reset a
+                # configured ``use_amp: false`` back to the True default and
+                # kept training in bfloat16 autocast.  Reading older records
+                # that lack the key still works -- deserialize passes config
+                # straight to __init__, which defaults it.
+                "use_amp": self.use_amp,
                 "trainable": self.trainable,
                 "seed": self.seed,
                 "inner_clamp_r_inner": self.inner_clamp_r_inner,
diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index d39f7a1028..1f2a1c3afd 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -2500,6 +2500,10 @@ def serialize(self) -> dict[str, Any]:
                 "mlp_bias": self.mlp_bias,
                 "exclude_types": self.exclude_types,
                 "eps": self.eps,
+                # Kept in lockstep with the dpmodel serialize contract so the
+                # two backends' records stay key-identical and ``use_amp``
+                # survives a cross-backend round-trip.
+                "use_amp": self.use_amp,
                 "trainable": self.trainable,
                 "seed": self.seed,
                 "inner_clamp_r_inner": self.inner_clamp_r_inner,
diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py
index d2b238ae8e..8ea752d640 100644
--- a/source/tests/common/dpmodel/test_descrpt_dpa4.py
+++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py
@@ -216,6 +216,29 @@ 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.
+
+        It was absent from the serialized config, so any backend that rebuilds
+        the descriptor from that dict (pt_expt does) silently reset a
+        configured ``use_amp: false`` to the True default and kept training
+        under bfloat16 autocast.  The forward-output round-trip test cannot
+        catch this: dpmodel never autocasts, so the outputs match either way --
+        only the attribute itself pins the contract.
+        """
+        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)

From 99d33ea407b29cad91c8217085ad5a9a775ec1d8 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 5 Aug 2026 16:06:51 +0800
Subject: [PATCH 06/16] feat(pt_expt): honor enable_tf32 / DP_TF32_INFER like
 the pt backend

pt_expt accepted `model.enable_tf32` and threw it away with a warning, so
DPA4/SeZM training always ran at "highest" matmul precision while the pt
backend -- reading the same input.json -- ran its training forwards under
`set_float32_matmul_precision("high")`.  On Ampere and later that is the
difference between TF32 tensor cores and fp32 CUDA cores for every matmul,
and GEMM is ~60% of compiled device time on this workload, so the two
backends were not comparable on that hardware at all.

Mirror pt's policy exactly: TRAINING forwards follow `enable_tf32`
(argcheck default True), EVAL forwards follow `DP_TF32_INFER` (0/1/2 ->
highest/high/medium, invalid values rejected).  Scope matches pt, where
argcheck declares the knob inside the dpa4 model arg block and only the
sezm builders wire it: pt_expt attaches it in `get_sezm_model` and
`get_native_spin_model`, and every other model keeps class defaults that
select full fp32 in both modes.

Ownership: `call_common` is the single owner for eager forwards -- every
pt_expt model's `forward` reaches the backbone through it, and the export
trace roots at `call_common_lower`, so the precision switch never enters an
exported graph.  The compiled path needs its own application because
`_CompiledModel.forward` bypasses `call_common` entirely; placing the
context only on the model would have left it dead on exactly the path this
is meant to speed up.  The context spans the lazy compile there, since
Inductor picks its GEMM backend while lowering.

Gating on `self.training` is what keeps the existing 1e-12 parity tests
valid: eval and export stay at "highest" unless DP_TF32_INFER asks
otherwise.
---
 deepmd/pt_expt/model/get_model.py             |  67 ++++++--
 deepmd/pt_expt/model/make_model.py            |  57 +++++++
 deepmd/pt_expt/train/training.py              |  19 ++-
 .../pt_expt/model/test_get_model_dpa4.py      | 148 +++++++++++++-----
 4 files changed, 233 insertions(+), 58 deletions(-)

diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py
index 7bdb07de57..424ddd7996 100644
--- a/deepmd/pt_expt/model/get_model.py
+++ b/deepmd/pt_expt/model/get_model.py
@@ -8,6 +8,7 @@
 
 import copy
 import logging
+import os
 
 from deepmd.dpmodel.atomic_model.dp_atomic_model import (
     DPAtomicModel,
@@ -52,8 +53,49 @@
 
 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, copied from the pt backend's
+#: ``deepmd.pt.model.model.sezm_model._TF32_INFER_PRECISION_CHOICES`` so the two
+#: backends read the same environment variable the same way.
+_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.
+
+    Mirrors the pt backend: TRAINING forwards follow ``model.enable_tf32``
+    (default ``True``) and EVAL forwards follow ``DP_TF32_INFER``.  The model
+    layer applies the policy -- see ``call_common`` in
+    :func:`deepmd.pt_expt.model.make_model.make_model` and
+    ``_CompiledModel.forward`` for the compiled path.
+
+    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(
@@ -82,17 +124,11 @@ def get_sezm_model(data: dict) -> EnergyModel:
 
     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`` follows the pt backend: TRAINING forwards run at TF32
+    ("high") matmul precision when it is true (the default), while 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(
@@ -192,8 +228,8 @@ def get_sezm_model(data: dict) -> EnergyModel:
             atom_exclude_types=data.get("atom_exclude_types", []),
             pair_exclude_types=pair_exclude_types,
         )
-        return LinearEnergyModel(atomic_model_=composed)
-    return model
+        return _apply_tf32_policy(LinearEnergyModel(atomic_model_=composed), data)
+    return _apply_tf32_policy(model, data)
 
 
 def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
@@ -257,7 +293,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:
diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py
index 663e7fb22f..13d090c0a2 100644
--- a/deepmd/pt_expt/model/make_model.py
+++ b/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,
 )
@@ -435,6 +439,59 @@ def get_min_nbor_dist(self) -> float | None:
             """Get the minimum distance between two atoms."""
             return self.min_nbor_dist
 
+        # === TF32 matmul precision ===================================
+        # Mirrors the pt backend's ``SeZMModel`` policy (see
+        # ``deepmd.pt.model.model.sezm_model.SeZMModel.tf32_precision_ctx``):
+        # TRAINING forwards follow ``model.enable_tf32`` and EVAL forwards
+        # follow ``DP_TF32_INFER``.  Both attributes are set by the DPA4/SeZM
+        # builders in ``deepmd.pt_expt.model.get_model``; every other pt_expt
+        # model keeps these defaults, which select full fp32 in both modes and
+        # therefore leave its numerics untouched.
+        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.
+
+            This is the ONE owner of matmul precision for eager forwards: every
+            pt_expt model's ``forward`` reaches the backbone through here, and
+            the export trace roots at ``call_common_lower`` instead, so the
+            precision switch never enters an exported graph.  The compiled
+            training path bypasses this method entirely and applies the same
+            policy at its own entry point (``_CompiledModel.forward``).
+
+            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().
 
diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py
index cb2fb141a6..f61cf6da3c 100644
--- a/deepmd/pt_expt/train/training.py
+++ b/deepmd/pt_expt/train/training.py
@@ -919,7 +919,24 @@ 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.
+
+        The compiled path never reaches ``call_common`` -- which owns matmul
+        precision for eager forwards -- so it applies the same policy here, at
+        its own entry point.  The context also spans the LAZY compile below:
+        Inductor selects its GEMM backend while lowering, so a precision set
+        only around the call would never reach the generated 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,
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 aa76fd7ecc..68055b1c63 100644
--- a/source/tests/pt_expt/model/test_get_model_dpa4.py
+++ b/source/tests/pt_expt/model/test_get_model_dpa4.py
@@ -271,51 +271,113 @@ def test_default_unsupported_values_pass(self) -> None:
         self.assertIsInstance(model, EnergyModel)
 
 
-# `enable_tf32` toggles TF32 matmul precision in pt but is ignored by pt_expt
-# (always "highest" precision); a truthy value must emit a warn-once message.
-@pytest.mark.parametrize("enable_tf32", [True, False])  # truthy warns, falsy silent
-def test_enable_tf32_warns_once(enable_tf32, monkeypatch) -> None:
-    import importlib
-
-    # the package __init__ rebinds the name ``get_model`` to the function, so
-    # ``import ...get_model as`` would shadow the submodule; load it explicitly
-    gm_mod = importlib.import_module("deepmd.pt_expt.model.get_model")
-
-    # reset the warn-once set so the assertion is deterministic regardless of
-    # test ordering (other get_sezm_model calls may have already warned)
-    monkeypatch.setattr(gm_mod, "_WARNED_ONCE", set())
-
-    # Count emissions on the EMITTING logger with our own handler rather than
-    # through caplog: caplog reads a root handler, so whatever global logging
-    # state earlier tests left behind (set_log_handles flips the ``deepmd``
-    # logger's propagate off and installs its own handlers) changes how many
-    # records reach it -- zero when propagation is off, more than one when the
-    # record is seen through several attached handlers.  A handler on the
-    # emitting logger sees exactly one record per ``log.warning`` call.
-    records: list[logging.LogRecord] = []
-
-    class _Collect(logging.Handler):
-        def emit(self, record: logging.LogRecord) -> None:
-            records.append(record)
-
-    handler = _Collect(level=logging.WARNING)
-    old_level = gm_mod.log.level
-    gm_mod.log.setLevel(logging.WARNING)
-    gm_mod.log.addHandler(handler)
+# === TF32 matmul precision ==================================================
+# pt_expt mirrors the pt backend (``SeZMModel.tf32_precision_ctx``): TRAINING
+# forwards follow ``model.enable_tf32`` (default True) and EVAL forwards follow
+# ``DP_TF32_INFER``.  The knob is DPA4/SeZM-scoped, matching pt, where argcheck
+# declares it inside the dpa4 model arg block.
+
+
+@pytest.mark.parametrize(
+    "enable_tf32",
+    [
+        True,  # the argcheck default; training must select TF32 ("high")
+        False,  # opt-out; training must stay at full fp32
+    ],
+)
+def test_enable_tf32_is_stored(enable_tf32) -> None:
+    """The config knob reaches the model instead of being warned away."""
+    model = get_model(_make_raw_model_config(enable_tf32=enable_tf32))
+    assert model.enable_tf32 is enable_tf32
+
+
+def test_enable_tf32_defaults_true() -> None:
+    """An absent key follows pt's ``default=True`` (argcheck.py `enable_tf32`)."""
+    raw = _make_raw_model_config()
+    assert "enable_tf32" not in raw
+    assert get_model(raw).enable_tf32 is True
+
+
+@pytest.mark.parametrize(
+    ("env_value", "expected"),
+    [
+        (None, "highest"),  # unset -> pt's "0" default, full fp32
+        ("0", "highest"),
+        ("1", "high"),
+        ("2", "medium"),
+    ],
+)
+def test_tf32_infer_precision_from_env(env_value, expected, monkeypatch) -> None:
+    """Eval precision follows ``DP_TF32_INFER``, as in the pt backend."""
+    if env_value is None:
+        monkeypatch.delenv("DP_TF32_INFER", raising=False)
+    else:
+        monkeypatch.setenv("DP_TF32_INFER", env_value)
+    assert get_model(_make_raw_model_config()).tf32_infer_precision == expected
+
+
+def test_tf32_infer_precision_rejects_garbage(monkeypatch) -> None:
+    """An unusable ``DP_TF32_INFER`` fails fast rather than silently defaulting."""
+    monkeypatch.setenv("DP_TF32_INFER", "yes")
+    with pytest.raises(ValueError, match="DP_TF32_INFER"):
+        get_model(_make_raw_model_config())
+
+
+@pytest.mark.parametrize(
+    ("enable_tf32", "training", "expected"),
+    [
+        (True, True, "high"),  # the only combination that selects TF32
+        (False, True, "highest"),  # opt-out keeps training at full fp32
+        (True, False, "highest"),  # eval ignores enable_tf32 (uses DP_TF32_INFER)
+        (False, False, "highest"),
+    ],
+)
+def test_tf32_precision_ctx_selects_and_restores(
+    enable_tf32, training, expected, monkeypatch
+) -> None:
+    """The context selects pt's precision for the mode and restores the old one.
+
+    ``torch.set_float32_matmul_precision`` is a process global, so a forward
+    that leaked its setting would silently change every later matmul in the
+    process; the restore is as much of the contract as the selection.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("tf32_precision_ctx is a no-op without CUDA")
+    monkeypatch.delenv("DP_TF32_INFER", raising=False)
+    model = get_model(_make_raw_model_config(enable_tf32=enable_tf32))
+    model.train(training)
+
+    torch.set_float32_matmul_precision("highest")
     try:
-        gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32))
-        matches = [r for r in records if "enable_tf32" in r.getMessage()]
-        if enable_tf32:
-            assert len(matches) == 1, [r.getMessage() for r in records]
-            # a second call must NOT warn again (warn-once per process)
-            records.clear()
-            gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32))
-            assert not [r for r in records if "enable_tf32" in r.getMessage()]
-        else:
-            assert not matches, [r.getMessage() for r in records]
+        with model.tf32_precision_ctx():
+            assert torch.get_float32_matmul_precision() == expected
+        assert torch.get_float32_matmul_precision() == "highest"
     finally:
-        gm_mod.log.removeHandler(handler)
-        gm_mod.log.setLevel(old_level)
+        torch.set_float32_matmul_precision("highest")
+
+
+def test_non_sezm_model_keeps_full_precision() -> None:
+    """The knob is DPA4/SeZM-scoped: other pt_expt models are untouched.
+
+    pt declares ``enable_tf32`` inside the dpa4 model arg block and wires it
+    only in its sezm builders, so a plain se_e2_a model must keep the class
+    defaults -- full fp32 in both train and eval.
+    """
+    model = get_model(
+        {
+            "type_map": ["O", "H"],
+            "descriptor": {
+                "type": "se_e2_a",
+                "sel": [4, 4],
+                "rcut": 4.0,
+                "rcut_smth": 3.5,
+                "seed": 1,
+            },
+            "fitting_net": {"seed": 1},
+        }
+    )
+    assert model.enable_tf32 is False
+    assert model.tf32_infer_precision == "highest"
 
 
 class TestNativeSpinErrorTranslation(unittest.TestCase):

From 504bb24309f332940e27f36017ed785895dd3a1a Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 5 Aug 2026 20:15:03 +0800
Subject: [PATCH 07/16] perf(dpmodel): stop spelling the DPA4 grid router as a
 degenerate GEMM

The GridBranch router contracts the branch axis H, and H is a handful (1 in
the water example).  Spelling it as `matmul(router.reshape(N, 1, F, 1, H),
value)` therefore asks cuBLAS for a batched GEMM with M=1 and K=H, which it
serves from its small-N kernels (`gemmSN_*`, `gemmk1`).

A shape-resolved profile of a compiled DPA4 training step found this to be
the single largest GEMM in the run:

    aten::bmm [[119808, 1, 1], [119808, 1, 96]]   0.0249 s/step forward

with its two backward siblings adding 0.0138 s/step -- together ~0.039 s/step
against a total pt-vs-pt_expt compiled gap of 0.055 s/step.  The batch is
N(1152) * G(104) and K is 1: no contraction is happening at all, it is a
scalar multiply routed through a GEMM kernel.

Micro-benchmarked fwd+bwd at those exact shapes:

    H=1:  matmul 7.523 ms   mul+sum 1.828 ms   (4.1x)
    H=3:  matmul 4.436 ms   mul+sum 4.453 ms   (equal)

so the broadcast form is never worse.  The comment this replaces claimed the
intermediate costs "H times the size of the result" -- true, but H is small,
and the measurement shows it does not pay for the degenerate GEMM.

This restores the spelling that 7518a417c replaced and 01c58e665 restored
once already; that revert was justified on a different workload (AMP-on
eager, where the site was invisible) and 75459610a then re-applied the
matmul as part of a broader sweep without re-measuring this site.  The
numbers above are what was missing both times.
---
 deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py | 21 +++++++++----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index e007932d0b..ef9802a4f0 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -436,18 +436,17 @@ 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 batched matmul.
+        # einsum "ngfhc,nfh->ngfc" as a broadcast multiply and reduction.
         #
-        # NOT as ``xp.sum(value * router[:, None, :, :, None], axis=3)``: that
-        # materialises the whole (N, G, F, H, C) product just to reduce it away,
-        # H times the size of the result.  ``matmul`` broadcasts its leading
-        # batch axes, so the router's (N, 1, F, 1, H) view contracts H in place
-        # against value's (N, G, F, H, C) with no permute and no intermediate.
-        router_row = xp.reshape(
-            router, (n_batch, 1, n_focus, 1, self.n_branches)
-        )  # (N, 1, F, 1, H)
-        out = xp.matmul(router_row, value)  # (N, G, F, 1, C)
-        out = xp.reshape(out, (n_batch, n_grid, n_focus, self.channels))
+        # NOT as ``matmul(reshape(router, (N, 1, F, 1, H)), value)``: the branch
+        # count H is a handful, so that spelling is a batched GEMM with M=1 and
+        # K=H, which cuBLAS serves from its small-N (``gemmSN_*`` /
+        # ``gemmk1``) kernels.  At the shapes a compiled DPA4 step actually
+        # runs -- N=1152, G=104, F=1, C=96, H=1 -- the matmul measured 7.52 ms
+        # forward+backward against 1.83 ms for this form, and it was the single
+        # largest GEMM of the step; at H=3 the two are equal (4.44 vs 4.45 ms).
+        # The intermediate this form materialises is only H 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)

From ae720432b5cf86b11164b8e593418cdaae11b487 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Thu, 6 Aug 2026 15:14:23 +0800
Subject: [PATCH 08/16] docs: shorten the comments added by this branch

The contraction and TF32 comments had grown into measurement essays. Keep
the part a reader needs -- why the obvious spelling is wrong -- and drop the
profiling detail, which belongs in the PR discussion rather than the source.

Also drops a `logging` import left unused when the enable_tf32 warn-once test
was replaced.
---
 deepmd/dpmodel/descriptor/dpa4.py             | 10 +++---
 deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py | 32 ++++++-------------
 deepmd/dpmodel/descriptor/dpa4_nn/lora.py     |  9 ++----
 deepmd/dpmodel/descriptor/dpa4_nn/so3.py      | 28 +++++-----------
 deepmd/pt/model/descriptor/sezm.py            |  5 ++-
 deepmd/pt_expt/model/get_model.py             | 22 ++++++-------
 deepmd/pt_expt/model/make_model.py            | 24 ++++++--------
 deepmd/pt_expt/train/training.py              |  9 +++---
 .../tests/common/dpmodel/test_descrpt_dpa4.py | 10 +++---
 .../pt_expt/model/test_get_model_dpa4.py      | 21 +++++-------
 10 files changed, 61 insertions(+), 109 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py
index 18277ed7c2..0e13a99f9f 100644
--- a/deepmd/dpmodel/descriptor/dpa4.py
+++ b/deepmd/dpmodel/descriptor/dpa4.py
@@ -2748,12 +2748,10 @@ def serialize(self) -> dict[str, Any]:
                 "mlp_bias": self.mlp_bias,
                 "exclude_types": self.exclude_types,
                 "eps": self.eps,
-                # ``use_amp`` must round-trip: the pt_expt backend rebuilds the
-                # descriptor from this dict, so omitting it silently reset a
-                # configured ``use_amp: false`` back to the True default and
-                # kept training in bfloat16 autocast.  Reading older records
-                # that lack the key still works -- deserialize passes config
-                # straight to __init__, which defaults it.
+                # 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,
diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index ef9802a4f0..404a23913d 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -118,12 +118,9 @@ def _degree_batched_matmul(xp: Any, coeff: Any, weight: Any) -> Any:
 
     Notes
     -----
-    The obvious spelling ``matmul(coeff, weight[None, ...])`` puts ``N`` in the
-    matmul batch, so ``weight`` is broadcast to ``(N, D, i, o)`` and autograd
-    must reduce that expanded gradient back to the parameter shape every step.
-    Batching over the small degree axis keeps ``N`` as matmul ROWS, so the
-    weight is used in place; the transposes touch only ``coeff``, which is far
-    smaller than the expanded weight would be.
+    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(
@@ -437,15 +434,10 @@ def call(
         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 multiply and reduction.
-        #
-        # NOT as ``matmul(reshape(router, (N, 1, F, 1, H)), value)``: the branch
-        # count H is a handful, so that spelling is a batched GEMM with M=1 and
-        # K=H, which cuBLAS serves from its small-N (``gemmSN_*`` /
-        # ``gemmk1``) kernels.  At the shapes a compiled DPA4 step actually
-        # runs -- N=1152, G=104, F=1, C=96, H=1 -- the matmul measured 7.52 ms
-        # forward+backward against 1.83 ms for this form, and it was the single
-        # largest GEMM of the step; at H=3 the two are equal (4.44 vs 4.45 ms).
-        # The intermediate this form materialises is only H times the result.
+        # 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 ===
@@ -537,12 +529,7 @@ 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 matmul batched over the DEGREE axis.
-        #
-        # NOT as ``matmul(coeff, weight[None, ...])``: that puts N in the matmul
-        # batch, so the weight broadcasts to (N, D, i, o) and autograd has to
-        # reduce that whole expanded gradient back to the parameter shape.  See
-        # the same fix in so3.py, where it was the costliest kernel of a step.
+        # 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]:
@@ -623,8 +610,7 @@ 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 matmul batched over the DEGREE axis; see
-        # the note in FrameContract.call for why N must not be the batch axis.
+        # 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]:
diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/lora.py b/deepmd/dpmodel/descriptor/dpa4_nn/lora.py
index 58f3297811..ba8dd1ab1a 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/lora.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/lora.py
@@ -189,12 +189,9 @@ 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 matmul batched over the (D, F) axes.
-        #
-        # NOT as ``matmul(x[:, :, :, None, :], weight_expanded[None, ...])``:
-        # that makes N the matmul batch, so the weight broadcasts to
-        # (N, D, F, Cin, Cout) and autograd reduces that whole expanded
-        # gradient every step.  This is the LoRA twin of the so3.py fix.
+        # 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(xp.permute_dims(x, (1, 2, 0, 3)), weight_expanded)
         out = xp.permute_dims(out, (2, 0, 1, 3))  # (N, D, F, Cout)
diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
index 15cd5c10a9..44c0e4cdc9 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
@@ -131,16 +131,9 @@ 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 matmul batched over the FOCUS axis.
-        #
-        # NOT as ``matmul(x[:, :, None, :], weight[None, ...])``: that makes B a
-        # batch axis, so matmul broadcasts the weight to (B, F, Cin, Cout) --
-        # inflating a few-hundred-KB parameter into hundreds of millions of
-        # elements per call, whose gradient autograd must then reduce back down
-        # (an ``ExpandBackward0`` reduce that measured as the single most
-        # expensive kernel of a DPA4 training step). Batching over F instead
-        # keeps B as matmul ROWS, so the weight is used in place and its
-        # gradient is an ordinary matmul.
+        # 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(xp.permute_dims(x, (1, 0, 2)), weight)  # (F, B, Cout)
         out = xp.permute_dims(out, (1, 0, 2))  # (B, F, Cout)
@@ -448,16 +441,11 @@ 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 matmul batched over the (D, F) axes.
-        #
-        # NOT as ``matmul(x[:, :, :, None, :], weight_expanded[None, ...])``:
-        # that makes N a batch axis, so matmul broadcasts the weight to
-        # (N, D, F, Cin, Cout) -- for the water DPA4 example, a 165K-element
-        # parameter expanded to 191M elements (~0.8 GB) on every call, whose
-        # gradient autograd then reduces back down. That ``ExpandBackward0``
-        # reduce measured at 45.6 ms per call, three calls per training step:
-        # the most expensive kernel in the run. Batching over the small (D, F)
-        # axes keeps N as matmul ROWS, so the weight is never expanded.
+        # 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)
diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index 438caf8f75..800da9453e 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -2561,9 +2561,8 @@ def serialize(self) -> dict[str, Any]:
                 "mlp_bias": self.mlp_bias,
                 "exclude_types": self.exclude_types,
                 "eps": self.eps,
-                # Kept in lockstep with the dpmodel serialize contract so the
-                # two backends' records stay key-identical and ``use_amp``
-                # survives a cross-backend round-trip.
+                # 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,
diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py
index a9e6404118..ed3a668ab5 100644
--- a/deepmd/pt_expt/model/get_model.py
+++ b/deepmd/pt_expt/model/get_model.py
@@ -58,9 +58,8 @@
 
 log = logging.getLogger(__name__)
 
-#: ``DP_TF32_INFER`` -> eval-time matmul precision, copied from the pt backend's
-#: ``deepmd.pt.model.model.sezm_model._TF32_INFER_PRECISION_CHOICES`` so the two
-#: backends read the same environment variable the same way.
+#: ``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",
@@ -71,11 +70,9 @@
 def _apply_tf32_policy(model: BaseModel, data: dict) -> BaseModel:
     """Attach the DPA4/SeZM TF32 matmul-precision policy to a built model.
 
-    Mirrors the pt backend: TRAINING forwards follow ``model.enable_tf32``
-    (default ``True``) and EVAL forwards follow ``DP_TF32_INFER``.  The model
-    layer applies the policy -- see ``call_common`` in
-    :func:`deepmd.pt_expt.model.make_model.make_model` and
-    ``_CompiledModel.forward`` for the compiled path.
+    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
     ----------
@@ -132,9 +129,9 @@ def get_sezm_model(data: dict) -> BaseModel:
 
     Notes
     -----
-    ``enable_tf32`` follows the pt backend: TRAINING forwards run at TF32
-    ("high") matmul precision when it is true (the default), while EVAL
-    forwards follow ``DP_TF32_INFER``.  See :func:`_apply_tf32_policy`.
+    ``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 "spin" in data:
@@ -208,8 +205,7 @@ def get_sezm_model(data: dict) -> BaseModel:
         pair_exclude_types=pair_exclude_types,
     )
     if bridging_enabled:
-        # Upstream factored the bridging composition into ``_compose_bridging``;
-        # the TF32 policy attaches to whichever model is returned.
+        # 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)
 
diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py
index 7ed37e156a..340acbe1e7 100644
--- a/deepmd/pt_expt/model/make_model.py
+++ b/deepmd/pt_expt/model/make_model.py
@@ -474,14 +474,11 @@ def get_min_nbor_dist(self) -> float | None:
             """Get the minimum distance between two atoms."""
             return self.min_nbor_dist
 
-        # === TF32 matmul precision ===================================
-        # Mirrors the pt backend's ``SeZMModel`` policy (see
-        # ``deepmd.pt.model.model.sezm_model.SeZMModel.tf32_precision_ctx``):
-        # TRAINING forwards follow ``model.enable_tf32`` and EVAL forwards
-        # follow ``DP_TF32_INFER``.  Both attributes are set by the DPA4/SeZM
-        # builders in ``deepmd.pt_expt.model.get_model``; every other pt_expt
-        # model keeps these defaults, which select full fp32 in both modes and
-        # therefore leave its numerics untouched.
+        # === 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"
 
@@ -512,12 +509,11 @@ def tf32_precision_ctx(self) -> Generator[None, None, None]:
         def call_common(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
             """Run the shared dense/graph forward under the TF32 policy.
 
-            This is the ONE owner of matmul precision for eager forwards: every
-            pt_expt model's ``forward`` reaches the backbone through here, and
-            the export trace roots at ``call_common_lower`` instead, so the
-            precision switch never enters an exported graph.  The compiled
-            training path bypasses this method entirely and applies the same
-            policy at its own entry point (``_CompiledModel.forward``).
+            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
             -------
diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py
index 97775133a4..e93c1a03aa 100644
--- a/deepmd/pt_expt/train/training.py
+++ b/deepmd/pt_expt/train/training.py
@@ -964,11 +964,10 @@ def __getattr__(self, name: str) -> Any:
     def forward(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
         """Run the compiled forward under the wrapped model's TF32 policy.
 
-        The compiled path never reaches ``call_common`` -- which owns matmul
-        precision for eager forwards -- so it applies the same policy here, at
-        its own entry point.  The context also spans the LAZY compile below:
-        Inductor selects its GEMM backend while lowering, so a precision set
-        only around the call would never reach the generated kernels.
+        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
         -------
diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py
index 3d017b060a..c64eb9152d 100644
--- a/source/tests/common/dpmodel/test_descrpt_dpa4.py
+++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py
@@ -275,12 +275,10 @@ def test_supported_feature_roundtrip(self, overrides) -> None:
     def test_use_amp_survives_roundtrip(self, use_amp) -> None:
         """``use_amp`` must round-trip through serialize/deserialize.
 
-        It was absent from the serialized config, so any backend that rebuilds
-        the descriptor from that dict (pt_expt does) silently reset a
-        configured ``use_amp: false`` to the True default and kept training
-        under bfloat16 autocast.  The forward-output round-trip test cannot
-        catch this: dpmodel never autocasts, so the outputs match either way --
-        only the attribute itself pins the contract.
+        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
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 68055b1c63..2e2dc50a02 100644
--- a/source/tests/pt_expt/model/test_get_model_dpa4.py
+++ b/source/tests/pt_expt/model/test_get_model_dpa4.py
@@ -2,7 +2,6 @@
 """Tests for the DPA4/SeZM model-type dispatch in pt_expt ``get_model``."""
 
 import copy
-import logging
 import unittest
 
 import pytest
@@ -271,11 +270,9 @@ def test_default_unsupported_values_pass(self) -> None:
         self.assertIsInstance(model, EnergyModel)
 
 
-# === TF32 matmul precision ==================================================
-# pt_expt mirrors the pt backend (``SeZMModel.tf32_precision_ctx``): TRAINING
-# forwards follow ``model.enable_tf32`` (default True) and EVAL forwards follow
-# ``DP_TF32_INFER``.  The knob is DPA4/SeZM-scoped, matching pt, where argcheck
-# declares it inside the dpa4 model arg block.
+# === TF32 matmul precision ===
+# Same policy as pt: training follows ``enable_tf32`` (default True), eval
+# follows ``DP_TF32_INFER``.  Like pt, the knob is DPA4/SeZM-scoped.
 
 
 @pytest.mark.parametrize(
@@ -335,11 +332,10 @@ def test_tf32_infer_precision_rejects_garbage(monkeypatch) -> None:
 def test_tf32_precision_ctx_selects_and_restores(
     enable_tf32, training, expected, monkeypatch
 ) -> None:
-    """The context selects pt's precision for the mode and restores the old one.
+    """The context picks the right precision per mode, then restores it.
 
-    ``torch.set_float32_matmul_precision`` is a process global, so a forward
-    that leaked its setting would silently change every later matmul in the
-    process; the restore is as much of the contract as the selection.
+    ``set_float32_matmul_precision`` is a process global, so a leaked setting
+    would change every later matmul; restoring matters as much as selecting.
     """
     if not torch.cuda.is_available():
         pytest.skip("tf32_precision_ctx is a no-op without CUDA")
@@ -359,9 +355,8 @@ def test_tf32_precision_ctx_selects_and_restores(
 def test_non_sezm_model_keeps_full_precision() -> None:
     """The knob is DPA4/SeZM-scoped: other pt_expt models are untouched.
 
-    pt declares ``enable_tf32`` inside the dpa4 model arg block and wires it
-    only in its sezm builders, so a plain se_e2_a model must keep the class
-    defaults -- full fp32 in both train and eval.
+    pt wires ``enable_tf32`` only in its sezm builders, so a plain se_e2_a
+    model keeps the class defaults: full fp32 in both train and eval.
     """
     model = get_model(
         {

From 1574442049b4cf3c66851c69e9c17eb75d10e325 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Thu, 6 Aug 2026 15:24:14 +0800
Subject: [PATCH 09/16] revert(dpmodel): restore master's GridBranch router
 line

The router site ends up byte-identical to master: 7518a417c replaced the
broadcast sum with a matmul, 504bb2430 put the sum back, and the net diff was
a one-line comment swapped for five -- losing master's (N, G, F, C) shape
annotation on the way.  Restore master's line exactly, so the branch touches
this site not at all.

The degenerate GEMM that profiling found there was self-inflicted: it existed
only on this branch, never on master, so "fixing" it delivered nothing.

Also corrects the so3 ChannelLinear comment, which claimed the contraction is
batched over the focus axis.  What matters is that B stays the GEMM rows; at
n_focus=1 -- every shipped config -- both permutes are contiguous views and
the whole thing is one (B, Cin) x (Cin, Cout) GEMM at no copy cost.
---
 deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py | 8 ++------
 deepmd/dpmodel/descriptor/dpa4_nn/so3.py      | 7 ++++---
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index 404a23913d..fcf7657e0b 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -433,12 +433,8 @@ 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 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)
+        # 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)
 
         # === Step 3. Project back to coefficients and mix output channels ===
         return _project_frames(from_grid(out), self.out_proj, self.n_frames)
diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
index 44c0e4cdc9..b2cd2a18fe 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
@@ -131,9 +131,10 @@ 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", 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.
+        # 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(xp.permute_dims(x, (1, 0, 2)), weight)  # (F, B, Cout)
         out = xp.permute_dims(out, (1, 0, 2))  # (B, F, Cout)

From 1e56cf6d5bc7b8b942b5edb92a965e86a422c8ca Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Thu, 6 Aug 2026 16:16:48 +0800
Subject: [PATCH 10/16] Revert "feat(pt_expt): honor enable_tf32 /
 DP_TF32_INFER like the pt backend"

This reverts the pt_expt TF32 policy (99d33ea40 plus its comment edits in
ae720432b), restoring the warn-and-ignore behavior on master.

The knob is unrelated to this PR's measured speedup (the benchmark card has
no TF32 silicon; the whole 1.69x/3.01x gain comes from the contraction fix),
its benefit was never measured, and PR #5958 owns the pt_expt training
runtime alignment -- including the documented position that pt_expt runs at
'highest' matmul precision.  Keeping a second, contradicting implementation
here would split ownership of the same policy across two PRs.
---
 deepmd/pt_expt/model/get_model.py             |  65 ++------
 deepmd/pt_expt/model/make_model.py            |  53 -------
 deepmd/pt_expt/train/training.py              |  18 +--
 .../pt_expt/model/test_get_model_dpa4.py      | 145 ++++++------------
 4 files changed, 59 insertions(+), 222 deletions(-)

diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py
index ed3a668ab5..50f60ecf49 100644
--- a/deepmd/pt_expt/model/get_model.py
+++ b/deepmd/pt_expt/model/get_model.py
@@ -8,7 +8,6 @@
 
 import copy
 import logging
-import os
 from typing import (
     TYPE_CHECKING,
 )
@@ -58,46 +57,8 @@
 
 log = logging.getLogger(__name__)
 
-#: ``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
+# Warn at most once per process for backend-ignored switches (keyed by name).
+_WARNED_ONCE: set[str] = set()
 
 
 _model_factory = BackendModelFactory(
@@ -129,11 +90,17 @@ def get_sezm_model(data: dict) -> BaseModel:
 
     Notes
     -----
-    ``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`.
+    ``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.
     """
     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(
@@ -205,9 +172,8 @@ def get_sezm_model(data: dict) -> BaseModel:
         pair_exclude_types=pair_exclude_types,
     )
     if bridging_enabled:
-        # 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)
+        return _compose_bridging(model, data, bridging_method)
+    return model
 
 
 def _compose_bridging(
@@ -372,10 +338,7 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
             "spin scheme 'native' requires an atomic model declaring "
             "supports_native_spin()"
         )
-    return _apply_tf32_policy(
-        NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin),
-        data,
-    )
+    return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin)
 
 
 def get_linear_model(model_params: dict) -> BaseModel:
diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py
index 340acbe1e7..c6def8f136 100644
--- a/deepmd/pt_expt/model/make_model.py
+++ b/deepmd/pt_expt/model/make_model.py
@@ -1,10 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import contextlib
 import math
 import types
-from collections.abc import (
-    Generator,
-)
 from typing import (
     Any,
 )
@@ -474,55 +470,6 @@ 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().
 
diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py
index e93c1a03aa..2a8165c90f 100644
--- a/deepmd/pt_expt/train/training.py
+++ b/deepmd/pt_expt/train/training.py
@@ -961,23 +961,7 @@ def __getattr__(self, name: str) -> Any:
         except AttributeError:
             return getattr(self.original_model, name)
 
-    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(
+    def forward(
         self,
         coord: torch.Tensor,
         atype: torch.Tensor,
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 2e2dc50a02..aa76fd7ecc 100644
--- a/source/tests/pt_expt/model/test_get_model_dpa4.py
+++ b/source/tests/pt_expt/model/test_get_model_dpa4.py
@@ -2,6 +2,7 @@
 """Tests for the DPA4/SeZM model-type dispatch in pt_expt ``get_model``."""
 
 import copy
+import logging
 import unittest
 
 import pytest
@@ -270,109 +271,51 @@ def test_default_unsupported_values_pass(self) -> None:
         self.assertIsInstance(model, EnergyModel)
 
 
-# === TF32 matmul precision ===
-# Same policy as pt: training follows ``enable_tf32`` (default True), eval
-# follows ``DP_TF32_INFER``.  Like pt, the knob is DPA4/SeZM-scoped.
-
-
-@pytest.mark.parametrize(
-    "enable_tf32",
-    [
-        True,  # the argcheck default; training must select TF32 ("high")
-        False,  # opt-out; training must stay at full fp32
-    ],
-)
-def test_enable_tf32_is_stored(enable_tf32) -> None:
-    """The config knob reaches the model instead of being warned away."""
-    model = get_model(_make_raw_model_config(enable_tf32=enable_tf32))
-    assert model.enable_tf32 is enable_tf32
-
-
-def test_enable_tf32_defaults_true() -> None:
-    """An absent key follows pt's ``default=True`` (argcheck.py `enable_tf32`)."""
-    raw = _make_raw_model_config()
-    assert "enable_tf32" not in raw
-    assert get_model(raw).enable_tf32 is True
-
-
-@pytest.mark.parametrize(
-    ("env_value", "expected"),
-    [
-        (None, "highest"),  # unset -> pt's "0" default, full fp32
-        ("0", "highest"),
-        ("1", "high"),
-        ("2", "medium"),
-    ],
-)
-def test_tf32_infer_precision_from_env(env_value, expected, monkeypatch) -> None:
-    """Eval precision follows ``DP_TF32_INFER``, as in the pt backend."""
-    if env_value is None:
-        monkeypatch.delenv("DP_TF32_INFER", raising=False)
-    else:
-        monkeypatch.setenv("DP_TF32_INFER", env_value)
-    assert get_model(_make_raw_model_config()).tf32_infer_precision == expected
-
-
-def test_tf32_infer_precision_rejects_garbage(monkeypatch) -> None:
-    """An unusable ``DP_TF32_INFER`` fails fast rather than silently defaulting."""
-    monkeypatch.setenv("DP_TF32_INFER", "yes")
-    with pytest.raises(ValueError, match="DP_TF32_INFER"):
-        get_model(_make_raw_model_config())
-
-
-@pytest.mark.parametrize(
-    ("enable_tf32", "training", "expected"),
-    [
-        (True, True, "high"),  # the only combination that selects TF32
-        (False, True, "highest"),  # opt-out keeps training at full fp32
-        (True, False, "highest"),  # eval ignores enable_tf32 (uses DP_TF32_INFER)
-        (False, False, "highest"),
-    ],
-)
-def test_tf32_precision_ctx_selects_and_restores(
-    enable_tf32, training, expected, monkeypatch
-) -> None:
-    """The context picks the right precision per mode, then restores it.
-
-    ``set_float32_matmul_precision`` is a process global, so a leaked setting
-    would change every later matmul; restoring matters as much as selecting.
-    """
-    if not torch.cuda.is_available():
-        pytest.skip("tf32_precision_ctx is a no-op without CUDA")
-    monkeypatch.delenv("DP_TF32_INFER", raising=False)
-    model = get_model(_make_raw_model_config(enable_tf32=enable_tf32))
-    model.train(training)
-
-    torch.set_float32_matmul_precision("highest")
+# `enable_tf32` toggles TF32 matmul precision in pt but is ignored by pt_expt
+# (always "highest" precision); a truthy value must emit a warn-once message.
+@pytest.mark.parametrize("enable_tf32", [True, False])  # truthy warns, falsy silent
+def test_enable_tf32_warns_once(enable_tf32, monkeypatch) -> None:
+    import importlib
+
+    # the package __init__ rebinds the name ``get_model`` to the function, so
+    # ``import ...get_model as`` would shadow the submodule; load it explicitly
+    gm_mod = importlib.import_module("deepmd.pt_expt.model.get_model")
+
+    # reset the warn-once set so the assertion is deterministic regardless of
+    # test ordering (other get_sezm_model calls may have already warned)
+    monkeypatch.setattr(gm_mod, "_WARNED_ONCE", set())
+
+    # Count emissions on the EMITTING logger with our own handler rather than
+    # through caplog: caplog reads a root handler, so whatever global logging
+    # state earlier tests left behind (set_log_handles flips the ``deepmd``
+    # logger's propagate off and installs its own handlers) changes how many
+    # records reach it -- zero when propagation is off, more than one when the
+    # record is seen through several attached handlers.  A handler on the
+    # emitting logger sees exactly one record per ``log.warning`` call.
+    records: list[logging.LogRecord] = []
+
+    class _Collect(logging.Handler):
+        def emit(self, record: logging.LogRecord) -> None:
+            records.append(record)
+
+    handler = _Collect(level=logging.WARNING)
+    old_level = gm_mod.log.level
+    gm_mod.log.setLevel(logging.WARNING)
+    gm_mod.log.addHandler(handler)
     try:
-        with model.tf32_precision_ctx():
-            assert torch.get_float32_matmul_precision() == expected
-        assert torch.get_float32_matmul_precision() == "highest"
+        gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32))
+        matches = [r for r in records if "enable_tf32" in r.getMessage()]
+        if enable_tf32:
+            assert len(matches) == 1, [r.getMessage() for r in records]
+            # a second call must NOT warn again (warn-once per process)
+            records.clear()
+            gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32))
+            assert not [r for r in records if "enable_tf32" in r.getMessage()]
+        else:
+            assert not matches, [r.getMessage() for r in records]
     finally:
-        torch.set_float32_matmul_precision("highest")
-
-
-def test_non_sezm_model_keeps_full_precision() -> None:
-    """The knob is DPA4/SeZM-scoped: other pt_expt models are untouched.
-
-    pt wires ``enable_tf32`` only in its sezm builders, so a plain se_e2_a
-    model keeps the class defaults: full fp32 in both train and eval.
-    """
-    model = get_model(
-        {
-            "type_map": ["O", "H"],
-            "descriptor": {
-                "type": "se_e2_a",
-                "sel": [4, 4],
-                "rcut": 4.0,
-                "rcut_smth": 3.5,
-                "seed": 1,
-            },
-            "fitting_net": {"seed": 1},
-        }
-    )
-    assert model.enable_tf32 is False
-    assert model.tf32_infer_precision == "highest"
+        gm_mod.log.removeHandler(handler)
+        gm_mod.log.setLevel(old_level)
 
 
 class TestNativeSpinErrorTranslation(unittest.TestCase):

From 63071be846af5d1a58e520837b9657bb3b125a71 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 14 Aug 2026 14:53:42 +0800
Subject: [PATCH 11/16] fix(dpmodel): preserve empty batches in the
 degree-batched contraction

Reshaping with -1 cannot be inferred from a zero-element array, so
FrameContract/FrameExpand raised on N == 0 (empty graph/edge set, or a
distributed rank owning no nodes) where the previous broadcasted matmul
returned an empty result. Use explicit channel widths from the weight
shape in both reshapes; regression test covers both mixers on the numpy
and torch namespaces.
---
 deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py |  9 ++++-
 .../common/dpmodel/test_dpa4_frame_mixers.py  | 39 +++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index fcf7657e0b..05098aecaf 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -123,11 +123,16 @@ def _degree_batched_matmul(xp: Any, coeff: Any, weight: Any) -> Any:
     every backward.  The transposes touch only ``coeff``, which is smaller.
     """
     n_batch, coeff_dim, n_focus, _ = coeff.shape
+    # explicit channel widths: `-1` cannot be inferred when N == 0 (empty
+    # graph/edge batch, or a distributed rank owning no nodes)
+    input_dim = weight.shape[-2]
+    output_dim = weight.shape[-1]
     coeff_d = xp.reshape(
-        xp.permute_dims(coeff, (1, 0, 2, 3)), (coeff_dim, n_batch * n_focus, -1)
+        xp.permute_dims(coeff, (1, 0, 2, 3)),
+        (coeff_dim, n_batch * n_focus, input_dim),
     )  # (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))
+    out = xp.reshape(out, (coeff_dim, n_batch, n_focus, output_dim))
     return xp.permute_dims(out, (1, 0, 2, 3))  # (N, D, F, o)
 
 
diff --git a/source/tests/common/dpmodel/test_dpa4_frame_mixers.py b/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
index 38498c84ec..84c71f8cf3 100644
--- a/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
+++ b/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
@@ -211,3 +211,42 @@ 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:
+    """``N == 0`` must yield an empty result, not a reshape error.
+
+    Reachable when the cross-grid leading axis is an empty graph/edge set
+    or a distributed rank owns no nodes; the degree-batched contraction
+    must keep the broadcasted matmul's empty-batch behavior (explicit
+    channel widths -- ``-1`` cannot be inferred from zero elements).
+    """
+    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)

From fa465706e73712af04feebe047db3100188be416 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 14 Aug 2026 15:51:57 +0800
Subject: [PATCH 12/16] fix(pt_expt): keep use_amp out of serialization, fix
 the assembly boundary

use_amp is a runtime/training policy, not model state: revert the
dpa4/sezm serialize additions (a use_amp record leaks a torch runtime
option into cross-backend records -- the jax deserializer rejects
use_amp=true -- and #5963 established that checkpoints must not carry
the AMP switch).

The real pt_expt bug is in model assembly: make_model handed the raw
dpmodel atomic class to the dpmodel CM, so the constructed atomic model
was converted through the auto-wrap serialize()/deserialize() round-trip
and every runtime-only option on the live descriptor was reset to its
constructor default. Hand the CM the auto-wrapped atomic class instead:
the atomic model is constructed directly as a torch module and the live
(already wrapped) descriptor/fitting are kept as-is -- no round-trip.

Regression tests exercise the public construction path (get_model with
descriptor.use_amp=false) and pin that the portable record does not
carry use_amp.
---
 deepmd/dpmodel/descriptor/dpa4.py             |  5 --
 deepmd/pt/model/descriptor/sezm.py            |  3 -
 deepmd/pt_expt/common.py                      | 38 +++++++++---
 deepmd/pt_expt/model/make_model.py            | 10 +++-
 .../tests/common/dpmodel/test_descrpt_dpa4.py | 30 ++++------
 .../pt_expt/model/test_get_model_dpa4.py      | 58 +++++++++++++++++++
 6 files changed, 108 insertions(+), 36 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py
index 8af0285128..f1a08123c4 100644
--- a/deepmd/dpmodel/descriptor/dpa4.py
+++ b/deepmd/dpmodel/descriptor/dpa4.py
@@ -2806,11 +2806,6 @@ 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,
diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index a3860dc2f6..6e05ae884f 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -2556,9 +2556,6 @@ 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,
diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py
index 93092ecd9c..0ce1c418ce 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.
 
+    Constructing this class DIRECTLY (instead of building the raw dpmodel
+    class and converting the instance afterwards) is the lossless path:
+    instance conversion round-trips through ``serialize()``/``deserialize()``,
+    which by design drops runtime-only configuration (e.g. the DPA4
+    ``use_amp`` switch) from 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/make_model.py b/deepmd/pt_expt/model/make_model.py
index f59d507727..93b459afd0 100644
--- a/deepmd/pt_expt/model/make_model.py
+++ b/deepmd/pt_expt/model/make_model.py
@@ -24,6 +24,7 @@
     cuda_infer_level,
 )
 from deepmd.pt_expt.common import (
+    auto_wrapped_class,
     torch_module,
 )
 from deepmd.pt_expt.utils.graph_builder import (
@@ -422,7 +423,14 @@ def make_model(
         The model.
 
     """
-    DPModel = make_model_dp(T_AtomicModel)
+    # Hand the dpmodel CM the WRAPPED atomic class so `self.atomic_model =
+    # T_AtomicModel(...)` constructs the pt_expt module directly with the
+    # live (already wrapped) descriptor/fitting. Passing the raw dpmodel
+    # class instead would build a raw atomic model and convert the instance
+    # through a serialize()/deserialize() round-trip, which drops
+    # runtime-only configuration (e.g. the DPA4 `use_amp` switch) that the
+    # portable record deliberately does not carry.
+    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 83a344a46d..cb1465ae8a 100644
--- a/source/tests/common/dpmodel/test_descrpt_dpa4.py
+++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py
@@ -265,26 +265,18 @@ 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.
+    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=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
+        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/pt_expt/model/test_get_model_dpa4.py b/source/tests/pt_expt/model/test_get_model_dpa4.py
index d0990c0aa1..fe30edf604 100644
--- a/source/tests/pt_expt/model/test_get_model_dpa4.py
+++ b/source/tests/pt_expt/model/test_get_model_dpa4.py
@@ -285,3 +285,61 @@ def test_unrelated_construction_error_propagates(self) -> None:
 
 if __name__ == "__main__":
     unittest.main()
+
+
+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

From 8bce829481aa89a62005960b553a2f4b8f034dce Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 14 Aug 2026 17:00:18 +0800
Subject: [PATCH 13/16] fix: address review round 2 (D,F batching, lossless
 compositions, test placement)

- Frame mixers now batch over (D, F) in BOTH dpmodel and pt: expanding
  the small weight across F (D*F*i*o elements) replaces the materialized
  permuted coefficient copy (N*D*F*i elements, ratio N/o) that both the
  previous helper and pt's einsum lowering incurred. No reshape is
  involved, so the N == 0 empty-batch case flows through naturally; the
  empty-axis regression stays and an F > 1 forward/backward parity test
  pins both lowerings against the einsum contract for input and weight
  gradients. The stale 'broadcast batched matmul' description in the
  mixers parity-test docstring is replaced by the backend-independent
  contract.
- Composition assembly is lossless like the standard path: pt_expt now
  constructs wrapped DPAtomicModel/PairTabAtomicModel/InnerPotential
  classes directly (module-level auto_wrapped_class bindings, also
  passed to the backend factory), and _compose_bridging passes
  constructor args instead of a populated raw atomic_model_ instance,
  so neither the ZBL composition nor the explicit linear_ener route
  round-trips live children through serialize()/deserialize().
  Regression tests cover both public routes with descriptor.use_amp
  false on the learned child.
- The TestUseAmpSurvivesAssembly class moved before the __main__ block
  it had landed after, so direct file execution discovers it too.
---
 deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py |  28 ++---
 .../pt/model/descriptor/sezm_nn/grid_net.py   |  18 +++-
 deepmd/pt_expt/model/get_model.py             |  31 ++++--
 .../common/dpmodel/test_dpa4_frame_mixers.py  | 101 +++++++++++++++++-
 .../pt_expt/model/test_get_model_dpa4.py      |  52 ++++++++-
 5 files changed, 194 insertions(+), 36 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index 05098aecaf..dca7d7a36f 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -118,22 +118,16 @@ def _degree_batched_matmul(xp: Any, coeff: Any, weight: Any) -> Any:
 
     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.
+    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.
     """
-    n_batch, coeff_dim, n_focus, _ = coeff.shape
-    # explicit channel widths: `-1` cannot be inferred when N == 0 (empty
-    # graph/edge batch, or a distributed rank owning no nodes)
-    input_dim = weight.shape[-2]
-    output_dim = weight.shape[-1]
-    coeff_d = xp.reshape(
-        xp.permute_dims(coeff, (1, 0, 2, 3)),
-        (coeff_dim, n_batch * n_focus, input_dim),
-    )  # (D, N*F, i)
-    out = xp.matmul(coeff_d, weight)  # (D, N*F, o)
-    out = xp.reshape(out, (coeff_dim, n_batch, n_focus, output_dim))
-    return xp.permute_dims(out, (1, 0, 2, 3))  # (N, D, F, o)
+    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:
@@ -530,7 +524,7 @@ 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)
-        # Batched over the degree axis, never over N -- see the helper's note.
+        # 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]:
@@ -611,7 +605,7 @@ 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)
-        # Batched over the degree axis, never over N -- see the helper's note.
+        # 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]:
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/model/get_model.py b/deepmd/pt_expt/model/get_model.py
index f0d87a9118..4ecaac0a16 100644
--- a/deepmd/pt_expt/model/get_model.py
+++ b/deepmd/pt_expt/model/get_model.py
@@ -13,14 +13,20 @@
 )
 
 from deepmd.dpmodel.atomic_model.dp_atomic_model import (
-    DPAtomicModel,
+    DPAtomicModel as DPAtomicModelDP,
+)
+from deepmd.dpmodel.atomic_model.inner_potential import (
+    InnerPotentialAtomicModel as InnerPotentialAtomicModelDP,
 )
 from deepmd.dpmodel.atomic_model.pairtab_atomic_model import (
-    PairTabAtomicModel,
+    PairTabAtomicModel as PairTabAtomicModelDP,
 )
 from deepmd.dpmodel.model.model_factory import (
     BackendModelFactory,
 )
+from deepmd.pt_expt.common import (
+    auto_wrapped_class,
+)
 from deepmd.dpmodel.model.model_factory import (
     get_spin_model as get_spin_model_from_factory,
 )
@@ -61,6 +67,14 @@
 _WARNED_ONCE: set[str] = set()
 
 
+# Constructing the WRAPPED atomic classes directly is the lossless assembly
+# rule: building a raw dpmodel instance and converting it afterwards
+# round-trips through serialize()/deserialize() and drops runtime-only
+# configuration (e.g. the DPA4 `use_amp` switch) from the live children.
+DPAtomicModel = auto_wrapped_class(DPAtomicModelDP)
+PairTabAtomicModel = auto_wrapped_class(PairTabAtomicModelDP)
+InnerPotentialAtomicModel = auto_wrapped_class(InnerPotentialAtomicModelDP)
+
 _model_factory = BackendModelFactory(
     descriptor_base=BaseDescriptor,
     fitting_base=BaseFitting,
@@ -205,12 +219,6 @@ def _compose_bridging(
     LinearEnergyModel
         A composition over ``[learned, InnerPotential]``.
     """
-    from deepmd.dpmodel.atomic_model.inner_potential import (
-        InnerPotentialAtomicModel,
-    )
-    from deepmd.dpmodel.atomic_model.linear_atomic_model import (
-        LinearEnergyAtomicModel,
-    )
     from deepmd.pt_expt.model.dp_linear_model import (
         LinearEnergyModel,
     )
@@ -222,7 +230,11 @@ def _compose_bridging(
         rcut=descriptor.get_rcut(),
         sel=descriptor.get_sel(),
     )
-    composed = LinearEnergyAtomicModel(
+    # Constructor-args form (NOT a pre-built `atomic_model_=` instance): the
+    # CM then constructs the WRAPPED composition class directly and keeps the
+    # live children -- assigning a populated raw dpmodel composition would
+    # convert it through the lossy serialize()/deserialize() round-trip.
+    return LinearEnergyModel(
         models=[model.atomic_model, zbl_atomic],
         type_map=data["type_map"],
         weights="sum",
@@ -231,7 +243,6 @@ def _compose_bridging(
         atom_exclude_types=data.get("atom_exclude_types", []),
         pair_exclude_types=data.get("pair_exclude_types", []),
     )
-    return LinearEnergyModel(atomic_model_=composed)
 
 
 def get_standard_model(data: dict) -> BaseModel:
diff --git a/source/tests/common/dpmodel/test_dpa4_frame_mixers.py b/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
index 84c71f8cf3..9faa71b4ac 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
@@ -250,3 +253,95 @@ def test_empty_batch_passes_through(cls) -> None:
     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
+    from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import (
+        _degree_batched_matmul,
+    )
+
+    import array_api_compat
+
+    coeff_dp = coeff.detach().clone().requires_grad_(True)
+    dp_out = _degree_batched_matmul(
+        array_api_compat.array_namespace(coeff_dp),
+        coeff_dp,
+        pt_mod.weight.index_select(0, pt_mod.degree_index).detach(),
+    )
+    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
+    )
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 fe30edf604..6d18699fd5 100644
--- a/source/tests/pt_expt/model/test_get_model_dpa4.py
+++ b/source/tests/pt_expt/model/test_get_model_dpa4.py
@@ -283,10 +283,6 @@ def test_unrelated_construction_error_propagates(self) -> None:
             get_model(raw)
 
 
-if __name__ == "__main__":
-    unittest.main()
-
-
 class TestUseAmpSurvivesAssembly(unittest.TestCase):
     """``use_amp`` is runtime policy: it must survive model ASSEMBLY without
     entering the portable serialization record (which the JAX deserializer
@@ -343,3 +339,51 @@ def test_standard_type_keeps_use_amp_false(self) -> None:
         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()

From e2f54c8fb639026946cb1e26cc2586f56ad1a934 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 14 Aug 2026 09:01:20 +0000
Subject: [PATCH 14/16] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt_expt/model/get_model.py                    | 10 ++++------
 .../tests/common/dpmodel/test_dpa4_frame_mixers.py   | 12 ++++--------
 2 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py
index 4ecaac0a16..85e24e25f5 100644
--- a/deepmd/pt_expt/model/get_model.py
+++ b/deepmd/pt_expt/model/get_model.py
@@ -12,9 +12,7 @@
     TYPE_CHECKING,
 )
 
-from deepmd.dpmodel.atomic_model.dp_atomic_model import (
-    DPAtomicModel as DPAtomicModelDP,
-)
+from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP
 from deepmd.dpmodel.atomic_model.inner_potential import (
     InnerPotentialAtomicModel as InnerPotentialAtomicModelDP,
 )
@@ -24,12 +22,12 @@
 from deepmd.dpmodel.model.model_factory import (
     BackendModelFactory,
 )
-from deepmd.pt_expt.common import (
-    auto_wrapped_class,
-)
 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,
 )
diff --git a/source/tests/common/dpmodel/test_dpa4_frame_mixers.py b/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
index 9faa71b4ac..09200500d5 100644
--- a/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
+++ b/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
@@ -276,15 +276,11 @@ def test_focus_batched_lowering_matches_einsum_backward(kind) -> None:
     rng = np.random.default_rng(2026)
 
     if kind == "contract":
-        from deepmd.pt.model.descriptor.sezm_nn.grid_net import (
-            FrameContract as PTMixer,
-        )
+        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,
-        )
+        from deepmd.pt.model.descriptor.sezm_nn.grid_net import FrameExpand as PTMixer
 
         in_dim = channels
     pt_mod = PTMixer(
@@ -326,12 +322,12 @@ def test_focus_batched_lowering_matches_einsum_backward(kind) -> None:
     )
 
     # 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,
     )
 
-    import array_api_compat
-
     coeff_dp = coeff.detach().clone().requires_grad_(True)
     dp_out = _degree_batched_matmul(
         array_api_compat.array_namespace(coeff_dp),

From 81e756d30c4dc270b24a4288eea0813de06cf4d2 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Sat, 15 Aug 2026 16:32:43 +0800
Subject: [PATCH 15/16] test,docs: pin the SO3/LoRA contraction contract and
 dedup the assembly invariant

- so3.py: replace the benchmark anecdote with the structural rationale for
  batching over (D, F) instead of broadcasting the weight over the node axis.
- frame-mixer tests: state the empty-batch contract without referring to a
  transient call site, and compare the dp lowering's weight gradient against
  the pt module's (was forward-only).
- add test_lora_so3_call_matches_einsum_contract: pins LoRASO3.call against
  the explicit einsum reference on both the numpy and torch namespaces.
- state the auto_wrapped_class invariant once in its docstring; the four
  call sites now reference it in one line each.
---
 deepmd/dpmodel/descriptor/dpa4_nn/so3.py      |  8 +--
 deepmd/pt_expt/common.py                      | 10 ++--
 deepmd/pt_expt/model/get_model.py             | 12 ++--
 deepmd/pt_expt/model/make_model.py            |  9 +--
 .../common/dpmodel/test_dpa4_frame_mixers.py  | 14 +++--
 source/tests/common/dpmodel/test_dpa4_lora.py | 58 ++++++++++++++++++-
 6 files changed, 80 insertions(+), 31 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
index b2cd2a18fe..a5f1694942 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py
@@ -442,11 +442,9 @@ 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", 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.
+        # 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)
diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py
index 0ce1c418ce..076b9b6a04 100644
--- a/deepmd/pt_expt/common.py
+++ b/deepmd/pt_expt/common.py
@@ -143,11 +143,11 @@ def auto_wrapped_class(cls: type) -> type:
     then applies ``torch_module`` to get full ``__setattr__`` / post-init
     list conversion.  The wrapped class is cached per dpmodel type.
 
-    Constructing this class DIRECTLY (instead of building the raw dpmodel
-    class and converting the instance afterwards) is the lossless path:
-    instance conversion round-trips through ``serialize()``/``deserialize()``,
-    which by design drops runtime-only configuration (e.g. the DPA4
-    ``use_amp`` switch) from the portable record.
+    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
     ----------
diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py
index 85e24e25f5..f352f38238 100644
--- a/deepmd/pt_expt/model/get_model.py
+++ b/deepmd/pt_expt/model/get_model.py
@@ -65,10 +65,8 @@
 _WARNED_ONCE: set[str] = set()
 
 
-# Constructing the WRAPPED atomic classes directly is the lossless assembly
-# rule: building a raw dpmodel instance and converting it afterwards
-# round-trips through serialize()/deserialize() and drops runtime-only
-# configuration (e.g. the DPA4 `use_amp` switch) from the live children.
+# 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)
 InnerPotentialAtomicModel = auto_wrapped_class(InnerPotentialAtomicModelDP)
@@ -228,10 +226,8 @@ def _compose_bridging(
         rcut=descriptor.get_rcut(),
         sel=descriptor.get_sel(),
     )
-    # Constructor-args form (NOT a pre-built `atomic_model_=` instance): the
-    # CM then constructs the WRAPPED composition class directly and keeps the
-    # live children -- assigning a populated raw dpmodel composition would
-    # convert it through the lossy serialize()/deserialize() round-trip.
+    # constructor-args form: the CM builds the wrapped composition around
+    # the live children (see the auto_wrapped_class invariant)
     return LinearEnergyModel(
         models=[model.atomic_model, zbl_atomic],
         type_map=data["type_map"],
diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py
index 93b459afd0..0768bb1858 100644
--- a/deepmd/pt_expt/model/make_model.py
+++ b/deepmd/pt_expt/model/make_model.py
@@ -423,13 +423,8 @@ def make_model(
         The model.
 
     """
-    # Hand the dpmodel CM the WRAPPED atomic class so `self.atomic_model =
-    # T_AtomicModel(...)` constructs the pt_expt module directly with the
-    # live (already wrapped) descriptor/fitting. Passing the raw dpmodel
-    # class instead would build a raw atomic model and convert the instance
-    # through a serialize()/deserialize() round-trip, which drops
-    # runtime-only configuration (e.g. the DPA4 `use_amp` switch) that the
-    # portable record deliberately does not carry.
+    # 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
diff --git a/source/tests/common/dpmodel/test_dpa4_frame_mixers.py b/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
index 09200500d5..8958b204dc 100644
--- a/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
+++ b/source/tests/common/dpmodel/test_dpa4_frame_mixers.py
@@ -224,12 +224,10 @@ def test_torch_namespace(cls) -> None:
     ],
 )
 def test_empty_batch_passes_through(cls) -> None:
-    """``N == 0`` must yield an empty result, not a reshape error.
+    """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; the degree-batched contraction
-    must keep the broadcasted matmul's empty-batch behavior (explicit
-    channel widths -- ``-1`` cannot be inferred from zero elements).
+    or a distributed rank owns no nodes.
     """
     import torch
 
@@ -329,10 +327,13 @@ def test_focus_batched_lowering_matches_einsum_backward(kind) -> None:
     )
 
     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,
-        pt_mod.weight.index_select(0, pt_mod.degree_index).detach(),
+        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
@@ -341,3 +342,6 @@ def test_focus_batched_lowering_matches_einsum_backward(kind) -> None:
     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
+    )

From d5bf63dae18a305f7e420fb808c6df38bd5bf4b8 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Sun, 16 Aug 2026 15:23:52 +0800
Subject: [PATCH 16/16] refactor: drop the unused inner_potential_model
 injection point

Only linear_atomic_model is load-bearing. Dropping the pt_expt override of
the analytical child changes nothing: its whole state is in the portable
record, and the wrapped composition's ModuleList conversion already turns
it into a torch module. Verified both ways -- without linear_atomic_model
the use_amp assembly tests fail, without inner_potential_model they pass
and the child is still an nn.Module.

Also removes the now-dead pt_expt InnerPotentialAtomicModel alias, whose
only caller (_compose_bridging) was deleted by the merge.
---
 deepmd/dpmodel/model/model_factory.py | 21 ++++++---------------
 deepmd/pt_expt/model/get_model.py     |  5 -----
 2 files changed, 6 insertions(+), 20 deletions(-)

diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py
index 3aa9a76cab..d570e12eb5 100644
--- a/deepmd/dpmodel/model/model_factory.py
+++ b/deepmd/dpmodel/model/model_factory.py
@@ -135,7 +135,6 @@ def get_linear_atomic_model(
     backend_name: str,
     atomic_model: type,
     pairtab_model: type,
-    inner_potential_model: type | None = None,
     linear_atomic_model: type | None = None,
     descriptor_child_builder: "Callable[[dict], Any | None] | None" = None,
 ) -> Any:
@@ -164,17 +163,13 @@ def get_linear_atomic_model(
         Backend learned atomic-model class.
     pairtab_model : type
         Backend pair-tabulation atomic-model class.
-    inner_potential_model : type, optional
-        Backend analytical-bridging atomic-model class. Defaults to the
-        dpmodel class; a backend that wraps dpmodel classes must pass its
-        own wrapper so the composition is assembled from backend-native
-        children rather than converted afterwards.
     linear_atomic_model : type, optional
         Backend linear composition atomic-model class. Defaults to the
-        dpmodel class, with the same obligation as
-        ``inner_potential_model``: a wrapping backend that leaves this
-        unset gets a composition that its model wrapper must convert, and
-        conversion keeps only what the portable record carries.
+        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
@@ -193,13 +188,12 @@ def get_linear_atomic_model(
         unsupported kind.
     """
     from deepmd.dpmodel.atomic_model.inner_potential import (
-        InnerPotentialAtomicModel as InnerPotentialAtomicModelDP,
+        InnerPotentialAtomicModel,
     )
     from deepmd.dpmodel.atomic_model.linear_atomic_model import (
         LinearEnergyAtomicModel as LinearEnergyAtomicModelDP,
     )
 
-    InnerPotentialAtomicModel = inner_potential_model or InnerPotentialAtomicModelDP
     LinearEnergyAtomicModel = linear_atomic_model or LinearEnergyAtomicModelDP
 
     data = copy.deepcopy(data)
@@ -449,7 +443,6 @@ def __init__(
         atomic_model: type | None = None,
         pairtab_model: type | None = None,
         zbl_model: type | None = None,
-        inner_potential_model: type | None = None,
         linear_atomic_model: type | None = None,
     ) -> None:
         """Store backend-native classes used by all model construction paths."""
@@ -460,7 +453,6 @@ def __init__(
         self.atomic_model = atomic_model
         self.pairtab_model = pairtab_model
         self.zbl_model = zbl_model
-        self.inner_potential_model = inner_potential_model
         self.linear_atomic_model = linear_atomic_model
 
     def get_model_components(self, data: dict) -> tuple[Any, Any, str]:
@@ -498,7 +490,6 @@ def get_linear_atomic_model(
             backend_name=self.backend_name,
             atomic_model=self.atomic_model,
             pairtab_model=self.pairtab_model,
-            inner_potential_model=self.inner_potential_model,
             linear_atomic_model=self.linear_atomic_model,
             descriptor_child_builder=descriptor_child_builder,
         )
diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py
index 0af935e5f1..4b8230f2e7 100644
--- a/deepmd/pt_expt/model/get_model.py
+++ b/deepmd/pt_expt/model/get_model.py
@@ -10,9 +10,6 @@
 import logging
 
 from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP
-from deepmd.dpmodel.atomic_model.inner_potential import (
-    InnerPotentialAtomicModel as InnerPotentialAtomicModelDP,
-)
 from deepmd.dpmodel.atomic_model.linear_atomic_model import (
     LinearEnergyAtomicModel as LinearEnergyAtomicModelDP,
 )
@@ -67,7 +64,6 @@
 # runtime state (see the auto_wrapped_class invariant)
 DPAtomicModel = auto_wrapped_class(DPAtomicModelDP)
 PairTabAtomicModel = auto_wrapped_class(PairTabAtomicModelDP)
-InnerPotentialAtomicModel = auto_wrapped_class(InnerPotentialAtomicModelDP)
 LinearEnergyAtomicModel = auto_wrapped_class(LinearEnergyAtomicModelDP)
 
 _model_factory = BackendModelFactory(
@@ -78,7 +74,6 @@
     atomic_model=DPAtomicModel,
     pairtab_model=PairTabAtomicModel,
     zbl_model=DPZBLModel,
-    inner_potential_model=InnerPotentialAtomicModel,
     linear_atomic_model=LinearEnergyAtomicModel,
 )
 get_zbl_model = _model_factory.get_zbl_model