Skip to content

Commit beb50da

Browse files
wanghan-iapcmHan Wang
wanghan-iapcm
and
Han Wang
authored
fix(pt): sort nlist for compressed se_e2_a in forward_lower (deepmodeling#5524)
## Summary Fixes wrong energy/forces (and unstable LAMMPS MD) for **compressed `se_e2_a` models evaluated through `forward_lower`** — the C++/LAMMPS inference path. Reported in discussion deepmodeling#5438. ## Root cause The compressed `tabulate_fusion_se_a` op (`source/lib/src/tabulate.cc`, forward and grad kernels) has an `is_sorted`-gated early-termination: ```cpp ago = em_x[ii * nnei + nnei - 1]; // last neighbor's em_x if (ago == xx && ll[1]==0 && ll[2]==0 && ll[3]==0 && is_sorted) break; ``` It stops accumulating at the first neighbor whose env-mat **direction is zero**. Both `-1` padding *and* out-of-`rcut` neighbors (`sw==0`) have zero direction and the same `em_x == -davg/dstd` (`== ago`), so the op assumes all such neighbors are **trailing**. `is_sorted` defaults to `true` and the PT op never overrides it. The C++/LAMMPS `forward_lower` neighbor list uses `rcut + skin` and is **not distance-sorted**. `_format_nlist` only filters out-of-`rcut` neighbors in its *sort* branch (`n_nnei > nnei`), which is skipped when the LAMMPS list is narrower than `sum(sel)` (pad-only branch). The zero-direction neighbors then land **before** real ones, so the op breaks early and silently drops real neighbors → wrong descriptor → wrong energy/forces → unstable MD. Only the **compressed** path is affected: - The uncompressed embedding-net path sums over neighbors and treats zero-direction fillers identically regardless of position, so it is order-invariant. - Only `tabulate_fusion_se_a` (forward + grad) has this early-termination; `se_t`/`se_r` forward kernels do not. It is **device-independent** (reproduces identically on CPU and GPU). ## Fix The wiring already exists — the model calls `format_nlist(..., extra_nlist_sort=self.need_sorted_nlist_for_lower())` — but `DescrptBlockSeA.need_sorted_nlist_for_lower()` always returned `False`. Make it return `self.compress`: ```python def need_sorted_nlist_for_lower(self) -> bool: return self.compress ``` When compression is enabled this forces the sort + `rcut`-filter branch (in-`rcut` neighbors first, all padding last), restoring the op's invariant. The standard (uncompressed) route is unchanged, so there is no added cost on the common path. ## Verification - In LAMMPS (CPU and GPU) the compressed model now matches the uncompressed model to ~2.5e-14 (was ~0.5 eV off with scrambled forces). - New regression test `source/tests/pt/model/test_compressed_se_a_forward_lower.py` runs compressed `forward_lower` with an unsorted, over-`rcut` neighbor list and compares energy + force to the uncompressed reference, parameterized over `type_one_side ∈ {True, False}`. It **fails without this fix and passes with it**; existing `test_compressed_descriptor_se_a.py` and `test_forward_lower.py` still pass. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Fixed compressed descriptor behavior so enabling compression preserves neighbor ordering invariants and no longer causes valid neighbors to be dropped; energies and forces remain correct with unsorted/padded neighbor lists. * **Tests** * Added regression tests for multiple descriptor variants that validate compressed mode against uncompressed baselines using unsorted/over-cut neighbor lists, asserting energy and force fidelity. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
1 parent 87d8557 commit beb50da

3 files changed

Lines changed: 296 additions & 1 deletion

File tree

‎deepmd/pt/model/descriptor/se_a.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -871,4 +871,13 @@ def has_message_passing(self) -> bool:
871871

872872
def need_sorted_nlist_for_lower(self) -> bool:
873873
"""Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
874-
return False
874+
# The compressed tabulate op (`tabulate_fusion_se_a`) uses an
875+
# `is_sorted` early-termination that stops accumulating as soon as it
876+
# meets a neighbor whose env-mat direction is zero (padding, or an
877+
# out-of-rcut neighbor with sw==0). It therefore assumes such neighbors
878+
# are trailing. The `forward_lower` neighbor list coming from C++/LAMMPS
879+
# (rcut+skin, not pre-sorted) can interleave zero-direction neighbors
880+
# before real ones, which would silently drop real neighbors. Forcing a
881+
# sorted nlist (extra_nlist_sort) filters out-of-rcut neighbors and puts
882+
# all padding last, restoring the op's invariant. See discussion #5438.
883+
return self.compress

‎source/tests/pt/model/test_compressed_se_a_forward_lower.py‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
# SPDX-License-Identifier: LGPL-3.0-or-later
2+
"""Regression test for discussion #5438.
3+
4+
A compressed ``se_e2_a`` model produced wrong energy/forces when evaluated
5+
through ``forward_lower`` with a neighbor list that is *not* pre-sorted and
6+
contains out-of-``rcut`` (``sw == 0``) neighbors before the real ones -- exactly
7+
what the C++/LAMMPS inference path provides (its neighbor list uses
8+
``rcut + skin`` and is not distance-sorted).
9+
10+
The compressed ``tabulate_fusion_se_a`` op uses an ``is_sorted``
11+
early-termination that stops accumulating at the first neighbor whose env-mat
12+
direction is zero (padding, or an out-of-``rcut`` neighbor with ``sw == 0``),
13+
assuming such neighbors are trailing. When they appear before real neighbors the
14+
op silently drops the real neighbors, giving wrong descriptors.
15+
16+
The fix makes ``DescrptBlockSeA.need_sorted_nlist_for_lower()`` return
17+
``self.compress`` so the model forces an ``extra_nlist_sort`` (which filters
18+
out-of-``rcut`` neighbors and moves all padding last) before the op runs.
19+
20+
Without the fix, ``test_unsorted_overcut_nlist`` fails (energy/force mismatch);
21+
with the fix the compressed result matches the uncompressed reference.
22+
"""
23+
24+
import copy
25+
import unittest
26+
27+
import torch
28+
29+
from deepmd.pt.cxx_op import (
30+
ENABLE_CUSTOMIZED_OP,
31+
)
32+
from deepmd.pt.model.model import (
33+
get_model,
34+
)
35+
from deepmd.pt.utils import (
36+
env,
37+
)
38+
from deepmd.pt.utils.nlist import (
39+
extend_input_and_build_neighbor_list,
40+
)
41+
42+
from ...consistent.common import (
43+
parameterized,
44+
)
45+
from ...seed import (
46+
GLOBAL_SEED,
47+
)
48+
from .test_forward_lower import (
49+
reduce_tensor,
50+
)
51+
from .test_permutation import (
52+
model_se_e2_a,
53+
)
54+
55+
dtype = torch.float64
56+
57+
58+
@parameterized((True, False)) # type_one_side
59+
@unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built")
60+
class TestCompressedSeAForwardLower(unittest.TestCase):
61+
def setUp(self) -> None:
62+
(self.type_one_side,) = self.param
63+
model_params = copy.deepcopy(model_se_e2_a)
64+
model_params["descriptor"]["type_one_side"] = self.type_one_side
65+
self.model = get_model(model_params).to(env.DEVICE)
66+
67+
def _make_system(self):
68+
# a sparse system: atoms have fewer neighbors than ``sel`` within
69+
# ``rcut``, and some neighbors fall in ``(rcut, rcut + buffer]`` so the
70+
# over-cut neighbor list contains zero-``sw`` real neighbors.
71+
natoms = 6
72+
cell = 6.0 * torch.eye(3, dtype=dtype, device=env.DEVICE)
73+
generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED)
74+
coord = 5.5 * torch.rand(
75+
[natoms, 3], dtype=dtype, device=env.DEVICE, generator=generator
76+
)
77+
atype = torch.tensor([0, 0, 1, 1, 2, 2], dtype=torch.int64, device=env.DEVICE)
78+
return coord, atype, cell
79+
80+
def _min_nbor_dist(self, coord, cell):
81+
# minimum image minimum pair distance, used as the compression lower bound
82+
box = torch.diagonal(cell)
83+
diff = coord[:, None, :] - coord[None, :, :]
84+
diff = diff - torch.round(diff / box) * box
85+
dist = torch.linalg.norm(diff, dim=-1)
86+
dist = dist + torch.eye(coord.shape[0], device=coord.device) * 1e10
87+
return float(dist.min())
88+
89+
def test_unsorted_overcut_nlist(self) -> None:
90+
coord, atype, cell = self._make_system()
91+
rcut = self.model.get_rcut()
92+
sel = self.model.get_sel()
93+
94+
# reference: uncompressed forward_lower with a clean rcut-bounded nlist
95+
ec, ea, mp, nlist = extend_input_and_build_neighbor_list(
96+
coord.unsqueeze(0),
97+
atype.unsqueeze(0),
98+
rcut,
99+
sel,
100+
mixed_types=self.model.mixed_types(),
101+
box=cell.unsqueeze(0),
102+
)
103+
ref = self.model.forward_lower(ec, ea, nlist, mp, do_atomic_virial=False)
104+
105+
# enable compression (lower bound below the true min distance -> no extrapolation)
106+
self.model.min_nbor_dist = torch.tensor(
107+
0.9 * self._min_nbor_dist(coord, cell),
108+
dtype=env.GLOBAL_PT_FLOAT_PRECISION,
109+
device=env.DEVICE,
110+
)
111+
self.model.enable_compression()
112+
113+
# over-rcut FLAT neighbor list (mimics LAMMPS rcut+skin), reversed so the
114+
# out-of-rcut / padding neighbors precede the real ones.
115+
ec2, ea2, mp2, nlist2 = extend_input_and_build_neighbor_list(
116+
coord.unsqueeze(0),
117+
atype.unsqueeze(0),
118+
rcut + 2.0,
119+
sum(sel),
120+
mixed_types=True,
121+
box=cell.unsqueeze(0),
122+
)
123+
nlist2 = torch.flip(nlist2, dims=[-1])
124+
out = self.model.forward_lower(ec2, ea2, nlist2, mp2, do_atomic_virial=False)
125+
126+
torch.testing.assert_close(out["energy"], ref["energy"], rtol=1e-10, atol=1e-10)
127+
natoms = coord.shape[0]
128+
f_ref = reduce_tensor(ref["extended_force"], mp, natoms)
129+
f_out = reduce_tensor(out["extended_force"], mp2, natoms)
130+
torch.testing.assert_close(f_out, f_ref, rtol=1e-10, atol=1e-10)
131+
132+
133+
if __name__ == "__main__":
134+
unittest.main()

‎source/tests/pt/model/test_compressed_se_r_forward_lower.py‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
# SPDX-License-Identifier: LGPL-3.0-or-later
2+
"""Companion to ``test_compressed_se_a_forward_lower.py`` for discussion #5438.
3+
4+
The compressed ``se_e2_a`` descriptor produced wrong energy/forces when its
5+
``forward_lower`` neighbor list was *not* pre-sorted and contained
6+
out-of-``rcut`` (``sw == 0``) neighbors before the real ones -- the root cause
7+
being the ``is_sorted`` early-termination in the ``tabulate_fusion_se_a`` op
8+
(``source/lib/src/tabulate.cc``), which stops accumulating at the first
9+
zero-direction neighbor.
10+
11+
``se_e2_r`` is expected to be *immune*: its ``tabulate_fusion_se_r`` kernel has
12+
no such early-termination (it takes no ``is_sorted`` argument and always
13+
iterates every neighbor), and the descriptor reduces over neighbors with an
14+
order-independent ``mean``. Consequently ``need_sorted_nlist_for_lower()``
15+
correctly stays ``False`` for se_r and needs no ``self.compress`` override.
16+
17+
This test locks in that immunity. It runs ``forward_lower`` on an over-``rcut``
18+
FLAT nlist reversed so padding / out-of-``rcut`` neighbors precede the real ones
19+
-- exactly the input that broke se_a -- and asserts that the *compressed* result
20+
matches the *uncompressed* result on that **identical** nlist (energy and force
21+
to machine precision). On that same input the buggy se_a op diverged grossly,
22+
so this guard would catch an analogous se_r regression.
23+
24+
Why compare on the *same* nlist (not against a clean rcut-bounded one): this
25+
over-cut nlist has width ``== nnei`` and contains out-of-``rcut`` neighbors, so
26+
``format_nlist`` takes its *pad* branch (``n_nnei > nnei`` is false), which does
27+
NOT re-sort or rcut-filter. The pad branch is therefore mildly order-dependent
28+
(``nlist_distinguish_types`` truncates per-type sections in raw nlist order and
29+
lets over-``rcut`` neighbors leak in), so reversing the nlist shifts the
30+
*uncompressed* energy by ~1e-4. This is a property of ``format_nlist``, NOT of
31+
the reduction, and it affects se_a and se_r identically (uncompressed se_a shifts
32+
~4e-6 on the same input). Comparing compressed vs uncompressed on the *identical*
33+
nlist cancels that shared pad-branch effect and isolates the compression op:
34+
verified ``rel == 0`` (energy + force) -- whereas the buggy se_a op diverged
35+
grossly on this same input.
36+
"""
37+
38+
import copy
39+
import unittest
40+
41+
import torch
42+
43+
from deepmd.pt.cxx_op import (
44+
ENABLE_CUSTOMIZED_OP,
45+
)
46+
from deepmd.pt.model.model import (
47+
get_model,
48+
)
49+
from deepmd.pt.utils import (
50+
env,
51+
)
52+
from deepmd.pt.utils.nlist import (
53+
extend_input_and_build_neighbor_list,
54+
)
55+
56+
from ...seed import (
57+
GLOBAL_SEED,
58+
)
59+
from .test_forward_lower import (
60+
reduce_tensor,
61+
)
62+
63+
dtype = torch.float64
64+
65+
model_se_r = {
66+
"type_map": ["O", "H", "B"],
67+
"descriptor": {
68+
"type": "se_e2_r",
69+
"sel": [46, 92, 4],
70+
"rcut_smth": 0.50,
71+
"rcut": 4.00,
72+
"neuron": [25, 50, 100],
73+
"resnet_dt": False,
74+
"seed": 1,
75+
},
76+
"fitting_net": {
77+
"neuron": [24, 24, 24],
78+
"resnet_dt": True,
79+
"seed": 1,
80+
},
81+
"data_stat_nbatch": 20,
82+
}
83+
84+
85+
@unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built")
86+
class TestCompressedSeRForwardLower(unittest.TestCase):
87+
def setUp(self) -> None:
88+
model_params = copy.deepcopy(model_se_r)
89+
self.model = get_model(model_params).to(env.DEVICE)
90+
91+
def _make_system(self):
92+
# a sparse system: atoms have fewer neighbors than ``sel`` within
93+
# ``rcut``, and some neighbors fall in ``(rcut, rcut + buffer]`` so the
94+
# over-cut neighbor list contains zero-``sw`` real neighbors.
95+
natoms = 6
96+
cell = 6.0 * torch.eye(3, dtype=dtype, device=env.DEVICE)
97+
generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED)
98+
coord = 5.5 * torch.rand(
99+
[natoms, 3], dtype=dtype, device=env.DEVICE, generator=generator
100+
)
101+
atype = torch.tensor([0, 0, 1, 1, 2, 2], dtype=torch.int64, device=env.DEVICE)
102+
return coord, atype, cell
103+
104+
def _min_nbor_dist(self, coord, cell):
105+
# minimum image minimum pair distance, used as the compression lower bound
106+
box = torch.diagonal(cell)
107+
diff = coord[:, None, :] - coord[None, :, :]
108+
diff = diff - torch.round(diff / box) * box
109+
dist = torch.linalg.norm(diff, dim=-1)
110+
dist = dist + torch.eye(coord.shape[0], device=coord.device) * 1e10
111+
return float(dist.min())
112+
113+
def test_unsorted_overcut_nlist(self) -> None:
114+
coord, atype, cell = self._make_system()
115+
rcut = self.model.get_rcut()
116+
sel = self.model.get_sel()
117+
118+
# over-rcut FLAT neighbor list (mimics LAMMPS rcut+skin), reversed so the
119+
# out-of-rcut / padding neighbors precede the real ones -- the exact input
120+
# that broke compressed se_a.
121+
ec, ea, mp, nlist = extend_input_and_build_neighbor_list(
122+
coord.unsqueeze(0),
123+
atype.unsqueeze(0),
124+
rcut + 2.0,
125+
sum(sel),
126+
mixed_types=True,
127+
box=cell.unsqueeze(0),
128+
)
129+
nlist = torch.flip(nlist, dims=[-1])
130+
131+
# reference: uncompressed forward_lower on this exact nlist
132+
ref = self.model.forward_lower(ec, ea, nlist, mp, do_atomic_virial=False)
133+
134+
# enable compression (lower bound below the true min distance -> no
135+
# extrapolation) and rerun forward_lower on the IDENTICAL nlist
136+
self.model.min_nbor_dist = torch.tensor(
137+
0.9 * self._min_nbor_dist(coord, cell),
138+
dtype=env.GLOBAL_PT_FLOAT_PRECISION,
139+
device=env.DEVICE,
140+
)
141+
self.model.enable_compression()
142+
out = self.model.forward_lower(ec, ea, nlist, mp, do_atomic_virial=False)
143+
144+
torch.testing.assert_close(out["energy"], ref["energy"], rtol=1e-10, atol=1e-10)
145+
natoms = coord.shape[0]
146+
f_ref = reduce_tensor(ref["extended_force"], mp, natoms)
147+
f_out = reduce_tensor(out["extended_force"], mp, natoms)
148+
torch.testing.assert_close(f_out, f_ref, rtol=1e-10, atol=1e-10)
149+
150+
151+
if __name__ == "__main__":
152+
unittest.main()

0 commit comments

Comments
 (0)