From d21fd40c7ddad2e6d6810fc35cb52b3e4ee4ae3c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 7 Jul 2026 11:34:08 +0800 Subject: [PATCH 1/7] feat(dpmodel): graph-native tebd_input_mode='strip' kernel for dpa1 --- deepmd/dpmodel/descriptor/dpa1.py | 111 +++++++++++++++--- .../dpmodel/test_dpa1_call_graph_block.py | 90 ++++++++++++++ 2 files changed, 184 insertions(+), 17 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 80ad0f1b75..c09c021779 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -1749,7 +1749,7 @@ def call_graph( Notes ----- Known limitations: - - ``tebd_input_mode == "concat"`` only (strip mode lands later); + - ``tebd_input_mode`` in {"concat", "strip"}; compressed descriptors stay dense; - ``exclude_types`` is not yet supported and raises (lands in a later PR). """ from deepmd.dpmodel.utils.neighbor_graph import ( @@ -1757,9 +1757,9 @@ def call_graph( segment_sum, ) - if self.tebd_input_mode not in ["concat"]: + if self.tebd_input_mode not in ["concat", "strip"]: raise NotImplementedError( - "graph path supports tebd_input_mode='concat' only (NeighborGraph PR-A)" + f"graph path does not support tebd_input_mode={self.tebd_input_mode!r}" ) if self.exclude_types: raise NotImplementedError( @@ -1794,20 +1794,25 @@ def call_graph( ) # (E, 4), (E, 1) sw zeroed on padding # radial channel ss = rr[:, 0:1] # (E, 1) - # neighbor / center type embeddings (concat mode); ghost type == owner type - # so gathering by the LOCAL owner (src) reproduces the dense neighbor tebd. - # NB: do NOT wrap in ``xp.asarray(..., device=dev)`` -- that DETACHES under - # torch and severs the type-embedding weight gradient (the tebd net would - # never train); type_embedding already lives on the model device. - tebd = type_embedding - atype_embd_nlist = xp.take(tebd, nei_type, axis=0) # (E, tebd_dim) - if not self.type_one_side: - atype_embd_nnei = xp.take(tebd, center_type, axis=0) # (E, tebd_dim) - ss = xp.concat([ss, atype_embd_nlist, atype_embd_nnei], axis=-1) - else: - ss = xp.concat([ss, atype_embd_nlist], axis=-1) - # embedding net (same weights as the dense path); applies on the last axis - gg = self.embeddings[0].call(ss) # (E, ng) + if self.tebd_input_mode == "concat": + # neighbor / center type embeddings; ghost type == owner type so + # gathering by the LOCAL owner (src) reproduces the dense neighbor tebd. + # NB: do NOT wrap in ``xp.asarray(..., device=dev)`` -- that DETACHES + # under torch and severs the type-embedding weight gradient (the tebd + # net would never train); type_embedding already lives on the device. + tebd = type_embedding + atype_embd_nlist = xp.take(tebd, nei_type, axis=0) # (E, tebd_dim) + if not self.type_one_side: + atype_embd_nnei = xp.take(tebd, center_type, axis=0) # (E, tebd_dim) + ss = xp.concat([ss, atype_embd_nlist, atype_embd_nnei], axis=-1) + else: + ss = xp.concat([ss, atype_embd_nlist], axis=-1) + # embedding net (same weights as the dense path); applies on last axis + gg = self.embeddings[0].call(ss) # (E, ng) + else: # strip: factorized gg_s*gg_t + gg_s (per-edge; no neighbor coupling) + gg = self._graph_edge_gg_strip( + ss, center_type, nei_type, type_embedding, sw_e + ) # transformer attention over each center's edges — mirrors the dense # self.dpa1_attention(gg, nlist_mask, input_r, sw), which also runs on # the UNMASKED gg (padding rows are neutralized afterwards). @@ -1835,6 +1840,78 @@ def call_graph( rot_mat = gr[:, :, 1:] return grrg, rot_mat + def _graph_edge_gg_strip( + self, + ss: Array, + center_type: Array, + nei_type: Array, + type_embedding: Array, + sw_e: Array, + ) -> Array: + """Per-edge stripped-tebd embedding, op-for-op vs the dense strip branch. + + Mirrors the ``tebd_input_mode == "strip"`` block of :meth:`call`: the + geometric net runs on the radial channel only (``gg_s``), the stripped + type-embedding net produces a per-type(-pair) factor (``gg_t``, + optionally switch-smoothed), and the two combine as + ``gg_s * gg_t + gg_s``. The compression branches (geo/tebd) are NOT + reached on the graph route: :meth:`DescrptDPA1.uses_graph_lower` + excludes compressed descriptors, so this kernel assumes no compression. + + Parameters + ---------- + ss + (E, 1) per-edge radial channel (``rr[:, 0:1]``). + center_type + (E,) center (dst) LOCAL atom type of each edge. + nei_type + (E,) neighbor (src) LOCAL atom type of each edge. + type_embedding + (ntypes_with_padding, tebd_dim) type-embedding table. + sw_e + (E, 1) smooth switch, zeroed on padding edges. + + Returns + ------- + gg + (E, ng) per-edge embedding feeding the attention / segment_sum. + """ + assert self.embeddings_strip is not None + xp = array_api_compat.array_namespace(ss) + nt = self.tebd_dim + ntypes_with_padding = type_embedding.shape[0] + # geometric net on the radial channel only (dense: gg_s = cal_g(ss_scalar)) + gg_s = self.embeddings[0].call(ss) # (E, ng) + if self.type_one_side: + # one-side strip table indexed by NEIGHBOR type only + tt_full = self.cal_g_strip(type_embedding, 0) # (ntypes_pad, ng) + gg_t = xp.take(tt_full, nei_type, axis=0) # (E, ng) + else: + # two-side type-pair table; row = center * ntypes_pad + nei + # (dense builds the same (ntypes_pad**2, 2*nt) table, nei-fastest). + type_embedding_nei = xp.tile( + xp.reshape(type_embedding, (1, ntypes_with_padding, nt)), + (ntypes_with_padding, 1, 1), + ) + type_embedding_center = xp.tile( + xp.reshape(type_embedding, (ntypes_with_padding, 1, nt)), + (1, ntypes_with_padding, 1), + ) + two_side_type_embedding = xp.reshape( + xp.concat([type_embedding_nei, type_embedding_center], axis=-1), + (-1, nt * 2), + ) + tt_full = self.cal_g_strip( + two_side_type_embedding, 0 + ) # (ntypes_pad**2, ng) + # int64 for torch take (take_along/take requires Long indices) + idx = xp.astype(center_type * ntypes_with_padding + nei_type, xp.int64) + gg_t = xp.take(tt_full, idx, axis=0) # (E, ng) + if self.smooth: + # dense: gg_t = gg_t * sw (per-neighbor); sw_e is (E, 1), zeroed on padding + gg_t = gg_t * sw_e + return gg_s * gg_t + gg_s + def _graph_attention( self, gg: Array, diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py index 9a984a30f3..fb01caec55 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py @@ -114,3 +114,93 @@ def test_exclude_types_raises(self) -> None: dd.se_atten.call_graph( ng, self.atype.reshape(-1), type_embedding=dd.type_embedding.call() ) + + +class TestDpa1BlockCallGraphStrip: + """Bit-exact parity between the graph-native ``call_graph`` and the dense + ``call`` for ``tebd_input_mode='strip'``. + + Strip mode factorizes the per-neighbor feature as ``gg = gg_s*gg_t + gg_s`` + (radial embedding times type-pair strip embedding); it has no neighbor-axis + coupling, so the graph translation is edge-for-edge and must be bit-exact + with the dense block on the SAME neighbor list. + """ + + def setup_method(self) -> None: + rng = np.random.default_rng(3) + self.nloc = 4 + self.coord = rng.normal(size=(1, self.nloc, 3)) * 1.5 + self.atype = np.array([[0, 1, 0, 1]], dtype=np.int64) + + def _make(self, type_one_side: bool, smooth: bool, attn_layer: int) -> DescrptDPA1: + return DescrptDPA1( + rcut=4.0, + rcut_smth=0.5, + sel=[20], # non-binding sel: carry-all graph == dense on real neighbors + ntypes=2, + attn_layer=attn_layer, + axis_neuron=2, + neuron=[6, 12], + tebd_input_mode="strip", + type_one_side=type_one_side, + smooth_type_embedding=smooth, + ) + + def _assert_parity(self, dd: DescrptDPA1, compact: bool) -> None: + ( + ext_coord, + ext_atype, + mapping, + nlist, + ) = extend_input_and_build_neighbor_list( + self.coord, + self.atype, + dd.get_rcut(), + dd.get_sel(), + mixed_types=dd.mixed_types(), + box=None, + ) + tebd = dd.type_embedding.call() + nf, nall = ext_atype.shape + atype_embd_ext = np.reshape( + np.take(tebd, np.reshape(ext_atype, (-1,)), axis=0), + (nf, nall, dd.tebd_dim), + ) + dense_g, *_ = dd.se_atten.call( + nlist, + ext_coord, + ext_atype, + atype_embd_ext=atype_embd_ext, + mapping=None, + type_embedding=tebd, + ) + ng = from_dense_quartet(ext_coord, nlist, mapping, compact=compact) + graph_g, _rot_mat = dd.se_atten.call_graph( + ng, + np.reshape(ext_atype, (-1,)), + type_embedding=tebd, + ) + assert not np.any(np.isnan(graph_g)) + np.testing.assert_allclose( + graph_g.reshape(dense_g.shape), + dense_g, + rtol=1e-12, + atol=1e-12, + ) + + @pytest.mark.parametrize( + "type_one_side", [False, True] + ) # two-side vs one-side strip table + @pytest.mark.parametrize("smooth", [False, True]) # gg_t switch-smoothing branch + def test_strip_attn0_equals_dense(self, type_one_side, smooth) -> None: + """attn_layer=0: no attention, so strip parity is bit-exact for both smooth values.""" + dd = self._make(type_one_side, smooth, attn_layer=0) + self._assert_parity(dd, compact=True) + + @pytest.mark.parametrize( + "type_one_side", [False, True] + ) # two-side vs one-side strip table + def test_strip_attn2_equals_dense(self, type_one_side) -> None: + """attn_layer=2, smooth=False: bit-exact (avoids by-design smooth softmax divergence).""" + dd = self._make(type_one_side, smooth=False, attn_layer=2) + self._assert_parity(dd, compact=False) From c0707375410d3132f4ab27507702c9a3fe1f0c88 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 7 Jul 2026 12:03:51 +0800 Subject: [PATCH 2/7] feat(dpmodel): route dpa1/se_atten_v2 strip through the graph lower (gate keeps exclude_types/compression dense) --- deepmd/dpmodel/descriptor/dpa1.py | 16 +++- .../test_dpa1_call_graph_descriptor.py | 92 ++++++++++++++++++- .../test_dpa1_graph_attention_parity.py | 39 +++++++- 3 files changed, 133 insertions(+), 14 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index c09c021779..bf26cff8a7 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -432,9 +432,9 @@ def uses_graph_lower(self) -> bool: The graph-native lower (``call_graph``) covers the factorizable path AND transformer attention (``attn_layer >= 0``, NeighborGraph PR-D) - with concat type-embedding and no type exclusion. Remaining ineligible - configs (``tebd_input_mode == "strip"``, ``exclude_types``) fall back - to the legacy dense path, so those models keep working unchanged. + with concat OR strip type-embedding. Remaining ineligible configs + (``exclude_types``, and compressed descriptors) fall back to the legacy + dense path, so those models keep working unchanged. Eligibility does NOT imply numerical interchangeability with the dense route for every config: with ``smooth_type_embedding=True`` @@ -442,8 +442,14 @@ def uses_graph_lower(self) -> bool: differs from the dense lower by up to ~1e-4 (see the Notes of :meth:`call_graph`). """ + # compressed descriptors have no graph kernel (geo/tebd tabulation is + # dense-only); keep them on the legacy dense path. + if self.compress: + return False + # exclude_types stays dense (graph exclusion is owned elsewhere); strip is + # now graph-eligible (per-edge factorized embedding, no neighbor coupling). return ( - self.se_atten.tebd_input_mode == "concat" + self.se_atten.tebd_input_mode in ("concat", "strip") and not self.se_atten.exclude_types ) @@ -575,7 +581,7 @@ def call( nall = xp.reshape(coord_ext, (nlist.shape[0], -1)).shape[1] // 3 # graph-eligible configs route through the graph-native adapter (decision # #14: graph = single math source, dense call = thin adapter). Ineligible - # configs (attention, strip tebd, exclude_types) and the ghost case with + # configs (exclude_types, compressed descriptors) and the ghost case with # no mapping fall back to the legacy dense body. The graph needs `mapping` # to fold ghosts to local owners; without it only nall == nloc is valid. if self.uses_graph_lower() and (mapping is not None or nall == nloc): diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py index dc1d51da91..a9233dd1aa 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py @@ -99,14 +99,13 @@ def test_descriptor_graph_equals_dense_full_tuple(self, sel) -> None: @pytest.mark.parametrize( "kwargs", [ - {"tebd_input_mode": "strip"}, # strip tebd: graph unsupported -> dense {"exclude_types": [(0, 1)]}, # type exclusion: graph unsupported -> dense ], ) def test_ineligible_config_falls_back_to_dense(self, kwargs) -> None: - """attn_layer=0 configs the graph can't handle (strip tebd, exclude_types) - must report uses_graph_lower()=False and run the dense body without - raising (regression: Task-3 routing previously raised NotImplementedError). + """attn_layer=0 configs the graph can't handle (exclude_types) must report + uses_graph_lower()=False and run the dense body without raising (strip is + now graph-eligible; its routing + parity is covered by TestDpa1StripRouting). """ dd = DescrptDPA1( rcut=4.0, rcut_smth=0.5, sel=[30], ntypes=2, attn_layer=0, **kwargs @@ -200,3 +199,88 @@ def test_call_graph_returns_flat_node_axis(self) -> None: n = atype_local.shape[0] assert grrg.shape[0] == n and grrg.ndim == 2 assert rot_mat.shape[0] == n and rot_mat.ndim == 3 + + +class TestDpa1StripRouting: + """Strip is now graph-eligible: ``uses_graph_lower()`` admits it, ``call`` + routes through the graph adapter, and the adapter (``static_nnei``) is + bit-exact with the legacy ``_call_dense`` for EVERY strip config -- including + ``smooth_type_embedding=True`` at ``attn_layer>0`` (the adapter reproduces + the dense phantom-neighbor terms, unlike the direct carry-all ``call_graph``). + """ + + def setup_method(self) -> None: + rng = np.random.default_rng(7) + self.nloc = 4 + self.coord = rng.normal(size=(1, self.nloc, 3)) * 1.5 + self.atype = np.array([[0, 1, 0, 1]], dtype=np.int64) + + def _make(self, type_one_side: bool, smooth: bool, attn_layer: int) -> DescrptDPA1: + return DescrptDPA1( + rcut=4.0, + rcut_smth=0.5, + sel=[20], # non-binding + ntypes=2, + attn_layer=attn_layer, + axis_neuron=2, + neuron=[6, 12], + tebd_input_mode="strip", + type_one_side=type_one_side, + smooth_type_embedding=smooth, + resnet_dt=False, + ) + + def test_uses_graph_lower_strip_gate(self) -> None: + """The gate admits non-compressed strip; excludes compressed and exclude_types.""" + dd = self._make(type_one_side=False, smooth=True, attn_layer=2) + assert dd.uses_graph_lower() is True # strip is now graph-eligible + # negative contract: compression keeps the descriptor on the dense path + dd.compress = True + assert dd.uses_graph_lower() is False + dd.compress = False + # negative contract: exclude_types still forces dense (separate feature, + # owned by the pair-exclude PR; strip does NOT change its eligibility) + dd_excl = DescrptDPA1( + rcut=4.0, + rcut_smth=0.5, + sel=[20], + ntypes=2, + attn_layer=2, + tebd_input_mode="strip", + exclude_types=[(0, 1)], + resnet_dt=False, + ) + assert dd_excl.uses_graph_lower() is False + + @pytest.mark.parametrize("type_one_side", [False, True]) # strip table branch + @pytest.mark.parametrize( + "smooth", [False, True] + ) # switch-smoothing + smooth attention + @pytest.mark.parametrize("attn_layer", [0, 2]) # no-attn and multi-layer attention + def test_call_strip_graph_equals_dense( + self, type_one_side, smooth, attn_layer + ) -> None: + """The routed ``call`` (graph adapter) is bit-exact with ``_call_dense``.""" + dd = self._make(type_one_side, smooth, attn_layer) + assert dd.uses_graph_lower() is True # precondition: call routes to graph + ( + ext_coord, + ext_atype, + mapping, + nlist, + ) = extend_input_and_build_neighbor_list( + self.coord, + self.atype, + dd.get_rcut(), + dd.get_sel(), + mixed_types=dd.mixed_types(), + box=None, + ) + routed = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + dense = dd._call_dense(ext_coord, ext_atype, nlist) + assert len(routed) == len(dense) + for r, d in zip(routed, dense, strict=True): + if r is None: + assert d is None + continue + np.testing.assert_allclose(r, d, rtol=1e-12, atol=1e-12) diff --git a/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py b/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py index 97d044862f..f9ef26b54d 100644 --- a/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py +++ b/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py @@ -214,17 +214,46 @@ class TestGraphEligibility: def test_attention_concat_is_graph_eligible(self) -> None: assert _make(2).uses_graph_lower() - def test_strip_mode_stays_dense(self) -> None: - """se_atten_v2 (tebd_input_mode='strip') is NOT graph-eligible yet: - strip-mode graph support is a later PR; it must keep the dense route - (the PR-D plan's 'se_atten_v2 inherits for free' did not hold). + def test_se_atten_v2_is_graph_eligible(self) -> None: + """se_atten_v2 (tebd_input_mode='strip', smooth=True) is now graph-eligible. + + It is a DescrptDPA1 subclass with no exclude_types and no routing override, + so admitting strip closes the 'se_atten_v2 is dense-only' gap. (Was: strip + stayed dense.) """ from deepmd.dpmodel.descriptor.se_atten_v2 import ( DescrptSeAttenV2, ) dd = DescrptSeAttenV2(rcut=4.0, rcut_smth=0.5, sel=[20], ntypes=2, attn_layer=2) - assert not dd.uses_graph_lower() + assert dd.uses_graph_lower() is True + + def test_se_atten_v2_graph_equals_dense(self) -> None: + """The graph-routed se_atten_v2 ``call`` is bit-exact with ``_call_dense`` + (the ``static_nnei`` adapter reproduces the dense phantom terms despite + smooth=True) at a non-binding sel. + """ + from deepmd.dpmodel.descriptor.se_atten_v2 import ( + DescrptSeAttenV2, + ) + + rng = np.random.default_rng(GLOBAL_SEED) + nloc = 4 + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype = np.array([[0, 1, 0, 1]], dtype=np.int64) + dd = DescrptSeAttenV2(rcut=4.0, rcut_smth=0.5, sel=[20], ntypes=2, attn_layer=2) + assert dd.uses_graph_lower() is True + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + coord, atype, dd.get_rcut(), dd.get_sel(), mixed_types=True, box=None + ) + routed = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + dense = dd._call_dense(ext_coord, ext_atype, nlist) + assert len(routed) == len(dense) + for r, d in zip(routed, dense, strict=True): + if r is None: + assert d is None + continue + np.testing.assert_allclose(r, d, rtol=1e-12, atol=1e-12) class TestBindingSelDivergence: From 4a68c864e9028235acbc244c903ea44fc5186c69 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 7 Jul 2026 12:11:55 +0800 Subject: [PATCH 3/7] test(pt_expt): cover the strip graph kernel in dpa1 make_fx export tests --- source/tests/pt_expt/descriptor/test_dpa1.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index cddd22419f..2bfde949c1 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -252,8 +252,9 @@ def fn(coord_ext, atype_ext, nlist): atol=atol, ) + @pytest.mark.parametrize("tm", ["concat", "strip"]) # tebd_input_mode @pytest.mark.parametrize("prec", ["float64"]) # precision - def test_make_fx_graph(self, prec) -> None: + def test_make_fx_graph(self, prec, tm) -> None: """make_fx (export-readiness) of the attn_layer=0 GRAPH forward. For ``attn_layer == 0`` the dense ``forward`` routes through the @@ -275,6 +276,7 @@ def test_make_fx_graph(self, prec) -> None: self.sel_mix, self.nt, attn_layer=0, + tebd_input_mode=tm, precision=prec, seed=GLOBAL_SEED, ).to(self.device) @@ -311,9 +313,10 @@ def fn(coord_ext, atype_ext, nlist, mapping): atol=atol, ) + @pytest.mark.parametrize("tm", ["concat", "strip"]) # tebd_input_mode @pytest.mark.parametrize("smooth", [False, True]) # smooth attention branch @pytest.mark.parametrize("prec", ["float64"]) # precision - def test_make_fx_graph_attn(self, prec, smooth) -> None: + def test_make_fx_graph_attn(self, prec, smooth, tm) -> None: """make_fx (export-readiness) of the GRAPH forward with attention. MERGE BLOCKER (NeighborGraph PR-D): pt_expt compiled training routes @@ -335,6 +338,7 @@ def test_make_fx_graph_attn(self, prec, smooth) -> None: self.sel_mix, self.nt, attn_layer=2, + tebd_input_mode=tm, attn_dotr=True, smooth_type_embedding=smooth, precision=prec, From c4b41622b522c9ae91c0919824ea7f9d04ceb089 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 7 Jul 2026 14:32:34 +0800 Subject: [PATCH 4/7] docs(dpmodel): correct _call_dense docstring for strip/attention now graph-eligible --- deepmd/dpmodel/descriptor/dpa1.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index bf26cff8a7..3008daac84 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -686,8 +686,8 @@ def _call_dense( atype_ext: Array, nlist: Array, ) -> Array: - """Legacy dense descriptor body (the ineligible ``call`` path: attention, - strip tebd, exclude_types, or the no-mapping ghost case). + """Legacy dense descriptor body (the ineligible ``call`` path: + compressed descriptors, exclude_types, or the no-mapping ghost case). Parameters ---------- From e9356680c3eeca54c143ce84e47446153cda98e2 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 8 Jul 2026 11:37:15 +0800 Subject: [PATCH 5/7] test(dpmodel): add NaN guard before strip/attention parity assert_allclose np.testing.assert_allclose treats NaN == NaN as equal by default, so a shared NaN between the graph and dense outputs would silently pass the parity check. Applies to the two strip/attention parity tests CodeRabbit flagged: test_call_strip_graph_equals_dense (test_dpa1_call_graph_descriptor.py) and test_se_atten_v2_graph_equals_dense (test_dpa1_graph_attention_parity.py). --- source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py | 1 + source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py | 1 + 2 files changed, 2 insertions(+) diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py index a9233dd1aa..6632d82894 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py @@ -283,4 +283,5 @@ def test_call_strip_graph_equals_dense( if r is None: assert d is None continue + assert not np.any(np.isnan(r)) np.testing.assert_allclose(r, d, rtol=1e-12, atol=1e-12) diff --git a/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py b/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py index f9ef26b54d..4eec03ba0b 100644 --- a/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py +++ b/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py @@ -253,6 +253,7 @@ def test_se_atten_v2_graph_equals_dense(self) -> None: if r is None: assert d is None continue + assert not np.any(np.isnan(r)) np.testing.assert_allclose(r, d, rtol=1e-12, atol=1e-12) From 66089f1ab875047cf96cc8d5f69e6a152fc7897c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 8 Jul 2026 11:47:42 +0800 Subject: [PATCH 6/7] test(pt_expt): exclude se_atten_v2 from the None-vs-explicit nlist dispatch test CI failure: test_default_fallback[se_atten_v2] failed (energy off by 1.2e-07, rtol=1e-10). Root cause is this PR's own change: uses_graph_lower() now admits tebd_input_mode='strip', which makes se_atten_v2 (always strip, always smooth_type_embedding=True) graph-eligible. call_common's dispatch (make_model.py) forces the dense route whenever an explicit neighbor_list is passed, but pt_expt's default-flip (decision #17) routes neighbor_list=None to the carry-all graph for graph-eligible mixed_types descriptors -- so 'None' and 'explicit DefaultNeighborList()' stop being the same computation for this model. model_dpa1 already works around the same issue by pinning smooth_type_embedding=False (with a comment explaining why); se_atten_v2 can't do that since DescrptSeAttenV2 hardcodes smooth_type_embedding=True, so the divergence is unavoidable and documented (NeighborGraph PR-D: dense keeps sel-padding phantom softmax terms the graph route doesn't). Exclude it from this specific parametrize with a docstring explaining why; it is still covered by test_pt_expt_equivalence/test_dpmodel_equivalence, which always pass an explicit neighbor_list on both sides and so never hit this route split. --- .../tests/pt_expt/utils/test_neighbor_list.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/source/tests/pt_expt/utils/test_neighbor_list.py b/source/tests/pt_expt/utils/test_neighbor_list.py index ff29ed50c4..ce3c9a0e95 100644 --- a/source/tests/pt_expt/utils/test_neighbor_list.py +++ b/source/tests/pt_expt/utils/test_neighbor_list.py @@ -479,7 +479,9 @@ def test_pt_expt_multiframe_equivalence(name: str) -> None: ) -@pytest.mark.parametrize("name", list(ALL_MODELS)) # descriptor family +@pytest.mark.parametrize( + "name", [n for n in ALL_MODELS if n != "se_atten_v2"] +) # descriptor family def test_default_fallback(name: str) -> None: """``neighbor_list=None`` dispatches to the same DefaultNeighborList builder. @@ -491,6 +493,19 @@ def test_default_fallback(name: str) -> None: GNN message-passing scatter (atomic adds) is not bit-reproducible run-to-run, so the virial can differ by ~1 ULP between the passes (a real dispatch bug would differ by orders of magnitude more). + + ``se_atten_v2`` is excluded: that equivalence assumption only holds for the + DENSE-nlist route. For a ``mixed_types`` descriptor with + ``uses_graph_lower() == True``, passing an explicit ``neighbor_list`` forces + the dense route (``call_common``'s ``neighbor_list is not None`` branch), + while ``None`` lets pt_expt's default-flip (decision #17) route to the + carry-all graph instead -- two genuinely different algorithms, not two + evaluations of one. ``se_atten_v2`` hardcodes ``smooth_type_embedding=True`` + (unlike ``model_dpa1`` above, which pins it ``False`` for exactly this + reason), so graph and dense intentionally diverge here (NeighborGraph PR-D: + dense keeps sel-padding phantom terms in the attention softmax denominator, + the graph route does not) -- there is no tolerance that would make this a + meaningful dispatch-equivalence check for it. """ coord_np, atype_np, box_np = _system() md = get_model(copy.deepcopy(ALL_MODELS[name])).to(env.DEVICE) From b77445503ccc3f22ef9b3d369ff01b56dbf569e6 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 8 Jul 2026 12:10:16 +0800 Subject: [PATCH 7/7] test(pt_expt): assert bounded closeness instead of excluding se_atten_v2 Follow-up to 66089f1ab. That commit fixed the CI failure by excluding se_atten_v2 from test_default_fallback's parametrize, but a plain exclusion silently drops coverage rather than documenting the expected behavior. Replace it with a positive assertion, mirroring the precedent already set by test_block_compact_graph_smooth_clean_divergence (test_dpa1_graph_attention_parity.py): for models where tebd_input_mode in {concat, strip} + attn_layer > 0 + smooth_type_embedding=True, dense and the carry-all graph default (neighbor_list=None) intentionally diverge (NeighborGraph PR-D: dense keeps sel-padding phantom terms in the attention softmax denominator, the graph route does not). Verified empirically this is not strip-specific: a plain concat se_atten with smooth_type_embedding=True shows the same ~1e-7 order-of-magnitude divergence at attn_layer>0, and exactly zero divergence at attn_layer=0 (no softmax involved there at all) -- confirming the gap lives entirely in the shared attention code, not in anything this PR added. Add model_dpa1_smooth (concat's counterpart to model_se_atten_v2) so both tebd_input_mode values are covered by the same mechanism, and a KNOWN_GRAPH_DENSE_DIVERGENT set used only by test_default_fallback to widen its tolerance for exactly these two models (atol=3e-5, rtol=1e-3 -- measured empirically: energy ~1e-7, force ~1e-6, virial up to ~1.3e-5, since virial is a derivative and amplifies the softmax-denominator perturbation more than the value itself). All other models keep the original tight tolerance. --- .../tests/pt_expt/utils/test_neighbor_list.py | 75 ++++++++++++++++--- 1 file changed, 64 insertions(+), 11 deletions(-) diff --git a/source/tests/pt_expt/utils/test_neighbor_list.py b/source/tests/pt_expt/utils/test_neighbor_list.py index ce3c9a0e95..c52bdf03c2 100644 --- a/source/tests/pt_expt/utils/test_neighbor_list.py +++ b/source/tests/pt_expt/utils/test_neighbor_list.py @@ -123,6 +123,36 @@ "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, } +model_dpa1_smooth = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "se_atten", + "sel": 40, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12, 24], + "axis_neuron": 4, + "attn": 16, + "attn_layer": 2, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": False, + "temperature": 1.0, + "set_davg_zero": True, + "type_one_side": True, + # concat's counterpart to model_se_atten_v2 below: smooth attention + # left ON (unlike model_dpa1 above, which pins it off), so the + # tebd_input_mode="concat" + attn_layer>0 carry-all-vs-dense + # divergence (see test_default_fallback's KNOWN_GRAPH_DENSE_DIVERGENT) + # is exercised for concat too, not just strip (se_atten_v2). + "smooth_type_embedding": True, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + model_se_atten_v2 = { "type_map": TYPE_MAP, "descriptor": { @@ -245,12 +275,22 @@ "se_r": model_se_r, "se_e3": model_se_e3, "dpa1": model_dpa1, + "dpa1_smooth": model_dpa1_smooth, "se_atten_v2": model_se_atten_v2, "dpa2": model_dpa2, "dpa3": model_dpa3, "hybrid": model_hybrid, } +# tebd_input_mode in {"concat", "strip"} with attn_layer > 0 and +# smooth_type_embedding=True: the carry-all graph default +# (neighbor_list=None) intentionally diverges from the dense route (see +# test_default_fallback's docstring). Both modes hit the same shared +# attention softmax mechanism (dpa1.py's `_graph_attention`, gated only on +# `attn_layer > 0`, entered identically regardless of concat/strip), so one +# tolerance covers both. +KNOWN_GRAPH_DENSE_DIVERGENT = {"dpa1_smooth", "se_atten_v2"} + def _system(natoms: int = 6, box_len: float = 10.0, seed: int = GLOBAL_SEED): """A small 3-type periodic system; returns numpy (coord, atype, box).""" @@ -479,9 +519,7 @@ def test_pt_expt_multiframe_equivalence(name: str) -> None: ) -@pytest.mark.parametrize( - "name", [n for n in ALL_MODELS if n != "se_atten_v2"] -) # descriptor family +@pytest.mark.parametrize("name", list(ALL_MODELS)) # descriptor family def test_default_fallback(name: str) -> None: """``neighbor_list=None`` dispatches to the same DefaultNeighborList builder. @@ -494,18 +532,29 @@ def test_default_fallback(name: str) -> None: so the virial can differ by ~1 ULP between the passes (a real dispatch bug would differ by orders of magnitude more). - ``se_atten_v2`` is excluded: that equivalence assumption only holds for the + ``KNOWN_GRAPH_DENSE_DIVERGENT`` models (``dpa1_smooth``, ``se_atten_v2``) + are a special case: that "same builder" premise only holds for the DENSE-nlist route. For a ``mixed_types`` descriptor with ``uses_graph_lower() == True``, passing an explicit ``neighbor_list`` forces the dense route (``call_common``'s ``neighbor_list is not None`` branch), while ``None`` lets pt_expt's default-flip (decision #17) route to the carry-all graph instead -- two genuinely different algorithms, not two - evaluations of one. ``se_atten_v2`` hardcodes ``smooth_type_embedding=True`` + evaluations of one. Both hardcode/pin ``smooth_type_embedding=True`` (unlike ``model_dpa1`` above, which pins it ``False`` for exactly this - reason), so graph and dense intentionally diverge here (NeighborGraph PR-D: - dense keeps sel-padding phantom terms in the attention softmax denominator, - the graph route does not) -- there is no tolerance that would make this a - meaningful dispatch-equivalence check for it. + reason), so graph and dense intentionally diverge (NeighborGraph PR-D: + dense keeps sel-padding phantom terms in the attention softmax + denominator, the graph route does not -- see + ``test_block_compact_graph_smooth_clean_divergence`` in + ``test_dpa1_graph_attention_parity.py`` for the same invariant at the + block level). At this test's non-binding ``sel=40`` (vs. <=5 real + neighbors), the gap is small but non-zero and deterministic (not CUDA + ULP-style non-determinism): energy differs by ~1e-7, force by ~1e-6, + virial (a derivative, so it amplifies the softmax-denominator + perturbation more than the value itself) by up to ~1.3e-5. We assert + BOUNDED closeness (atol=3e-5, rtol=1e-3 -- individual virial/force + components can be near-zero, so atol dominates) rather than bit-identity + for these two, keeping the check meaningful rather than silently dropping + coverage. """ coord_np, atype_np, box_np = _system() md = get_model(copy.deepcopy(ALL_MODELS[name])).to(env.DEVICE) @@ -521,11 +570,15 @@ def test_default_fallback(name: str) -> None: coord_np, dtype=torch.float64, device=env.DEVICE ).requires_grad_(True) outs[tag] = md.forward(coord_t, atype_t, box=box_t, do_atomic_virial=True, **kw) + tol = ( + {"rtol": 1e-3, "atol": 3e-5} + if name in KNOWN_GRAPH_DENSE_DIVERGENT + else {"rtol": 1e-10, "atol": 1e-12} + ) for k in ("energy", "force", "virial"): np.testing.assert_allclose( outs["none"][k].detach().cpu().numpy(), outs["explicit"][k].detach().cpu().numpy(), - rtol=1e-10, - atol=1e-12, err_msg=f"{name} {k}", + **tol, )