From dba50c5a43a26d9697b4e186ba353fa106d6817c Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 25 Jul 2026 13:58:09 +0800 Subject: [PATCH 1/2] fix(pt): honor is_sorted in SE-Attention tabulation Forward the public sorting flag through the PyTorch forward and first-backward helpers so unsorted excluded-type rows are evaluated consistently with grad-grad. Fixes #5890 Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/pt/tabulate_multi_device.cc | 34 ++--- .../tests/pt/test_tabulate_fusion_se_atten.py | 126 ++++++++++++++++++ 2 files changed, 145 insertions(+), 15 deletions(-) diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc index cede1d03d9..ad05d68144 100644 --- a/source/op/pt/tabulate_multi_device.cc +++ b/source/op/pt/tabulate_multi_device.cc @@ -25,6 +25,7 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor, const torch::Tensor& em_tensor, const torch::Tensor& two_embed_tensor, int64_t last_layer_size, + bool is_sorted, torch::Tensor& descriptor_tensor) { // check input shape if (table_tensor.dim() != 2) { @@ -60,7 +61,8 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor, if (device == "GPU") { #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM deepmd::tabulate_fusion_se_a_gpu(descriptor, table, table_info, em_x, em, - two_embed, nloc, nnei, last_layer_size); + two_embed, nloc, nnei, last_layer_size, + is_sorted); #else throw std::runtime_error( "The input tensor is on the GPU, but the GPU support for the " @@ -68,7 +70,8 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor, #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM } else if (device == "CPU") { deepmd::tabulate_fusion_se_a_cpu(descriptor, table, table_info, em_x, em, - two_embed, nloc, nnei, last_layer_size); + two_embed, nloc, nnei, last_layer_size, + is_sorted); } } @@ -80,6 +83,7 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor, const torch::Tensor& two_embed_tensor, const torch::Tensor& dy_tensor, const torch::Tensor& descriptor_tensor, + bool is_sorted, torch::Tensor& dy_dem_x_tensor, torch::Tensor& dy_dem_tensor, torch::Tensor& dy_dtwo_tensor) { @@ -111,18 +115,18 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor, // compute if (device == "GPU") { #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM - deepmd::tabulate_fusion_se_a_grad_gpu(dy_dem_x, dy_dem, dy_dtwo, table, - table_info, em_x, em, two_embed, dy, - nloc, nnei, last_layer_size); + deepmd::tabulate_fusion_se_a_grad_gpu( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); #else throw std::runtime_error( "The input tensor is on the GPU, but the GPU support for the " "customized OP library is not enabled."); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM } else if (device == "CPU") { - deepmd::tabulate_fusion_se_a_grad_cpu(dy_dem_x, dy_dem, dy_dtwo, table, - table_info, em_x, em, two_embed, dy, - nloc, nnei, last_layer_size); + deepmd::tabulate_fusion_se_a_grad_cpu( + dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy, + nloc, nnei, last_layer_size, is_sorted); } } @@ -643,7 +647,7 @@ class TabulateFusionSeAGradOp // compute TabulateFusionSeAGradForward( table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(), - dy_tensor, descriptor_tensor, dy_dem_x_tensor, dy_dem_tensor, + dy_tensor, descriptor_tensor, true, dy_dem_x_tensor, dy_dem_tensor, dy_dtwo_tensor); // save data ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor, @@ -785,7 +789,7 @@ class TabulateFusionSeAOp // compute TabulateFusionSeAForward(table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(), - last_layer_size, descriptor_tensor); + last_layer_size, true, descriptor_tensor); // save data ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor, em_tensor, descriptor_tensor}); @@ -870,8 +874,8 @@ class TabulateFusionSeAttenGradOp torch::Tensor dy_dtwo_tensor = torch::zeros_like(two_embed_tensor); TabulateFusionSeAGradForward( table_tensor, table_info_tensor, em_x_tensor, em_tensor, - two_embed_tensor, dy_tensor, descriptor_tensor, dy_dem_x_tensor, - dy_dem_tensor, dy_dtwo_tensor); + two_embed_tensor, dy_tensor, descriptor_tensor, is_sorted, + dy_dem_x_tensor, dy_dem_tensor, dy_dtwo_tensor); ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor, em_tensor, two_embed_tensor, descriptor_tensor}); @@ -969,9 +973,9 @@ class TabulateFusionSeAttenOp torch::Tensor descriptor_tensor = torch::empty({em_tensor.size(0), 4, last_layer_size}, options); // compute - TabulateFusionSeAForward(table_tensor, table_info_tensor, - em_x_tensor, em_tensor, two_embed_tensor, - last_layer_size, descriptor_tensor); + TabulateFusionSeAForward( + table_tensor, table_info_tensor, em_x_tensor, em_tensor, + two_embed_tensor, last_layer_size, is_sorted, descriptor_tensor); // save data ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor, em_tensor, two_embed_tensor, descriptor_tensor}); diff --git a/source/tests/pt/test_tabulate_fusion_se_atten.py b/source/tests/pt/test_tabulate_fusion_se_atten.py index f46d0b1761..e2243a4ffa 100644 --- a/source/tests/pt/test_tabulate_fusion_se_atten.py +++ b/source/tests/pt/test_tabulate_fusion_se_atten.py @@ -1645,5 +1645,131 @@ def test_second_order_backward(self) -> None: ) +@parameterized((torch.float64, torch.float32)) +@unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built") +class TestTabulateFusionSeAttenUnsortedOp(unittest.TestCase): + """Verify that an unsorted attention neighbor row is not folded as padding.""" + + def setUp(self) -> None: + (self.dtype,) = self.param + self.last_layer_size = 8 + coefficients = torch.arange( + 1, + self.last_layer_size + 1, + dtype=self.dtype, + device=env.DEVICE, + ) + table = torch.zeros( + (2, self.last_layer_size, 6), + dtype=self.dtype, + device=env.DEVICE, + ) + table[0, :, 0] = coefficients + table[1, :, 0] = 2.0 * coefficients + self.table_tensor = table.reshape(self.last_layer_size, -1) + self.table_info_tensor = torch.tensor( + [0.0, 1.0, 2.0, 1.0, 1.0, -1.0], + dtype=self.dtype, + device="cpu", + ) + # The first row looks like sorted padding because its scalar input + # matches the final row and its directional components are zero. The + # middle row is nevertheless a real neighbor and must still contribute + # when ``is_sorted`` is false. + self.em_x_tensor = torch.tensor( + [[0.5, 1.5, 0.5]], + dtype=self.dtype, + device=env.DEVICE, + requires_grad=True, + ) + self.em_tensor = torch.tensor( + [[[0.2, 0.0, 0.0, 0.0], [0.3, 0.1, 0.0, 0.0], [0.4, 0.0, 0.0, 0.0]]], + dtype=self.dtype, + device=env.DEVICE, + requires_grad=True, + ) + self.two_embed_tensor = torch.tensor( + [ + [0.0] * self.last_layer_size, + [0.5] * self.last_layer_size, + [-0.25] * self.last_layer_size, + ], + dtype=self.dtype, + device=env.DEVICE, + requires_grad=True, + ) + + def _forward(self, is_sorted: bool) -> torch.Tensor: + return torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor, + self.em_x_tensor, + self.em_tensor, + self.two_embed_tensor, + self.last_layer_size, + is_sorted, + )[0] + + def test_forward_uses_is_sorted(self) -> None: + coefficients = torch.arange( + 1, + self.last_layer_size + 1, + dtype=self.dtype, + device=env.DEVICE, + ) + expected_unsorted = torch.zeros( + (1, 4, self.last_layer_size), + dtype=self.dtype, + device=env.DEVICE, + ) + expected_unsorted[0, 0] = 1.4 * coefficients + expected_unsorted[0, 1] = 0.3 * coefficients + expected_sorted = torch.zeros_like(expected_unsorted) + expected_sorted[0, 0] = 0.6 * coefficients + + torch.testing.assert_close(self._forward(False), expected_unsorted) + torch.testing.assert_close(self._forward(True), expected_sorted) + + def test_backward_uses_is_sorted(self) -> None: + unsorted_grads = torch.autograd.grad( + self._forward(False).sum(), + (self.em_x_tensor, self.em_tensor, self.two_embed_tensor), + retain_graph=True, + ) + sorted_dem = torch.autograd.grad( + self._forward(True).sum(), + self.em_tensor, + )[0] + + coefficients = torch.arange( + 1, + self.last_layer_size + 1, + dtype=self.dtype, + device=env.DEVICE, + ) + expected_dem = torch.tensor( + [[[36.0] * 4, [108.0] * 4, [27.0] * 4]], + dtype=self.dtype, + device=env.DEVICE, + ) + expected_dtwo = torch.stack( + (0.2 * coefficients, 0.8 * coefficients, 0.4 * coefficients) + ) + + torch.testing.assert_close( + unsorted_grads[0], torch.zeros_like(unsorted_grads[0]) + ) + torch.testing.assert_close(unsorted_grads[1], expected_dem) + torch.testing.assert_close(unsorted_grads[2], expected_dtwo) + self.assertFalse(torch.equal(unsorted_grads[1], sorted_dem)) + + def test_second_order_backward_uses_is_sorted(self) -> None: + descriptor_tensor = self._forward(False) + assert_second_order_backward_matches_finite_difference( + descriptor_tensor, + (self.em_x_tensor, self.em_tensor, self.two_embed_tensor), + ) + + if __name__ == "__main__": unittest.main() From 8a3a02381e822536ae353a9ece7bb0375116c85d Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 00:42:54 +0800 Subject: [PATCH 2/2] fix(pt): sort compressed se_atten lower nlist Make geometrically compressed DPA1 request the sorted lower-neighbor-list path, and add operator, forward_lower, and end-to-end regression coverage for the sorting invariant. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/pt/model/descriptor/se_atten.py | 6 +- source/op/pt/tabulate_multi_device.cc | 6 + .../test_compressed_se_atten_forward_lower.py | 130 ++++++++++++++++++ .../pt/test_model_compression_se_atten.py | 7 +- .../tests/pt/test_tabulate_fusion_se_atten.py | 21 +-- 5 files changed, 159 insertions(+), 11 deletions(-) create mode 100644 source/tests/pt/model/test_compressed_se_atten_forward_lower.py diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 3f18b069d8..76c1db65e4 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -813,7 +813,11 @@ def has_message_passing(self) -> bool: def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" - return False + # Geometric compression uses the tabulate op's sorted-neighbor fold, + # which assumes padding and out-of-cutoff neighbors are trailing. + # `forward_lower` may receive an unsorted rcut+skin list from LAMMPS, + # so request the filtering/sorting pass whenever that op is active. + return self.geo_compress class NeighborGatedAttention(nn.Module): diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc index ad05d68144..355d1d89f6 100644 --- a/source/op/pt/tabulate_multi_device.cc +++ b/source/op/pt/tabulate_multi_device.cc @@ -645,6 +645,10 @@ class TabulateFusionSeAGradOp torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor); torch::Tensor dy_dtwo_tensor = at::Tensor(); // compute + // The non-attention se_a path invokes this op per type-pair block, so + // exclusions cannot interleave zero rows. Compressed forward_lower also + // requests a sorted nlist through the Python-side + // DescrptBlockSeA.need_sorted_nlist_for_lower() contract. TabulateFusionSeAGradForward( table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(), dy_tensor, descriptor_tensor, true, dy_dem_x_tensor, dy_dem_tensor, @@ -787,6 +791,8 @@ class TabulateFusionSeAOp torch::Tensor descriptor_tensor = torch::empty({em_tensor.size(0), 4, last_layer_size}, options); // compute + // Keep the sorted fold enabled: exclusions are uniform within each se_a + // type-pair invocation, and compressed forward_lower sorts its nlist first. TabulateFusionSeAForward(table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(), last_layer_size, true, descriptor_tensor); diff --git a/source/tests/pt/model/test_compressed_se_atten_forward_lower.py b/source/tests/pt/model/test_compressed_se_atten_forward_lower.py new file mode 100644 index 0000000000..cfa3d1807d --- /dev/null +++ b/source/tests/pt/model/test_compressed_se_atten_forward_lower.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression coverage for compressed DPA1 on an unsorted lower nlist. + +The geometric tabulation kernel folds trailing padding rows when its input is +sorted. LAMMPS supplies ``forward_lower`` with an rcut+skin neighbor list that +may put zero-switch, out-of-cutoff rows before real neighbors. A compressed +DPA1 model must therefore request the extra filtering/sorting pass before the +kernel applies that fold. +""" + +import copy +import unittest + +import torch + +from deepmd.pt.cxx_op import ( + ENABLE_CUSTOMIZED_OP, +) +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + +from ...seed import ( + GLOBAL_SEED, +) +from .test_forward_lower import ( + reduce_tensor, +) +from .test_permutation import ( + model_dpa1, +) + +dtype = torch.float64 + + +@unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built") +class TestCompressedSeAttenForwardLower(unittest.TestCase): + def setUp(self) -> None: + model_params = copy.deepcopy(model_dpa1) + # Geometric compression is available only for the strip representation + # without attention layers. + model_params["descriptor"]["tebd_input_mode"] = "strip" + model_params["descriptor"]["attn_layer"] = 0 + self.model = get_model(model_params).to(env.DEVICE) + + def _make_system(self): + """Create a sparse periodic system with neighbors in the skin region.""" + natoms = 6 + cell = 6.0 * torch.eye(3, dtype=dtype, device=env.DEVICE) + generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) + coord = 5.5 * torch.rand( + [natoms, 3], dtype=dtype, device=env.DEVICE, generator=generator + ) + atype = torch.tensor([0, 0, 1, 1, 2, 2], dtype=torch.int64, device=env.DEVICE) + return coord, atype, cell + + def _min_nbor_dist(self, coord, cell) -> float: + """Return the periodic minimum pair distance for table construction.""" + box = torch.diagonal(cell) + diff = coord[:, None, :] - coord[None, :, :] + diff = diff - torch.round(diff / box) * box + dist = torch.linalg.norm(diff, dim=-1) + dist = dist + torch.eye(coord.shape[0], device=coord.device) * 1e10 + return float(dist.min()) + + def test_unsorted_overcut_nlist(self) -> None: + coord, atype, cell = self._make_system() + rcut = self.model.get_rcut() + sel = self.model.get_sel() + + # Use a clean cutoff-bounded list as the uncompressed reference. + ec, ea, mapping, nlist = extend_input_and_build_neighbor_list( + coord.unsqueeze(0), + atype.unsqueeze(0), + rcut, + sel, + mixed_types=self.model.mixed_types(), + box=cell.unsqueeze(0), + ) + ref = self.model.forward_lower(ec, ea, nlist, mapping, do_atomic_virial=False) + + self.model.min_nbor_dist = torch.tensor( + 0.9 * self._min_nbor_dist(coord, cell), + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=env.DEVICE, + ) + self.model.enable_compression() + self.assertTrue(self.model.need_sorted_nlist_for_lower()) + + # Mimic the unsorted rcut+skin list from LAMMPS and deliberately move + # its out-of-cutoff/padding rows ahead of the in-cutoff neighbors. + ec2, ea2, mapping2, nlist2 = extend_input_and_build_neighbor_list( + coord.unsqueeze(0), + atype.unsqueeze(0), + rcut + 2.0, + sum(sel), + mixed_types=True, + box=cell.unsqueeze(0), + ) + safe_nlist = torch.clamp_min(nlist2, 0) + gather_index = safe_nlist.reshape(1, -1, 1).expand(-1, -1, 3) + neighbor_coord = torch.gather(ec2, 1, gather_index).view( + 1, coord.shape[0], -1, 3 + ) + center_coord = ec2[:, : coord.shape[0], :].unsqueeze(2) + distance = torch.linalg.norm(neighbor_coord - center_coord, dim=-1) + real_neighbor = nlist2 >= 0 + self.assertTrue(torch.any(real_neighbor & (distance <= rcut)).item()) + self.assertTrue(torch.any(real_neighbor & (distance > rcut)).item()) + + nlist2 = torch.flip(nlist2, dims=[-1]) + out = self.model.forward_lower( + ec2, ea2, nlist2, mapping2, do_atomic_virial=False + ) + + torch.testing.assert_close(out["energy"], ref["energy"], rtol=1e-10, atol=1e-10) + natoms = coord.shape[0] + ref_force = reduce_tensor(ref["extended_force"], mapping, natoms) + out_force = reduce_tensor(out["extended_force"], mapping2, natoms) + torch.testing.assert_close(out_force, ref_force, rtol=1e-10, atol=1e-10) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/test_model_compression_se_atten.py b/source/tests/pt/test_model_compression_se_atten.py index 41818f81d7..b3f9414b9e 100644 --- a/source/tests/pt/test_model_compression_se_atten.py +++ b/source/tests/pt/test_model_compression_se_atten.py @@ -84,10 +84,13 @@ def _init_models_exclude_types(): INPUT = str(tests_path / "input.json") jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - # Configure se_atten descriptor with exclude_types + # Plain se_atten defaults to a zero average, so excluded rows share the + # padding sentinel and exercise the unsorted path in the compressed op. jdata["model"]["descriptor"] = { - "type": "se_atten_v2", + "type": "se_atten", "exclude_types": [[0, 1]], + "set_davg_zero": True, + "tebd_input_mode": "strip", "sel": 120, "rcut_smth": 0.50, "rcut": 6.00, diff --git a/source/tests/pt/test_tabulate_fusion_se_atten.py b/source/tests/pt/test_tabulate_fusion_se_atten.py index e2243a4ffa..925d346880 100644 --- a/source/tests/pt/test_tabulate_fusion_se_atten.py +++ b/source/tests/pt/test_tabulate_fusion_se_atten.py @@ -1665,7 +1665,9 @@ def setUp(self) -> None: device=env.DEVICE, ) table[0, :, 0] = coefficients + table[0, :, 1] = coefficients table[1, :, 0] = 2.0 * coefficients + table[1, :, 1] = 2.0 * coefficients self.table_tensor = table.reshape(self.last_layer_size, -1) self.table_info_tensor = torch.tensor( [0.0, 1.0, 2.0, 1.0, 1.0, -1.0], @@ -1722,10 +1724,10 @@ def test_forward_uses_is_sorted(self) -> None: dtype=self.dtype, device=env.DEVICE, ) - expected_unsorted[0, 0] = 1.4 * coefficients - expected_unsorted[0, 1] = 0.3 * coefficients + expected_unsorted[0, 0] = 2.1 * coefficients + expected_unsorted[0, 1] = 0.45 * coefficients expected_sorted = torch.zeros_like(expected_unsorted) - expected_sorted[0, 0] = 0.6 * coefficients + expected_sorted[0, 0] = 0.9 * coefficients torch.testing.assert_close(self._forward(False), expected_unsorted) torch.testing.assert_close(self._forward(True), expected_sorted) @@ -1748,17 +1750,20 @@ def test_backward_uses_is_sorted(self) -> None: device=env.DEVICE, ) expected_dem = torch.tensor( - [[[36.0] * 4, [108.0] * 4, [27.0] * 4]], + [[[54.0] * 4, [162.0] * 4, [40.5] * 4]], dtype=self.dtype, device=env.DEVICE, ) expected_dtwo = torch.stack( - (0.2 * coefficients, 0.8 * coefficients, 0.4 * coefficients) + (0.3 * coefficients, 1.2 * coefficients, 0.6 * coefficients) ) - - torch.testing.assert_close( - unsorted_grads[0], torch.zeros_like(unsorted_grads[0]) + expected_dem_x = torch.tensor( + [[7.2, 43.2, 10.8]], + dtype=self.dtype, + device=env.DEVICE, ) + + torch.testing.assert_close(unsorted_grads[0], expected_dem_x) torch.testing.assert_close(unsorted_grads[1], expected_dem) torch.testing.assert_close(unsorted_grads[2], expected_dtwo) self.assertFalse(torch.equal(unsorted_grads[1], sorted_dem))