From 7518a417cbebf1367c262b082c2dcce818250d8b Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 5 Aug 2026 01:37:00 +0800 Subject: [PATCH 01/10] 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 Date: Wed, 5 Aug 2026 02:13:02 +0800 Subject: [PATCH 02/10] 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 Date: Wed, 5 Aug 2026 02:14:40 +0800 Subject: [PATCH 03/10] 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 Date: Wed, 5 Aug 2026 09:40:13 +0800 Subject: [PATCH 04/10] 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 Date: Wed, 5 Aug 2026 11:12:33 +0800 Subject: [PATCH 05/10] 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 Date: Wed, 5 Aug 2026 16:06:51 +0800 Subject: [PATCH 06/10] 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 Date: Wed, 5 Aug 2026 20:15:03 +0800 Subject: [PATCH 07/10] 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 Date: Thu, 6 Aug 2026 15:14:23 +0800 Subject: [PATCH 08/10] 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 Date: Thu, 6 Aug 2026 15:24:14 +0800 Subject: [PATCH 09/10] 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 Date: Thu, 6 Aug 2026 16:16:48 +0800 Subject: [PATCH 10/10] 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):