From 5c14c07a1abd2648324bf69e0566fa2e6e68e266 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 13 Jul 2026 20:22:03 +0800 Subject: [PATCH] fix(dpmodel): dpa1 dense call must not route through divergent graph adapter DescrptDPA1.call routed graph-eligible configs through _call_graph_adapter (decision #14: dense call = thin adapter). The adapter is bit-exact vs _call_dense ONLY in the trivial-statistics regime (davg == 0): the dense se_atten body leaks a phantom padding-neighbor -davg/dstd residual (EnvMat.call subtracts davg AFTER the padding rows' geometry is weight-zeroed; with empty exclude_types nothing re-masks it, at ANY attn_layer) that the graph path deliberately omits. Measured on the consistency fixture: dense vs adapter agree at 0.0 for davg=0, diverge by ~20 (abs) for nonzero davg, both attn_layer=0 and 2. The gate also made 'mapping' -- an argument that only enables ghost folding on graph routes -- silently change dense numerics. The divergence never surfaced because the pt/pd consistency tests invoke dd2.call() WITHOUT mapping (gate off -> dense route), while production model-level forwards (DPAtomicModel.forward_atomic) always pass mapping: dpmodel/jax/tf2 inference of a trained dpa1 with set_davg_zero=False (the constructor default) silently used the adapter and differed from the pt/tf dense backends. Cross-backend suites stayed green only because fresh models sit in the trivial-stat regime. Fix (same pattern as the dpa2 graph PR): - DescrptDPA1.call always runs _call_dense; the graph-native route is reached exclusively through call_graph (pt_expt forward_atomic_graph and the graph .pt2). _call_graph_adapter is retained as the bit-exact-regime reference and its docstring now states the actual regime (incl. the slot-0 statistics simplification). - pt/pd consistency tests now pass mapping (the production invocation); red on the old routing, green now. - New regression test: call(mapping) == call(None) == _call_dense with nonzero injected davg and ghosts present. - Adapter parity tests exercise _call_graph_adapter directly instead of relying on the removed call routing. --- deepmd/dpmodel/descriptor/dpa1.py | 52 +++++++++---- .../test_dpa1_call_graph_descriptor.py | 78 ++++++++++++------- .../test_dpa1_graph_attention_parity.py | 12 +-- source/tests/pd/model/test_dpa1.py | 7 +- source/tests/pt/model/test_dpa1.py | 7 +- 5 files changed, 104 insertions(+), 52 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 3a47a7319c..01d5dd8983 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -595,18 +595,25 @@ def call( sw The smooth switch function. """ - xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) - nloc = nlist.shape[1] - 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 (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): - return self._call_graph_adapter(coord_ext, atype_ext, nlist, mapping) - else: - return self._call_dense(coord_ext, atype_ext, nlist) + # The dense ``call`` always runs the legacy dense body -- it is the + # cross-backend consistency reference and must match the tf/pt/pd/jax + # dense descriptors bit-for-bit. It previously routed graph-eligible + # configs through ``_call_graph_adapter`` (decision #14), but the + # adapter is bit-exact ONLY in the trivial-statistics regime + # (``davg == 0``): the dense se_atten body leaks a phantom + # padding-neighbor ``-davg/dstd`` residual (``EnvMat.call`` subtracts + # ``davg`` AFTER the padding rows' geometry is weight-zeroed, and + # neither the empty ``exclude_types`` mask nor the attention layers + # re-mask it) that the graph path deliberately omits -- the graph + # output is the physically correct one, but ``call`` must reproduce + # the dense reference. The former gate also made ``mapping`` -- an + # argument that only enables ghost folding on graph routes -- silently + # change the numerics of a dense call. The graph-native route is + # reached exclusively through :meth:`call_graph` (pt_expt + # ``forward_atomic_graph`` and the graph ``.pt2``), never through + # ``call``. ``_call_graph_adapter`` is retained as the + # bit-exact-regime reference exercised by the adapter parity tests. + return self._call_dense(coord_ext, atype_ext, nlist) def _call_graph_adapter( self, @@ -615,14 +622,27 @@ def _call_graph_adapter( nlist: Array, mapping: Array | None, ) -> Array: - """Regime-1 dense->graph adapter (the eligible ``call`` path). + """Regime-1 dense->graph adapter. Builds a NeighborGraph from the dense quartet with the SHAPE-STATIC converter (``compact=False``, so this is jit/export-traceable -- no ``nonzero``), runs :meth:`call_graph`, and reconstructs the dense-shaped - ``sw``. Preserves the dense 5-tuple ABI exactly; masked invalid edges - contribute zero in ``call_graph``'s ``segment_sum`` so the output is - identical to the legacy dense body. + ``sw``. Preserves the dense 5-tuple ABI; masked invalid edges + contribute zero in ``call_graph``'s ``segment_sum``. + + Bit-exact vs :meth:`_call_dense` **only in the trivial-statistics + regime** (``davg == 0``). For nonzero ``davg`` the dense body leaks a + phantom padding-neighbor ``-davg/dstd`` residual into every padding + slot (``EnvMat.call`` subtracts ``davg`` AFTER the padding geometry is + weight-zeroed; with empty ``exclude_types`` nothing re-masks it, at + any ``attn_layer``) that the graph path deliberately omits. The graph + kernel additionally applies the slot-0 statistics ``[:, 0, :]`` to + every edge -- exact for real stat-computed tables (slot-uniform by + construction), not for artificially slot-varying ones. This is why + the dense :meth:`call` does NOT route here: it is the cross-backend + consistency reference. This method is retained as the + bit-exact-regime reference exercised by the adapter parity tests; the + production graph route is :meth:`call_graph`. Parameters ---------- 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 b889644abc..ffc67bec9e 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Full 5-tuple ABI parity between the graph-routed ``DescrptDPA1.call`` -(attn_layer=0, which now goes ``from_dense_quartet -> call_graph``) and the -legacy dense descriptor output captured BEFORE the swap, for binding AND -non-binding ``sel``. +"""Full 5-tuple ABI parity between ``DescrptDPA1._call_graph_adapter`` +(``from_dense_quartet -> call_graph``) and the legacy dense descriptor +output, for binding AND non-binding ``sel``. -The dense reference is reconstructed by calling the BLOCK directly -(``dd.se_atten.call``) and applying the descriptor-level ``concat_output_tebd`` -step by hand (mirroring dpa1.py), because ``dd.call`` itself now routes through -the graph for ``attn_layer == 0``. +The public ``dd.call`` does NOT route through the adapter (it is the +cross-backend dense reference; the adapter is bit-exact only in the +trivial-statistics regime these fresh models are in), so the adapter is +exercised directly. The dense reference is reconstructed by calling the +BLOCK directly (``dd.se_atten.call``) and applying the descriptor-level +``concat_output_tebd`` step by hand (mirroring dpa1.py). """ import numpy as np @@ -66,7 +67,7 @@ def _dense_reference(self, dd, ext_coord, ext_atype, nlist): @pytest.mark.parametrize("sel", [[30], [4]]) # non-binding AND binding def test_descriptor_graph_equals_dense_full_tuple(self, sel) -> None: - """Graph-routed dd.call() returns the identical dense 5-tuple ABI.""" + """``_call_graph_adapter`` returns the identical dense 5-tuple ABI.""" dd = self._make(sel) ( ext_coord, @@ -81,10 +82,9 @@ def test_descriptor_graph_equals_dense_full_tuple(self, sel) -> None: mixed_types=dd.mixed_types(), box=None, ) - # dense reference captured via the block (pre-swap behaviour) + # dense reference captured via the block ref = self._dense_reference(dd, ext_coord, ext_atype, nlist) - # the swapped public ABI: routes through the graph - out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + out = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) assert len(out) == 5 # grrg np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) @@ -138,8 +138,8 @@ def test_exclude_types_graph_eligible_and_parity(self, exclude_types) -> None: ) # dense reference (calls block directly) ref = self._dense_reference(dd, ext_coord, ext_atype, nlist) - # graph-routed public call - out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + # graph adapter, exercised directly (the public ``call`` is dense) + out = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) assert len(out) == 5 np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) np.testing.assert_allclose(out[1], ref[1], rtol=1e-12, atol=1e-12) @@ -159,13 +159,26 @@ def test_exclude_types_graph_eligible_and_parity(self, exclude_types) -> None: assert not np.any(np.isnan(grrg_g)) assert not np.any(np.isinf(grrg_g)) - def test_eligible_no_mapping_with_ghosts_falls_back(self) -> None: - """An eligible (concat) attn_layer=0 descriptor called with mapping=None - on a PERIODIC system (nall > nloc ghosts) must fall back to the dense - body and match it (regression: the graph needs mapping for ghosts, the - identity-mapping default previously indexed out of range). + def test_call_is_mapping_insensitive_with_nontrivial_stats(self) -> None: + """``call(..., mapping)`` == ``call(..., None)`` == ``_call_dense`` with + NONZERO ``davg`` and ghosts present. + + Regression for the removed graph-adapter routing gate: ``mapping`` only + enables ghost folding on graph routes and must never change the dense + numerics. Under the old gate, passing ``mapping`` silently switched an + eligible ``call`` to ``_call_graph_adapter``, whose output differs from + the dense reference whenever ``davg != 0`` (the dense phantom + padding-slot ``-davg/dstd`` leak the graph path deliberately omits) -- + so the same physical system gave different descriptors depending on + whether the caller supplied ``mapping``. Nonzero stats are injected + precisely because the fresh-model default (``davg == 0``) is the one + regime where the two routes coincide and the bug is invisible. """ + rng = np.random.default_rng(7) dd = self._make([30]) + nnei = sum(dd.get_sel()) + dd.se_atten.mean = rng.normal(size=(2, nnei, 4)) + dd.se_atten.stddev = 0.1 + np.abs(rng.normal(size=(2, nnei, 4))) box = np.eye(3, dtype=np.float64)[None] * 6.0 ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( self.coord, @@ -176,9 +189,13 @@ def test_eligible_no_mapping_with_ghosts_falls_back(self) -> None: box=box, ) assert ext_atype.shape[1] > self.nloc # ghosts present - ref = self._dense_reference(dd, ext_coord, ext_atype, nlist) - out = dd.call(ext_coord, ext_atype, nlist, mapping=None) # must not IndexError - np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) + with_mapping = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + without_mapping = dd.call(ext_coord, ext_atype, nlist, mapping=None) + dense = dd._call_dense(ext_coord, ext_atype, nlist) + np.testing.assert_allclose( + with_mapping[0], without_mapping[0], rtol=0.0, atol=0.0 + ) + np.testing.assert_allclose(with_mapping[0], dense[0], rtol=0.0, atol=0.0) def test_single_rank_extension_keeps_type_invariant(self) -> None: """The ghost-free graph types a neighbor as ``atype[mapping[neighbor]]`` @@ -188,7 +205,7 @@ def test_single_rank_extension_keeps_type_invariant(self) -> None: a periodic image of its owner and shares its type. This test pins that invariant (an inconsistent ``mapping`` like the universal fixture's old buggy permutation is NOT a valid single-rank extension) and confirms the - graph-routed ``call`` matches dense on the resulting quartet. + graph adapter matches dense on the resulting quartet. """ dd = self._make([30]) box = np.eye(3, dtype=np.float64)[None] * 6.0 @@ -208,7 +225,7 @@ def test_single_rank_extension_keeps_type_invariant(self) -> None: ext_atype[f], ext_atype[f][mapping[f]] ) # atype_ext[k] == atype[mapping[k]] ref = self._dense_reference(dd, ext_coord, ext_atype, nlist) - out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + out = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) np.testing.assert_allclose(out[1], ref[1], rtol=1e-12, atol=1e-12) np.testing.assert_allclose(out[4], ref[4], rtol=1e-12, atol=1e-12) @@ -299,9 +316,12 @@ def test_uses_graph_lower_strip_gate(self) -> None: 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``.""" + """The strip graph adapter is bit-exact with ``_call_dense`` for a + fresh model (trivial statistics -- the adapter's bit-exact regime; + the public ``call`` itself is always dense). + """ dd = self._make(type_one_side, smooth, attn_layer) - assert dd.uses_graph_lower() is True # precondition: call routes to graph + assert dd.uses_graph_lower() is True # strip is graph-eligible ( ext_coord, ext_atype, @@ -315,10 +335,10 @@ def test_call_strip_graph_equals_dense( mixed_types=dd.mixed_types(), box=None, ) - routed = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + adapter = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) dense = dd._call_dense(ext_coord, ext_atype, nlist) - assert len(routed) == len(dense) - for r, d in zip(routed, dense, strict=True): + assert len(adapter) == len(dense) + for r, d in zip(adapter, dense, strict=True): if r is None: assert d is None continue 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 4eec03ba0b..4256c4cd3d 100644 --- a/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py +++ b/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py @@ -229,9 +229,11 @@ def test_se_atten_v2_is_graph_eligible(self) -> None: 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 se_atten_v2 graph adapter is bit-exact with ``_call_dense`` (the ``static_nnei`` adapter reproduces the dense phantom terms despite - smooth=True) at a non-binding sel. + smooth=True) at a non-binding sel, for a fresh model (trivial + statistics -- the adapter's bit-exact regime; the public ``call`` + itself is always dense). """ from deepmd.dpmodel.descriptor.se_atten_v2 import ( DescrptSeAttenV2, @@ -246,10 +248,10 @@ def test_se_atten_v2_graph_equals_dense(self) -> None: 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) + adapter = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) dense = dd._call_dense(ext_coord, ext_atype, nlist) - assert len(routed) == len(dense) - for r, d in zip(routed, dense, strict=True): + assert len(adapter) == len(dense) + for r, d in zip(adapter, dense, strict=True): if r is None: assert d is None continue diff --git a/source/tests/pd/model/test_dpa1.py b/source/tests/pd/model/test_dpa1.py index 285dd3d4cd..952f14c880 100644 --- a/source/tests/pd/model/test_dpa1.py +++ b/source/tests/pd/model/test_dpa1.py @@ -97,12 +97,17 @@ def test_consistency( atol=atol, err_msg=err_msg, ) - # dp impl + # dp impl. `mapping` is passed because that is the production + # invocation (DPAtomicModel.forward_atomic always forwards it); + # the dense `.call()` must give the same answer with or without + # it -- mapping only enables ghost folding on graph routes, it + # must never change the dense numerics. dd2 = DPDescrptDPA1.deserialize(dd0.serialize()) rd2, _, _, _, _ = dd2.call( self.coord_ext, self.atype_ext, self.nlist, + self.mapping, ) np.testing.assert_allclose( rd0.detach().cpu().numpy(), diff --git a/source/tests/pt/model/test_dpa1.py b/source/tests/pt/model/test_dpa1.py index 36b7373bed..bdb306f4de 100644 --- a/source/tests/pt/model/test_dpa1.py +++ b/source/tests/pt/model/test_dpa1.py @@ -104,12 +104,17 @@ def test_consistency( atol=atol, err_msg=err_msg, ) - # dp impl + # dp impl. `mapping` is passed because that is the production + # invocation (DPAtomicModel.forward_atomic always forwards it); + # the dense `.call()` must give the same answer with or without + # it -- mapping only enables ghost folding on graph routes, it + # must never change the dense numerics. dd2 = DPDescrptDPA1.deserialize(dd0.serialize()) rd2, _, _, _, _ = dd2.call( self.coord_ext, self.atype_ext, self.nlist, + self.mapping, ) np.testing.assert_allclose( rd0.detach().cpu().numpy(),