From b0dd642122bcfb64e8922eec4b70402585579b01 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 16 Jun 2026 14:53:53 +0800 Subject: [PATCH 1/3] refactor(dpa4): output ffn --- deepmd/pt/model/descriptor/sezm.py | 55 ++++++++++++++++++++++++------ deepmd/utils/argcheck.py | 19 +++++++++++ examples/water/dpa4/input.json | 1 + 3 files changed, 65 insertions(+), 10 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index ba3ef38e65..9a56d7b1c4 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -320,6 +320,16 @@ class DescrptSeZM(BaseDescriptor, nn.Module): message_node_so3 If True, use the corresponding post-aggregation SO(3) Wigner-D grid-net branch. The message is the query and the node state is the context. + so3_readout + Read-out FFN mode for the final ``l=0`` descriptor. ``"none"`` applies a + degree-0 scalar FFN to the ``l=0`` slice only; ``l>0`` coefficients are + discarded before the read-out. ``"glu"`` and ``"mlp"`` apply a full + equivariant FFN whose degree equals the node degree of the last + interaction block, driven by the SO(3) Wigner-D grid, so ``l>0`` geometry + is folded into ``l=0`` before the scalar is extracted. The value selects + the quadratic grid product (``"glu"``) or the polynomial point-wise grid + MLP (``"mlp"``). The Wigner-D frame order follows ``kmax``. The residual + stays on the ``l=0`` channel. lebedev_quadrature Either one boolean applied to both S2 branches, or two booleans ``[so2_enabled, ffn_enabled]`` aligned with ``s2_activation``. If @@ -425,6 +435,7 @@ def __init__( node_wise_so3: bool = False, message_node_s2: bool = False, message_node_so3: bool = False, + so3_readout: str = "none", lebedev_quadrature: bool | list[bool] | None = True, activation_function: str = "silu", glu_activation: bool = True, @@ -512,6 +523,9 @@ def __init__( self.node_wise_so3 = bool(node_wise_so3) self.message_node_s2 = bool(message_node_s2) self.message_node_so3 = bool(message_node_so3) + self.so3_readout = str(so3_readout).lower() + if self.so3_readout not in {"none", "glu", "mlp"}: + raise ValueError("`so3_readout` must be one of 'none', 'glu', or 'mlp'") if lebedev_quadrature is None: lebedev_quadrature = [True, True] elif isinstance(lebedev_quadrature, bool): @@ -932,13 +946,21 @@ def __init__( ) # === Final FFN for l=0 output mixing === + # ``so3_readout="none"`` runs a degree-0 scalar FFN on the l=0 slice. + # ``"glu"``/``"mlp"`` run a full FFN at the last block's node degree whose + # SO(3) Wigner-D grid folds l>0 geometry into l=0; the value selects the + # quadratic grid product or the point-wise grid MLP. + readout_lmax = self.node_l_schedule[-1] self.output_ffn = EquivariantFFN( - lmax=0, + lmax=0 if self.so3_readout == "none" else readout_lmax, channels=self.channels, hidden_channels=self.out_ffn_neurons, - grid_mlp=False, + kmax=min(self.kmax, readout_lmax), + grid_mlp=self.so3_readout == "mlp", + grid_branch=0, dtype=self.compute_dtype, s2_activation=False, + ffn_so3_grid=self.so3_readout != "none", activation_function=self.out_activation_function, glu_activation=self.out_glu_activation, mlp_bias=self.mlp_bias, @@ -1205,15 +1227,20 @@ def forward( x = self._forward_blocks(x, edge_cache, rad_feat_per_block) # === Step 11. Final l=0 output mixing === - # Extract l=0 scalar features and apply FFN in promoted dtype. - # Residual keeps the output close to identity with zero-initialized FFN output. + # ``none`` feeds the l=0 slice only; ``glu``/``mlp`` feed the full + # (N, D, 1, C) node tensor so the SO(3) grid folds l>0 into l=0. The + # residual is added on the full coefficient tensor before extracting + # l=0: slicing the summed tensor rather than the FFN output keeps the + # saved degree-axis stride static under torch.compile dynamic shapes. with nvtx_range("output_ffn"): - x_scalar = ( + ffn_in = ( x[:, 0:1, :, :] .reshape(n_nodes, 1, 1, self.channels) .to(dtype=self.compute_dtype) - ) # (N, 1, 1, C) - x_scalar = x_scalar + self.output_ffn(x_scalar) + if self.so3_readout == "none" + else x.to(dtype=self.compute_dtype) + ) + x_scalar = (ffn_in + self.output_ffn(ffn_in))[:, 0:1, :, :] # === Step 12. Reshape to (nf, nloc, channels) and return === descriptor = rearrange( @@ -1380,13 +1407,20 @@ def forward_with_edges( x = self._forward_blocks(x, edge_cache, rad_feat_per_block) # === Step 10. Final l=0 output mixing === + # ``none`` feeds the l=0 slice only; ``glu``/``mlp`` feed the full + # (N, D, 1, C) node tensor so the SO(3) grid folds l>0 into l=0. The + # residual is added on the full coefficient tensor before extracting + # l=0: slicing the summed tensor rather than the FFN output keeps the + # saved degree-axis stride static under torch.compile dynamic shapes. with nvtx_range("output_ffn"): - x_scalar = ( + ffn_in = ( x[:, 0:1, :, :] .reshape(n_nodes, 1, 1, self.channels) .to(dtype=self.compute_dtype) - ) # (N, 1, 1, C) - x_scalar = x_scalar + self.output_ffn(x_scalar) + if self.so3_readout == "none" + else x.to(dtype=self.compute_dtype) + ) + x_scalar = (ffn_in + self.output_ffn(ffn_in))[:, 0:1, :, :] # === Step 11. Reshape to (nf, nloc, channels) and return === descriptor = x_scalar.reshape(nf, nloc, self.channels) # (nf, nloc, C) @@ -2043,6 +2077,7 @@ def serialize(self) -> dict[str, Any]: "node_wise_so3": self.node_wise_so3, "message_node_s2": self.message_node_s2, "message_node_so3": self.message_node_so3, + "so3_readout": self.so3_readout, "lebedev_quadrature": self.lebedev_quadrature, "activation_function": self.activation_function, "glu_activation": self.glu_activation, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index bb763098d5..48c711a10b 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -551,6 +551,16 @@ def descrpt_se_zm_args() -> list[Argument]: "context. When enabled together with `message_node_s2`, the SO(3) " "branch is used for this path." ) + doc_so3_readout = ( + "Read-out FFN mode for the final l=0 descriptor. `none` applies a " + "degree-0 scalar FFN to the l=0 slice only; l>0 coefficients are " + "discarded before the read-out. `glu` and `mlp` apply a full equivariant " + "FFN on the SO(3) Wigner-D grid so l>0 geometry is folded into l=0 " + "before the scalar is extracted; the value selects the quadratic grid " + "product (`glu`) or the polynomial point-wise grid MLP (`mlp`). The " + "read-out degree equals the node degree of the last interaction block; " + "the Wigner-D frame order follows `kmax`." + ) doc_lebedev_quadrature = ( "Either one boolean applied to both S2 branches, or two booleans " "`[so2_enabled, ffn_enabled]` aligned with `s2_activation`. If a branch " @@ -881,6 +891,15 @@ def descrpt_se_zm_args() -> list[Argument]: default=False, doc=doc_only_pt_supported + doc_message_node_so3, ), + Argument( + "so3_readout", + str, + optional=True, + default="none", + extra_check=lambda x: x in ("none", "glu", "mlp"), + extra_check_errmsg="must be one of 'none', 'glu', or 'mlp'", + doc=doc_only_pt_supported + doc_so3_readout, + ), Argument( "lebedev_quadrature", [bool, list[bool]], diff --git a/examples/water/dpa4/input.json b/examples/water/dpa4/input.json index 34e316086a..415f6a5be0 100644 --- a/examples/water/dpa4/input.json +++ b/examples/water/dpa4/input.json @@ -23,6 +23,7 @@ "n_atten_head": 1, "ffn_neurons": 0, "ffn_so3_grid": true, + "so3_readout": "mlp", "grid_mlp": [ false, false, From bdefabe7a2be7d63400d54a11e4a5ff9bc691e54 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 20 Jun 2026 09:07:54 +0800 Subject: [PATCH 2/3] feat(dpmodel): port so3_readout to DescrptDPA4 (cross-backend with pt #5556) --- deepmd/dpmodel/descriptor/dpa4.py | 34 +++++++++--- .../tests/consistent/descriptor/test_dpa4.py | 8 +++ .../pt/model/test_dpa4_dpmodel_parity.py | 54 +++++++++++++++++++ 3 files changed, 90 insertions(+), 6 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 5c61dbc388..6a7f0a9545 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -194,6 +194,7 @@ def __init__( node_wise_so3: bool = False, message_node_s2: bool = False, message_node_so3: bool = False, + so3_readout: str = "none", lebedev_quadrature: bool | list[bool] | None = True, activation_function: str = "silu", glu_activation: bool = True, @@ -275,6 +276,9 @@ def __init__( self.node_wise_so3 = bool(node_wise_so3) self.message_node_s2 = bool(message_node_s2) self.message_node_so3 = bool(message_node_so3) + self.so3_readout = str(so3_readout).lower() + if self.so3_readout not in {"none", "glu", "mlp"}: + raise ValueError("`so3_readout` must be one of 'none', 'glu', or 'mlp'") if lebedev_quadrature is None: lebedev_quadrature = [True, True] elif isinstance(lebedev_quadrature, bool): @@ -639,12 +643,20 @@ def __init__( self.blocks = blocks # === Final FFN for l=0 output mixing (fp32+) === + # ``so3_readout="none"`` runs a degree-0 scalar FFN on the l=0 slice. + # ``"glu"``/``"mlp"`` run a full FFN at the last block's node degree whose + # SO(3) Wigner-D grid folds l>0 geometry into l=0; the value selects the + # quadratic grid product or the point-wise grid MLP. + readout_lmax = self.node_l_schedule[-1] self.output_ffn = EquivariantFFN( - lmax=0, + lmax=0 if self.so3_readout == "none" else readout_lmax, channels=self.channels, hidden_channels=self.out_ffn_neurons, - grid_mlp=False, + kmax=min(self.kmax, readout_lmax), + grid_mlp=self.so3_readout == "mlp", + grid_branch=0, s2_activation=False, + ffn_so3_grid=self.so3_readout != "none", activation_function=self.out_activation_function, glu_activation=self.out_glu_activation, mlp_bias=self.mlp_bias, @@ -926,10 +938,19 @@ def call( x = block(x, edge_cache, rad_feat_per_block[i])[0] # === Step 10. Final l=0 output mixing === - x_scalar = xp.reshape( - x[:, 0:1, :, :], (n_nodes, 1, 1, self.channels) - ) # (N, 1, 1, C) - x_scalar = x_scalar + self.output_ffn(x_scalar) + # ``none`` feeds the l=0 slice only; ``glu``/``mlp`` feed the full + # (N, D, 1, C) node tensor so the SO(3) grid folds l>0 into l=0. The + # residual is added on the full coefficient tensor before extracting + # l=0 to mirror pt. + compute_prec = get_xp_precision(xp, self.compute_precision) + if self.so3_readout == "none": + ffn_in = xp.astype( + xp.reshape(x[:, 0:1, :, :], (n_nodes, 1, 1, self.channels)), + compute_prec, + ) # (N, 1, 1, C) + else: + ffn_in = xp.astype(x, compute_prec) # (N, D, 1, C) + x_scalar = (ffn_in + self.output_ffn(ffn_in))[:, 0:1, :, :] # === Step 11. Reshape and return === descriptor = xp.reshape(x_scalar, (nf, nloc, self.channels)) @@ -1231,6 +1252,7 @@ def serialize(self) -> dict[str, Any]: "node_wise_so3": self.node_wise_so3, "message_node_s2": self.message_node_s2, "message_node_so3": self.message_node_so3, + "so3_readout": self.so3_readout, "lebedev_quadrature": self.lebedev_quadrature, "activation_function": self.activation_function, "glu_activation": self.glu_activation, diff --git a/source/tests/consistent/descriptor/test_dpa4.py b/source/tests/consistent/descriptor/test_dpa4.py index 18b8ac1ee5..e6f3216bd4 100644 --- a/source/tests/consistent/descriptor/test_dpa4.py +++ b/source/tests/consistent/descriptor/test_dpa4.py @@ -49,6 +49,7 @@ "ffn_so3_grid", "message_node_so3", "grid_mlp", + "so3_readout", ) DPA4_BASELINE_CASE = { @@ -59,6 +60,7 @@ "ffn_so3_grid": False, "message_node_so3": False, "grid_mlp": False, + "so3_readout": "none", } @@ -99,6 +101,10 @@ def dpa4_case(**overrides: Any) -> tuple: dpa4_case(ffn_so3_grid=True, message_node_so3=True), # polynomial grid MLP op (grid_branch=0 so grid_mlp takes effect) dpa4_case(grid_mlp=True, grid_branch=[0, 0, 0]), + # SO(3) readout: GLU grid product folds l>0 into the l=0 output + dpa4_case(so3_readout="glu"), + # SO(3) readout: point-wise grid MLP folds l>0 into the l=0 output + dpa4_case(so3_readout="mlp"), ) @@ -114,6 +120,7 @@ def data(self) -> dict: ffn_so3_grid, message_node_so3, grid_mlp, + so3_readout, ) = self.param return { "ntypes": self.ntypes, @@ -130,6 +137,7 @@ def data(self) -> dict: "ffn_so3_grid": ffn_so3_grid, "message_node_so3": message_node_so3, "grid_mlp": grid_mlp, + "so3_readout": so3_readout, "random_gamma": False, "precision": precision, "trainable": False, diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py index 751c6a9001..8009f81bc0 100644 --- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py +++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py @@ -3578,6 +3578,60 @@ def test_descriptor_extra_node_l(self) -> None: pt_mod, dp_mod, _ = self._build_descr_pair(extra_node_l=1) self._assert_descr_parity(pt_mod, dp_mod) + @pytest.mark.parametrize( + "so3_readout", ["glu", "mlp"] + ) # SO(3) grid readout: quadratic grid product vs point-wise grid MLP + def test_descriptor_so3_readout(self, so3_readout) -> None: + # so3_readout!="none" feeds the full (N, D, 1, C) node tensor to the + # output FFN so the SO(3) Wigner-D grid folds l>0 into l=0. The pt + # reference is pinned to CPU so the parity holds under the CUDA default + # device; the gate stays at the strict fp64 descriptor tolerance. + from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, + ) + from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, + ) + + kwargs = self._descr_kwargs(so3_readout=so3_readout) + pt_mod = DescrptSeZM(**kwargs).double().eval().to("cpu") + # so3_linear_2 / output projections are zero-initialized; perturb so the + # readout output is nontrivial (otherwise it is identically ~0) + rng = np.random.default_rng(2160) + with torch.no_grad(): + for p in pt_mod.parameters(): + p += torch.from_numpy(0.05 * rng.normal(size=tuple(p.shape))).to("cpu") + dp_mod = DescrptDPA4.deserialize(pt_mod.serialize()) + assert dp_mod.so3_readout == so3_readout + + inp = self._inputs() + coord, atype_ext, nlist, mp = ( + inp["coord"], + inp["atype_ext"], + inp["nlist"], + inp["mapping"], + ) + nf = coord.shape[0] + out_dp = np.asarray( + dp_mod.call(coord.reshape(nf, -1), atype_ext, nlist, mapping=mp)[0] + ) + out_pt = ( + pt_mod( + torch.from_numpy(coord).to("cpu"), + torch.from_numpy(atype_ext.astype(np.int64)).to("cpu"), + torch.from_numpy(nlist.astype(np.int64)).to("cpu"), + mapping=torch.from_numpy(mp.astype(np.int64)).to("cpu"), + )[0] + .detach() + .cpu() + .numpy() + ) + assert out_dp.shape == out_pt.shape + # nontrivial output magnitude (guards against a trivially-zero readout) + assert np.abs(out_dp).max() > 1e-6 + # strict fp64 descriptor-level gate + np.testing.assert_allclose(out_dp, out_pt, rtol=1e-10, atol=1e-12) + def test_descriptor_torch_namespace(self) -> None: # the dp descriptor must run under the torch array namespace as well: # feeding torch tensors must yield a torch tensor matching the numpy From 01ba3764376cbf671baddb84b06890bdbc0977c0 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 20 Jun 2026 09:37:23 +0800 Subject: [PATCH 3/3] fix(dpa4): truncate so3_readout input to final node degree (empty-edge path) AI-review (CodeRabbit/codex) finding on #5561: with so3_readout=glu/mlp and a shrinking l_schedule, the empty-edge path skips _forward_blocks, leaving x at the initial node degree (node_ebed_dims[0]); the full x was fed to output_ffn built for node_ebed_dims[-1] -> SO3Linear einsum shape mismatch on isolated atoms. Truncate the readout input to node_ebed_dims[-1] (no-op once blocks ran). - pt sezm.py: slice x to node_ebed_dims[-1] in both readout sites (forward, forward_with_edges). - dpmodel dpa4.py: same truncation for symmetry/robustness (no-op there since padded-edge blocks always shrink x). - regression test: so3_readout glu/mlp + shrinking schedule + all-isolated nlist (proven to fail pre-fix with the einsum size 4 vs 9 mismatch). --- deepmd/dpmodel/descriptor/dpa4.py | 7 +++++- deepmd/pt/model/descriptor/sezm.py | 10 ++++++-- source/tests/pt/model/test_descriptor_sezm.py | 24 +++++++++++++++++++ 3 files changed, 38 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 6a7f0a9545..de96ee8c0c 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -949,7 +949,12 @@ def call( compute_prec, ) # (N, 1, 1, C) else: - ffn_in = xp.astype(x, compute_prec) # (N, D, 1, C) + # truncate to the final node degree (what output_ffn is built for); + # no-op in the normal path (blocks already shrank x), defensive vs + # any path that leaves x at the initial degree. Mirrors pt. + ffn_in = xp.astype( + x[:, : self.node_ebed_dims[-1], :, :], compute_prec + ) # (N, D, 1, C) x_scalar = (ffn_in + self.output_ffn(ffn_in))[:, 0:1, :, :] # === Step 11. Reshape and return === diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 9a56d7b1c4..165fef0aaa 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -1238,7 +1238,10 @@ def forward( .reshape(n_nodes, 1, 1, self.channels) .to(dtype=self.compute_dtype) if self.so3_readout == "none" - else x.to(dtype=self.compute_dtype) + # truncate to the final node degree: the empty-edge path + # skips the blocks, leaving x at node_ebed_dims[0]; output_ffn + # is built for node_ebed_dims[-1]. No-op when blocks ran. + else x[:, : self.node_ebed_dims[-1], :, :].to(dtype=self.compute_dtype) ) x_scalar = (ffn_in + self.output_ffn(ffn_in))[:, 0:1, :, :] @@ -1418,7 +1421,10 @@ def forward_with_edges( .reshape(n_nodes, 1, 1, self.channels) .to(dtype=self.compute_dtype) if self.so3_readout == "none" - else x.to(dtype=self.compute_dtype) + # truncate to the final node degree: the empty-edge path + # skips the blocks, leaving x at node_ebed_dims[0]; output_ffn + # is built for node_ebed_dims[-1]. No-op when blocks ran. + else x[:, : self.node_ebed_dims[-1], :, :].to(dtype=self.compute_dtype) ) x_scalar = (ffn_in + self.output_ffn(ffn_in))[:, 0:1, :, :] diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py index 2be18bfcca..ec5a549903 100644 --- a/source/tests/pt/model/test_descriptor_sezm.py +++ b/source/tests/pt/model/test_descriptor_sezm.py @@ -155,6 +155,30 @@ def _assert_forward_backward_smoke(self, **model_kwargs) -> DescrptSeZM: self.assertTrue(torch.all(torch.isfinite(extended_coord.grad))) return model + def test_so3_readout_empty_edge_shrinking_schedule(self) -> None: + """so3_readout glu/mlp must handle the empty-edge path. + + With a shrinking ``l_schedule`` and no edges (every atom isolated), + ``_forward_blocks`` is skipped so ``x`` keeps the *initial* node degree + ``node_ebed_dims[0]``; the readout must truncate it to the final degree + ``node_ebed_dims[-1]`` (what ``output_ffn`` is built for) before the FFN. + Regression for the readout shape mismatch on isolated atoms. + """ + coord, atype, _ = _tiny_two_atom_system(self.device, dtype=torch.float32) + extended_coord = coord.reshape(1, -1).detach().requires_grad_(True) + # all neighbors masked out -> edge_cache.src.numel() == 0 -> blocks skipped + nlist = torch.full((1, 2, 2), -1, dtype=torch.int64, device=self.device) + for readout in ("glu", "mlp"): + with self.subTest(so3_readout=readout): + model = DescrptSeZM( + **_descriptor_kwargs(l_schedule=[2, 1], so3_readout=readout) + ) + desc, *_ = model( + extended_coord, atype, nlist, mapping=None, comm_dict=None + ) + self.assertEqual(desc.shape, (1, 2, 4)) + self.assertTrue(torch.all(torch.isfinite(desc))) + def test_forward_with_descriptor_variants(self) -> None: """Test forward/backward smoke paths for compact descriptor variants.""" cases = {