diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 801b4c34f2..0e13a99f9f 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -2748,6 +2748,11 @@ def serialize(self) -> dict[str, Any]: "mlp_bias": self.mlp_bias, "exclude_types": self.exclude_types, "eps": self.eps, + # Must round-trip: pt_expt rebuilds the descriptor from this + # dict, so omitting the key silently reset a configured + # ``use_amp: false`` to True and kept training in bfloat16. + # Older records without it still load (__init__ defaults it). + "use_amp": self.use_amp, "trainable": self.trainable, "seed": self.seed, "inner_clamp_r_inner": self.inner_clamp_r_inner, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py index 4c63d1a492..fcf7657e0b 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py @@ -99,6 +99,38 @@ def _build_frame_degree_index( raise ValueError("`coefficient_layout` must be either 'packed' or 'm_major'") +def _degree_batched_matmul(xp: Any, coeff: Any, weight: Any) -> Any: + """Contract ``einsum("ndfi,dio->ndfo")`` batched over the degree axis. + + Parameters + ---------- + xp : Any + The array namespace of ``coeff``. + coeff : Array + Coefficients with shape ``(N, D, F, i)``. + weight : Array + Per-degree weights with shape ``(D, i, o)``. + + Returns + ------- + Array + Contracted coefficients with shape ``(N, D, F, o)``. + + Notes + ----- + Batching over the degree axis, not over ``N``: the latter would broadcast + ``weight`` to ``(N, D, i, o)`` and make autograd reduce that expansion on + every backward. The transposes touch only ``coeff``, which is smaller. + """ + n_batch, coeff_dim, n_focus, _ = coeff.shape + coeff_d = xp.reshape( + xp.permute_dims(coeff, (1, 0, 2, 3)), (coeff_dim, n_batch * n_focus, -1) + ) # (D, N*F, i) + out = xp.matmul(coeff_d, weight) # (D, N*F, o) + out = xp.reshape(out, (coeff_dim, n_batch, n_focus, -1)) + return xp.permute_dims(out, (1, 0, 2, 3)) # (N, D, F, o) + + def _project_frames(coeff: Any, proj: ChannelLinear, n_frames: int) -> Any: """ Apply a channel-only linear map to each Wigner-D frame independently. @@ -493,9 +525,8 @@ def call(self, coeff: Any) -> Any: weight = xp_asarray_nodetach(xp, self.weight[...], device=device) degree_index = xp_asarray_nodetach(xp, self.degree_index, device=device) weight = xp.take(weight, degree_index, axis=0) - # einsum "ndfi,dio->ndfo" as a broadcast batched matmul: - # (N, D, F, i) @ (1, D, i, o) -> (N, D, F, o) - return xp.matmul(coeff, weight[None, ...]) + # Batched over the degree axis, never over N -- see the helper's note. + return _degree_batched_matmul(xp, coeff, weight) def serialize(self) -> dict[str, Any]: """Serialize the FrameContract to a dict.""" @@ -575,9 +606,8 @@ def call(self, coeff: Any) -> Any: weight = xp_asarray_nodetach(xp, self.weight[...], device=device) degree_index = xp_asarray_nodetach(xp, self.degree_index, device=device) weight = xp.take(weight, degree_index, axis=0) - # einsum "ndfi,dio->ndfo" as a broadcast batched matmul: - # (N, D, F, i) @ (1, D, i, o) -> (N, D, F, o) - return xp.matmul(coeff, weight[None, ...]) + # Batched over the degree axis, never over N -- see the helper's note. + return _degree_batched_matmul(xp, coeff, weight) def serialize(self) -> dict[str, Any]: """Serialize the FrameExpand to a dict.""" diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/lora.py b/deepmd/dpmodel/descriptor/dpa4_nn/lora.py index a4600b7dbe..ba8dd1ab1a 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/lora.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/lora.py @@ -189,10 +189,12 @@ def call(self, x: Array) -> Array: ) expand_index = xp_asarray_nodetach(xp, self.expand_index, device=device) weight_expanded = xp.take(weight, expand_index, axis=0) - # einsum "ndfi,difo->ndfo" as a broadcast batched matmul: - # (N, D, F, 1, Cin) @ (1, D, F, Cin, Cout) -> (N, D, F, 1, Cout) + # einsum "ndfi,difo->ndfo", batched over the small (D, F) axes rather + # than over N, which would broadcast the weight and make autograd + # reduce the expansion. LoRA twin of the so3.py contraction. weight_expanded = xp.permute_dims(weight_expanded, (0, 2, 1, 3)) - out = xp.matmul(x[:, :, :, None, :], weight_expanded[None, ...])[..., 0, :] + out = xp.matmul(xp.permute_dims(x, (1, 2, 0, 3)), weight_expanded) + out = xp.permute_dims(out, (2, 0, 1, 3)) # (N, D, F, Cout) if self.mlp_bias: bias = xp.reshape( xp_asarray_nodetach(xp, self.bias[...], device=device), diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py index 8ca2dbc855..b2cd2a18fe 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/so3.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so3.py @@ -131,10 +131,13 @@ def call(self, x: Any) -> Any: xp, self.weight[...], device=array_api_compat.device(x) ) weight = xp.reshape(weight, (self.in_channels, self.n_focus, self.out_channels)) - # einsum "bfi,ifo->bfo" as a broadcast batched matmul: - # (B, F, 1, Cin) @ (1, F, Cin, Cout) -> (B, F, 1, Cout) + # einsum "bfi,ifo->bfo" as F independent (B, Cin) x (Cin, Cout) GEMMs. + # B stays the GEMM rows so the weight is used in place; making B the + # batch axis would broadcast it to (B, F, Cin, Cout) and leave autograd + # reducing that expansion. At n_focus=1 both permutes are free views. weight = xp.permute_dims(weight, (1, 0, 2)) # (F, Cin, Cout) - out = xp.matmul(x[:, :, None, :], weight[None, ...])[..., 0, :] + out = xp.matmul(xp.permute_dims(x, (1, 0, 2)), weight) # (F, B, Cout) + out = xp.permute_dims(out, (1, 0, 2)) # (B, F, Cout) if self.use_bias: bias = xp_asarray_nodetach( xp, self.bias[...], device=array_api_compat.device(x) @@ -439,12 +442,16 @@ def call(self, x: Any) -> Any: weight_expanded = xp.take(weight, expand_index, axis=0) # (D, Cin, F, Cout) # === Step 2. Per-focus, per-degree channel mixing === - # einsum "ndfi,difo->ndfo" as a broadcast batched matmul: - # (N, D, F, 1, Cin) @ (1, D, F, Cin, Cout) -> (N, D, F, 1, Cout) + # einsum "ndfi,difo->ndfo", batched over the small (D, F) axes. + # Batching over the node axis N instead would broadcast the weight to + # (N, D, F, Cin, Cout) -- for the water example a 165K-element parameter + # blown up to 191M elements per call -- and autograd would then reduce + # that expansion back down. It was the costliest kernel of a step. weight_expanded = xp.permute_dims( weight_expanded, (0, 2, 1, 3) ) # (D, F, Cin, Cout) - out = xp.matmul(x[:, :, :, None, :], weight_expanded[None, ...])[..., 0, :] + out = xp.matmul(xp.permute_dims(x, (1, 2, 0, 3)), weight_expanded) + out = xp.permute_dims(out, (2, 0, 1, 3)) # (N, D, F, Cout) # === Step 3. Add l=0 bias === if self.mlp_bias: diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 084f75b14d..800da9453e 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -2561,6 +2561,9 @@ def serialize(self) -> dict[str, Any]: "mlp_bias": self.mlp_bias, "exclude_types": self.exclude_types, "eps": self.eps, + # Kept in step with the dpmodel serialize contract so both + # backends' records carry the same keys. + "use_amp": self.use_amp, "trainable": self.trainable, "seed": self.seed, "inner_clamp_r_inner": self.inner_clamp_r_inner, diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py index 93b94f76e7..c64eb9152d 100644 --- a/source/tests/common/dpmodel/test_descrpt_dpa4.py +++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py @@ -265,6 +265,27 @@ def test_supported_feature_roundtrip(self, overrides) -> None: out2 = np.asarray(dd2.call(coord.reshape(nf, -1), atype, nlist)[0]) np.testing.assert_array_equal(out1, out2) + @pytest.mark.parametrize( + "use_amp", + [ + True, # the constructor default; must not be clobbered either + False, # the value that was silently lost, re-enabling autocast + ], + ) + def test_use_amp_survives_roundtrip(self, use_amp) -> None: + """``use_amp`` must round-trip through serialize/deserialize. + + The key was missing from the config, so a backend that rebuilds from + it (pt_expt does) reset ``use_amp: false`` to True and kept training in + bfloat16. The forward-output round-trip test can't catch this -- + dpmodel never autocasts, so outputs match either way. + """ + dd = make_descriptor(use_amp=use_amp) + assert dd.use_amp is use_amp + assert dd.serialize()["config"]["use_amp"] is use_amp + dd2 = DescrptDPA4.deserialize(dd.serialize()) + assert dd2.use_amp is use_amp + def test_value_errors(self) -> None: with pytest.raises(ValueError): # kmax must be <= lmax make_descriptor(kmax=4, lmax=3)