From ce027648a2117c1a967f35d640551db61be6c83f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 08:55:46 +0800 Subject: [PATCH 01/26] refactor(descriptor): split needs-exchange from supports-edge-parallel has_message_passing_across_ranks conflated the two; the SFPG bridging veto moves to the new supports_edge_parallel() (issue #5906 Task 4 groundwork). --- deepmd/dpmodel/descriptor/dpa4.py | 20 +++++++++----- .../descriptor/make_base_descriptor.py | 11 ++++++++ deepmd/pt/model/descriptor/sezm.py | 22 ++++++++++------ deepmd/pt/model/model/sezm_model.py | 2 +- .../tests/common/dpmodel/test_descrpt_dpa4.py | 26 +++++++++++++------ 5 files changed, 57 insertions(+), 24 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 2687d9535d..95f43e09ca 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -2278,19 +2278,25 @@ def has_message_passing(self) -> bool: return True def has_message_passing_across_ranks(self) -> bool: - """Whether multi-rank inference needs cross-rank ghost exchange. + """SeZM reads ghost-neighbour features at every interaction block. - SeZM reads ghost-neighbour features at every interaction block; the - GRAPH lower implements the exchange via per-block ``border_op`` - (pt_expt ``exchange_ghost_features``). Source Freeze Propagation - bridging is excluded: its per-node gate folds a node's entire - outgoing-edge set, which a single rank cannot observe for ghost - owners, so bridging models fail fast on multi-rank instead. + The GRAPH lower implements the exchange via per-block ``border_op`` + (pt_expt ``exchange_ghost_features``), so multi-rank inference always + needs the with-comm artifact. Whether multi-rank is POSSIBLE at all + is :meth:`supports_edge_parallel` (bridging vetoes it there). The DENSE (nlist) lower remains comm-less — see :meth:`dense_lower_supports_comm`; the freeze machinery consults both so nlist-kind artifacts carry ``has_comm_artifact=False``. """ + return True + + def supports_edge_parallel(self) -> bool: + """Bridging vetoes multi-rank until the SFPG exchange lands. + + The Source Freeze Propagation Gate folds each node's full + outgoing-edge set, which no single rank observes (issue #5906). + """ return self.bridging_switch is None def dense_lower_supports_comm(self) -> bool: diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index 6594e4c632..3e160782ed 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -137,6 +137,17 @@ def has_message_passing_across_ranks(self) -> bool: """ return False + def supports_edge_parallel(self) -> bool: + """Whether this descriptor can run under MPI domain decomposition. + + Distinct from :meth:`has_message_passing_across_ranks` (whether + multi-rank inference NEEDS a per-block ghost exchange): this asks + whether any part of the computation folds state that a single + rank cannot observe. Default ``True``: an ordinary descriptor + reads only rank-local neighbourhoods. + """ + return True + def supports_native_spin(self) -> bool: """Returns whether the descriptor natively conditions on per-atom spin. diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index d39f7a1028..9b0dabd052 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -2280,14 +2280,20 @@ def has_message_passing(self) -> bool: return True def has_message_passing_across_ranks(self) -> bool: - """Whether multi-rank inference needs cross-rank ghost-feature exchange. - - SeZM reads ghost-neighbour features at every interaction block, so a - domain-decomposed run must exchange them through ``border_op``. Source - Freeze Propagation bridging is excluded: its per-node gate folds a - node's entire outgoing-edge set, which a single rank cannot observe for - ghost owners, so the edge-based with-comm artifact is not exported for - bridging models and multi-rank inference fails fast instead. + """SeZM reads ghost-neighbour features at every interaction block. + + A domain-decomposed run must exchange them through ``border_op``, so + multi-rank inference always needs the with-comm exchange. Whether + multi-rank is POSSIBLE at all is :meth:`supports_edge_parallel` + (bridging vetoes it there). + """ + return True + + def supports_edge_parallel(self) -> bool: + """Bridging vetoes multi-rank until the SFPG exchange lands. + + The Source Freeze Propagation Gate folds each node's full + outgoing-edge set, which no single rank observes (issue #5906). """ return self.bridging_switch is None diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index b623222a41..98780899db 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -3041,7 +3041,7 @@ def supports_edge_parallel(self) -> bool: if self.inter_potential is not None: return False descriptor = self.atomic_model.descriptor - return bool(descriptor.has_message_passing_across_ranks()) + return bool(descriptor.supports_edge_parallel()) def export_lower_input_kind(self) -> str: """Return the ABI consumed by the exported ``.pt2`` lower graph. diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py index d2b238ae8e..51e717124a 100644 --- a/source/tests/common/dpmodel/test_descrpt_dpa4.py +++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py @@ -95,22 +95,32 @@ def test_shapes_and_interface(self) -> None: def test_message_passing_semantics(self) -> None: # SeZM always resolves ghost neighbours on the lower path, so it always # reports message passing. The GRAPH lower implements the cross-rank - # exchange via a real per-layer border_op, so a plain (non-bridging) - # descriptor reports across_ranks True; its DENSE lower has no + # exchange via a real per-layer border_op, so every SeZM descriptor + # (bridged or not) reports across_ranks True; its DENSE lower has no # comm_dict implementation (the dense adapter raises on it), so # dense_lower_supports_comm() is False and the freeze machinery - # skips the dead dense with-comm artifact. Source Freeze Propagation - # bridging is excluded from across_ranks: its per-node gate folds a - # node's entire outgoing-edge set, which a single rank cannot - # observe for ghost owners, so bridging models fail fast on - # multi-rank instead. + # skips the dead dense with-comm artifact. Whether multi-rank is + # POSSIBLE at all is supports_edge_parallel (see + # test_capability_split_needs_vs_supports). dd = make_descriptor() assert dd.has_message_passing() is True assert dd.has_message_passing_across_ranks() is True assert dd.dense_lower_supports_comm() is False dd_bridge = make_descriptor(inner_clamp_r_inner=0.5, inner_clamp_r_outer=1.0) assert dd_bridge.has_message_passing() is True - assert dd_bridge.has_message_passing_across_ranks() is False + assert dd_bridge.has_message_passing_across_ranks() is True + + def test_capability_split_needs_vs_supports(self) -> None: + """has_message_passing_across_ranks = NEEDS exchange (always True for + SeZM); supports_edge_parallel = CAN run multi-rank (bridging vetoes, + until the SFPG exchange lands -- issue #5906). + """ + dd_plain = make_descriptor() + dd_bridged = make_descriptor(inner_clamp_r_inner=0.5, inner_clamp_r_outer=1.0) + assert dd_plain.has_message_passing_across_ranks() is True + assert dd_bridged.has_message_passing_across_ranks() is True + assert dd_plain.supports_edge_parallel() is True + assert dd_bridged.supports_edge_parallel() is False def test_serialize_roundtrip_exact(self) -> None: dd = make_descriptor() From 357fd272a0573c6348c4ff539cf587825fac64b5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 09:00:25 +0800 Subject: [PATCH 02/26] feat(dpmodel): promote export-time questions to atomic-model capabilities Six capabilities with concrete defaults on BaseAtomicModel, descriptor delegation on DPAtomicModel, and any/all aggregation on LinearEnergyAtomicModel (issue #5906 Task 4). --- .../dpmodel/atomic_model/base_atomic_model.py | 48 ++++ .../dpmodel/atomic_model/dp_atomic_model.py | 24 ++ .../atomic_model/linear_atomic_model.py | 25 ++ deepmd/dpmodel/descriptor/dpa1.py | 15 ++ deepmd/dpmodel/descriptor/dpa2.py | 1 + .../descriptor/make_base_descriptor.py | 26 ++ deepmd/pt_expt/descriptor/dpa1.py | 6 + deepmd/pt_expt/descriptor/dpa2.py | 4 + .../dpmodel/test_atomic_model_capabilities.py | 230 ++++++++++++++++++ 9 files changed, 379 insertions(+) create mode 100644 source/tests/common/dpmodel/test_atomic_model_capabilities.py diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index f2f2218443..b9e9532309 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -190,6 +190,54 @@ def supports_native_spin(self) -> bool: """ return False + def has_message_passing_across_ranks(self) -> bool: + """Whether multi-rank inference needs a cross-rank ghost exchange. + + Generic capability (concrete default ``False``): the export layer + consults it instead of reaching for a descriptor, so the answer + stays correct for descriptor-less models and compositions. + """ + return False + + def supports_edge_parallel(self) -> bool: + """Whether this atomic model can run under MPI domain decomposition. + + Default ``True``; a model folding state no single rank observes + (e.g. SFPG bridging before its exchange lands) overrides to False. + """ + return True + + def dense_lower_supports_comm(self) -> bool: + """Whether the DENSE (nlist) lower implements comm_dict exchange. + + Default ``True`` — dense comm is the production multi-rank path for + dpa2/dpa3; DPA4's dense adapter raises on comm_dict and overrides + via its descriptor. + """ + return True + + def uses_compact_edge_pairs(self) -> bool: + """Whether the graph lower emits compact ``center_edge_pairs`` + (drives the torch>=2.6 unbacked-SymInt export guard). + """ + return False + + def graph_edge_dtype(self) -> str: + """Edge-geometry dtype the graph deployment artifact accepts. + + ``"float64"`` is the model-agnostic ABI; geometrically compressed + float32 descriptors override to ``"float32"``. + """ + return "float64" + + def supports_graph_export(self) -> bool: + """Whether an exportable graph-lower implementation exists. + + A compressed descriptor without its fused opaque operator cannot be + traced through the reference tabulation kernel. + """ + return True + def get_default_fparam(self) -> list[float] | None: """Get the default frame parameters.""" return None diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 1752aabef9..c5cf0b157c 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -167,6 +167,30 @@ def uses_graph_lower(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.uses_graph_lower()) + def has_message_passing_across_ranks(self) -> bool: + """Delegates to this model's own descriptor.""" + return bool(self.descriptor.has_message_passing_across_ranks()) + + def supports_edge_parallel(self) -> bool: + """Delegates to this model's own descriptor.""" + return bool(self.descriptor.supports_edge_parallel()) + + def dense_lower_supports_comm(self) -> bool: + """Delegates to this model's own descriptor.""" + return bool(self.descriptor.dense_lower_supports_comm()) + + def uses_compact_edge_pairs(self) -> bool: + """Delegates to this model's own descriptor.""" + return bool(self.descriptor.uses_compact_edge_pairs()) + + def graph_edge_dtype(self) -> str: + """Delegates to this model's own descriptor.""" + return str(self.descriptor.graph_edge_dtype()) + + def supports_graph_export(self) -> bool: + """Delegates to this model's own descriptor.""" + return bool(self.descriptor.supports_graph_export()) + def supports_native_spin(self) -> bool: """Delegates to this model's own descriptor (cached at construction).""" return self._supports_native_spin diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 293ad005b5..4999bc15a7 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -182,6 +182,31 @@ def has_message_passing(self) -> bool: """Returns whether the atomic model has message passing.""" return any(model.has_message_passing() for model in self.models) + def has_message_passing_across_ranks(self) -> bool: + """ANY child needing the exchange makes the composition need it.""" + return any(m.has_message_passing_across_ranks() for m in self.models) + + def supports_edge_parallel(self) -> bool: + """EVERY child must tolerate decomposition; one veto vetoes all.""" + return all(m.supports_edge_parallel() for m in self.models) + + def dense_lower_supports_comm(self) -> bool: + """The shared dense lower is only comm-capable if every child's is.""" + return all(m.dense_lower_supports_comm() for m in self.models) + + def uses_compact_edge_pairs(self) -> bool: + """The export guard fires if ANY child emits compact pairs.""" + return any(m.uses_compact_edge_pairs() for m in self.models) + + def graph_edge_dtype(self) -> str: + """One shared edge tensor: float32 only if EVERY child accepts it.""" + dtypes = {m.graph_edge_dtype() for m in self.models} + return "float32" if dtypes == {"float32"} else "float64" + + def supports_graph_export(self) -> bool: + """All children trace into one artifact; each must be exportable.""" + return all(m.supports_graph_export() for m in self.models) + def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" return True diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index dcb499fe4e..a70e6d91e0 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -418,6 +418,21 @@ def has_message_passing_across_ranks(self) -> bool: """ return False + def graph_edge_dtype(self) -> str: + """float32 edges iff geometric compression runs float32 statistics. + + Compressed DPA1 evaluates both descriptor directions in the + statistics dtype and accepts float32 geometry directly. + """ + mean = getattr(self.se_atten, "mean", None) + if ( + self.geo_compress + and mean is not None + and str(mean.dtype).endswith("float32") + ): + return "float32" + return "float64" + def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" return self.se_atten.need_sorted_nlist_for_lower() diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index ccbfba085d..7ce2e82e48 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -603,6 +603,7 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.trainable = trainable self.add_tebd_to_repinit_out = add_tebd_to_repinit_out self.compress = False + self.geo_compress = False # graph-native lower opt-out flag (mirrors DescrptDPA1); not # serialized, re-derived structurally at construction/deserialization. self._graph_lower_disabled = False diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index 3e160782ed..2d058a3398 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -220,6 +220,32 @@ def uses_compact_edge_pairs(self) -> bool: """ return False + def dense_lower_supports_comm(self) -> bool: + """Whether the DENSE (nlist) lower implements comm_dict exchange. + + Default ``True`` — dense comm is the production multi-rank path + for dpa2/dpa3 (previously this was probed by method absence in + the freeze machinery); a descriptor whose dense adapter raises + on ``comm_dict`` (DPA4) overrides to ``False``. + """ + return True + + def graph_edge_dtype(self) -> str: + """Edge-geometry dtype the graph deployment artifact accepts. + + ``"float64"`` is the model-agnostic ABI; geometrically + compressed float32 descriptors override to ``"float32"``. + """ + return "float64" + + def supports_graph_export(self) -> bool: + """Whether an exportable graph-lower implementation exists. + + A compressed descriptor without its fused opaque operator cannot + be traced through the reference tabulation kernel. + """ + return True + def graph_type_embedding_table(self) -> Any | None: """Full type-embedding table consumed by the graph-route forward. diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index b6a7f9f84f..a2a179653a 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -577,6 +577,12 @@ def call_graph( comm_dict=comm_dict, ) + def supports_graph_export(self) -> bool: + """Compressed DPA1 must trace its fused opaque operator.""" + if not self.geo_compress: + return True + return self._fused_eligible("cuda") + def _fused_eligible(self, backend: str) -> bool: """Whether a fused descriptor kernel can serve this block. diff --git a/deepmd/pt_expt/descriptor/dpa2.py b/deepmd/pt_expt/descriptor/dpa2.py index a9943635e7..4cee82d6b5 100644 --- a/deepmd/pt_expt/descriptor/dpa2.py +++ b/deepmd/pt_expt/descriptor/dpa2.py @@ -49,6 +49,10 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: torch.zeros((), dtype=torch.bool, device="cpu"), ) + def supports_graph_export(self) -> bool: + """Compressed DPA2 has no fused graph operator yet.""" + return not self.geo_compress + def disable_graph_lower(self) -> None: """Persisted variant of the dpmodel escape hatch (see base class). diff --git a/source/tests/common/dpmodel/test_atomic_model_capabilities.py b/source/tests/common/dpmodel/test_atomic_model_capabilities.py new file mode 100644 index 0000000000..145e03bf9c --- /dev/null +++ b/source/tests/common/dpmodel/test_atomic_model_capabilities.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Issue #5906 Task 4: export-time questions are atomic-model capabilities. + +Compositions must answer by aggregation, never by wrapper type or by +reaching into a single ``.descriptor``. +""" + +import copy + +import pytest + +from deepmd.dpmodel.atomic_model import ( + DPAtomicModel, +) +from deepmd.dpmodel.atomic_model.inter_potential import ( + InterPotentialAtomicModel, +) +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel, +) +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.dpmodel.fitting import ( + InvarFitting, +) +from deepmd.dpmodel.model.model import ( + get_model, +) + +from .test_zbl_bridging import ( + ZBL_CONFIG, +) + +TYPE_MAP = ["Ni", "O"] + + +def _inter_potential() -> InterPotentialAtomicModel: + """Simplest concrete BaseAtomicModel: no descriptor at all.""" + return InterPotentialAtomicModel(type_map=TYPE_MAP, rcut=4.0, sel=[8]) + + +def _dp_atomic_model(descriptor) -> DPAtomicModel: + fitting = InvarFitting( + "energy", + len(TYPE_MAP), + descriptor.get_dim_out(), + 1, + mixed_types=descriptor.mixed_types(), + ) + return DPAtomicModel(descriptor, fitting, type_map=TYPE_MAP) + + +def _dpa4_descriptor(bridging: bool) -> DescrptDPA4: + kwargs = { + "ntypes": len(TYPE_MAP), + "sel": 8, + "rcut": 4.0, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 7, + "random_gamma": False, + } + if bridging: + kwargs.update(inner_clamp_r_inner=0.8, inner_clamp_r_outer=1.2) + return DescrptDPA4(**kwargs) + + +@pytest.mark.parametrize( + "cap,default", + [ + ("has_message_passing_across_ranks", False), # needs-exchange: opt-in + ("supports_edge_parallel", True), # nothing to veto by default + ("dense_lower_supports_comm", True), # dense comm is the norm (dpa2/dpa3) + ("uses_compact_edge_pairs", False), # torch>=2.6 guard: opt-in + ("supports_graph_export", True), # only compression restricts export + ], +) +def test_base_defaults(cap, default) -> None: + """BaseAtomicModel answers each capability with a concrete default.""" + model = _inter_potential() + assert getattr(model, cap)() is default + + +def test_base_graph_edge_dtype_default() -> None: + """float64 is the model-agnostic edge-geometry ABI.""" + assert _inter_potential().graph_edge_dtype() == "float64" + + +def test_dp_atomic_model_delegates_to_descriptor() -> None: + """DPAtomicModel answers every capability from its own descriptor.""" + bridged = _dp_atomic_model(_dpa4_descriptor(bridging=True)) + plain = _dp_atomic_model(_dpa4_descriptor(bridging=False)) + local_only = _dp_atomic_model(DescrptSeA(rcut=4.0, rcut_smth=3.5, sel=[8, 8])) + assert bridged.has_message_passing_across_ranks() is True + assert bridged.supports_edge_parallel() is False + assert plain.has_message_passing_across_ranks() is True + assert plain.supports_edge_parallel() is True + assert local_only.has_message_passing_across_ranks() is False + assert local_only.supports_edge_parallel() is True + # dense-lower comm and graph-export ride the same delegation + assert plain.dense_lower_supports_comm() is False # DPA4 dense adapter raises + assert local_only.dense_lower_supports_comm() is True + assert plain.graph_edge_dtype() == "float64" + assert plain.supports_graph_export() is True + assert plain.uses_compact_edge_pairs() is bool( + plain.descriptor.uses_compact_edge_pairs() + ) + + +class _EdgeParallelChild(InterPotentialAtomicModel): + """Stub child with settable capabilities (test_zbl_bridging.py pattern).""" + + def __init__( + self, + *args, + needs_exchange: bool = False, + edge_parallel: bool = True, + compact_pairs: bool = False, + edge_dtype: str = "float64", + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self._needs_exchange = needs_exchange + self._edge_parallel = edge_parallel + self._compact_pairs = compact_pairs + self._edge_dtype = edge_dtype + + def has_message_passing_across_ranks(self) -> bool: + return self._needs_exchange + + def supports_edge_parallel(self) -> bool: + return self._edge_parallel + + def uses_compact_edge_pairs(self) -> bool: + return self._compact_pairs + + def graph_edge_dtype(self) -> str: + return self._edge_dtype + + +def _stub_linear(**child_kwargs_pair) -> LinearEnergyAtomicModel: + """Two stub children; child_kwargs_pair maps kwarg -> (child0, child1).""" + kwargs0 = {k: v[0] for k, v in child_kwargs_pair.items()} + kwargs1 = {k: v[1] for k, v in child_kwargs_pair.items()} + return LinearEnergyAtomicModel( + [ + _EdgeParallelChild(type_map=TYPE_MAP, rcut=4.0, sel=[8], **kwargs0), + _EdgeParallelChild(type_map=TYPE_MAP, rcut=4.0, sel=[8], **kwargs1), + ], + type_map=TYPE_MAP, + weights="sum", + ) + + +def test_linear_aggregation_any_all() -> None: + """Real ZBL composition: the bridged DP child sets needs=True (any) and + vetoes edge-parallel (all). + """ + am = get_model(copy.deepcopy(ZBL_CONFIG)).atomic_model + assert isinstance(am, LinearEnergyAtomicModel) + assert am.has_message_passing_across_ranks() is True + assert am.supports_edge_parallel() is False + # ZBL rides the graph route with the learned child + assert am.supports_graph_export() is True + assert am.graph_edge_dtype() == "float64" + + +def test_linear_aggregation_mixed_children() -> None: + """Both boolean branches of every any/all rule via stub children.""" + # has_message_passing_across_ranks: ANY + assert ( + _stub_linear(needs_exchange=(True, False)).has_message_passing_across_ranks() + is True + ) + assert ( + _stub_linear(needs_exchange=(False, False)).has_message_passing_across_ranks() + is False + ) + # supports_edge_parallel: ALL (one veto vetoes all) + assert _stub_linear(edge_parallel=(True, False)).supports_edge_parallel() is False + assert _stub_linear(edge_parallel=(True, True)).supports_edge_parallel() is True + # uses_compact_edge_pairs: ANY + assert _stub_linear(compact_pairs=(True, False)).uses_compact_edge_pairs() is True + assert _stub_linear(compact_pairs=(False, False)).uses_compact_edge_pairs() is False + # graph_edge_dtype: float32 iff ALL children float32 + assert ( + _stub_linear(edge_dtype=("float32", "float64")).graph_edge_dtype() == "float64" + ) + assert ( + _stub_linear(edge_dtype=("float32", "float32")).graph_edge_dtype() == "float32" + ) + + +def test_linear_aggregation_dense_comm_and_graph_export() -> None: + """dense_lower_supports_comm and supports_graph_export are ALL rules.""" + + class _NoDenseComm(InterPotentialAtomicModel): + def dense_lower_supports_comm(self) -> bool: + return False + + class _NoGraphExport(InterPotentialAtomicModel): + def supports_graph_export(self) -> bool: + return False + + plain = _inter_potential() + mixed_comm = LinearEnergyAtomicModel( + [_NoDenseComm(type_map=TYPE_MAP, rcut=4.0, sel=[8]), _inter_potential()], + type_map=TYPE_MAP, + weights="sum", + ) + assert mixed_comm.dense_lower_supports_comm() is False + all_comm = LinearEnergyAtomicModel( + [plain, _inter_potential()], type_map=TYPE_MAP, weights="sum" + ) + assert all_comm.dense_lower_supports_comm() is True + mixed_export = LinearEnergyAtomicModel( + [_NoGraphExport(type_map=TYPE_MAP, rcut=4.0, sel=[8]), _inter_potential()], + type_map=TYPE_MAP, + weights="sum", + ) + assert mixed_export.supports_graph_export() is False + assert all_comm.supports_graph_export() is True From 7c256d5ae3ec505d1655cdfb661c5797c5d1ec2e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 09:11:19 +0800 Subject: [PATCH 03/26] refactor(pt_expt): hoist graph with-comm exportable into make_model LinearEnergyModel compositions need it once the with-comm gate opens for them (issue #5906 Task 4); one owner, next to the non-comm twin. _translate_energy_keys moves along (make_model cannot import from ener_model without a cycle); importers re-pointed. --- deepmd/pt_expt/model/ener_model.py | 221 +---------------------------- deepmd/pt_expt/model/make_model.py | 220 ++++++++++++++++++++++++++++ deepmd/pt_expt/train/training.py | 2 +- 3 files changed, 222 insertions(+), 221 deletions(-) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index b5b650ab40..77b6a84bac 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -23,6 +23,7 @@ ) from .make_model import ( + _translate_energy_keys, make_model, ) from .model import ( @@ -32,41 +33,6 @@ DPEnergyModel_ = make_model(DPEnergyAtomicModel, T_Bases=(BaseModel,)) -def _translate_energy_keys( - model_ret: dict[str, torch.Tensor], - *, - do_grad_r: bool, - do_grad_c: bool, - do_atomic_virial: bool, - local: bool, -) -> dict[str, torch.Tensor]: - """Map internal fitting keys -> public energy-model keys (shared by the - dense and graph ``forward_lower`` export traces). - - Operates on plain dicts (make_fx-safe). ``local=True`` is the GRAPH path - (per-node ``N == sum(n_node)`` local atoms, no ghost/extended region) and - emits ``force``/``atom_virial``; ``local=False`` is the DENSE extended-region - path and emits ``extended_force``/``extended_virial`` (folded to local by - ``communicate_extended_output`` at inference). - """ - out: dict[str, torch.Tensor] = {} - out["atom_energy"] = model_ret["energy"] - out["energy"] = model_ret["energy_redu"] - if do_grad_r: - out["force" if local else "extended_force"] = model_ret[ - "energy_derv_r" - ].squeeze(-2) - if do_grad_c: - out["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial: - out["atom_virial" if local else "extended_virial"] = model_ret[ - "energy_derv_c" - ].squeeze(-2) - if "mask" in model_ret: - out["mask"] = model_ret["mask"] - return out - - @BaseModel.register("ener") class EnergyModel(DPModelCommon, DPEnergyModel_): def __init__( @@ -584,188 +550,3 @@ def fn( aparam, charge_spin, ) - - def forward_lower_graph_exportable_with_comm( - self, - atype: torch.Tensor, - n_node: torch.Tensor, - n_local: torch.Tensor, - edge_index: torch.Tensor, - edge_vec: torch.Tensor, - edge_mask: torch.Tensor, - destination_order: torch.Tensor, - destination_row_ptr: torch.Tensor, - source_order: torch.Tensor, - source_row_ptr: torch.Tensor, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None, - send_list: torch.Tensor, - send_proc: torch.Tensor, - recv_proc: torch.Tensor, - send_num: torch.Tensor, - recv_num: torch.Tensor, - communicator: torch.Tensor, - nlocal: torch.Tensor, - nghost: torch.Tensor, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - """Trace ``forward_common_lower_graph`` with comm_dict tensors as - additional positional inputs -- the with-comm counterpart of - :meth:`forward_lower_graph_exportable` for message-passing graph - descriptors (dpa2's repformer block drives cross-rank ghost refresh - via ``deepmd_export::border_op``, see - :meth:`~deepmd.pt_expt.descriptor.repformers. - DescrptBlockRepformers._exchange_ghosts_graph`). - - Mirrors the dense ``forward_common_lower_exportable_with_comm`` - (``pt_expt/model/make_model.py``): packs the 8 trailing positional - comm tensors into a ``comm_dict`` inside the traced function. Also - derives ``n_local`` (the per-frame OWNED node count, reshaped to - ``(1,)``; single-frame -- LAMMPS always drives inference with - ``nf=1``) from the scalar ``nlocal`` tensor, so the differentiated - reduction excludes ghost (not-owned) nodes (see - :meth:`forward_common_lower_graph`'s ``n_local`` parameter). Unlike - the plain-graph export path (which traces - ``forward_common_lower_graph_exportable`` and then wraps a SECOND - make_fx trace around the key-translation closure), this method - traces ONCE: the comm-dict packing, ``n_local`` derivation, the - ``forward_common_lower_graph`` call and the key translation all live - in a single traced ``fn`` -- following the dense with-comm - precedent, which is also a single trace. - - Parameters - ---------- - atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin, do_atomic_virial - As in :meth:`forward_lower_graph_exportable`. - send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost - The 8 comm tensors (see ``_make_comm_sample_inputs`` in - ``serialization.py``), packed into ``comm_dict`` inside the - traced function. - - Runtime device contract: ALL 8 stay on CPU, symmetric with - the dense with-comm artifact -- they are consumed only by the - opaque ``border_op`` whose HOST code dereferences their - ``data_ptr`` (``send_list`` carries raw host pointers) and - reads ``nlocal``/``nghost`` via cheap host ``.item()`` calls. - Deriving the in-graph owned count from a device-placed - ``nlocal`` instead (the previous design) made every per-layer - ``border_op`` forward AND custom backward pull the scalars - back with synchronizing D2H reads (``4 * nlayers`` per MD - step). The C++ ``run_model_graph_with_comm`` implements this - placement. - n_local - (1,) int64 ON THE MODEL DEVICE: the per-frame OWNED node - count consumed IN-GRAPH by the owned-node energy mask (it - becomes a device kernel operand after - ``move_to_device_pass``, like ``n_node``; a CPU tensor fed - there is read as a device pointer -- CUDA illegal memory - access). Carries the same value as the ``nlocal`` comm - tensor; the two inputs exist precisely to separate the - device-compute role from the host-MPI-control role. - **make_fx_kwargs - Extra keyword arguments forwarded to ``make_fx`` - (e.g. ``tracing_mode="symbolic"``). - - Returns - ------- - torch.nn.Module - A traced module whose ``forward`` accepts ``(atype, n_node, - n_local, edge_index, edge_vec, edge_mask, destination_order, - destination_row_ptr, source_order, source_row_ptr, fparam, - aparam, charge_spin, send_list, send_proc, recv_proc, send_num, - recv_num, communicator, nlocal, nghost)`` and returns a dict with the - SAME public keys as :meth:`forward_lower_graph_exportable` - (``atom_energy``, ``energy``, ``force``, ``virial``, - ``atom_virial`` when ``do_atomic_virial``). - """ - model = self - do_grad_r = self.do_grad_r("energy") - do_grad_c = self.do_grad_c("energy") - - def fn( - atype: torch.Tensor, - n_node: torch.Tensor, - n_local: torch.Tensor, - edge_index: torch.Tensor, - edge_vec: torch.Tensor, - edge_mask: torch.Tensor, - destination_order: torch.Tensor, - destination_row_ptr: torch.Tensor, - source_order: torch.Tensor, - source_row_ptr: torch.Tensor, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None, - send_list: torch.Tensor, - send_proc: torch.Tensor, - recv_proc: torch.Tensor, - send_num: torch.Tensor, - recv_num: torch.Tensor, - communicator: torch.Tensor, - nlocal: torch.Tensor, - nghost: torch.Tensor, - ) -> dict[str, torch.Tensor]: - comm_dict = { - "send_list": send_list, - "send_proc": send_proc, - "recv_proc": recv_proc, - "send_num": send_num, - "recv_num": recv_num, - "communicator": communicator, - "nlocal": nlocal, - "nghost": nghost, - } - # ``n_local`` (slot 2, DEVICE) is the owned-count input consumed - # by the in-graph owned-node mask; the CPU ``nlocal`` comm - # tensor is host control metadata for border_op only. - model_ret = model.forward_common_lower_graph( - atype, - n_node, - n_local, - edge_index, - edge_vec, - edge_mask, - destination_order, - destination_row_ptr, - source_order, - source_row_ptr, - destination_sorted=True, - do_atomic_virial=do_atomic_virial, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - comm_dict=comm_dict, - ) - return _translate_energy_keys( - model_ret, - do_grad_r=do_grad_r, - do_grad_c=do_grad_c, - do_atomic_virial=do_atomic_virial, - local=True, - ) - - return make_fx(fn, **make_fx_kwargs)( - atype, - n_node, - n_local, - edge_index, - edge_vec, - edge_mask, - destination_order, - destination_row_ptr, - source_order, - source_row_ptr, - fparam, - aparam, - charge_spin, - send_list, - send_proc, - recv_proc, - send_num, - recv_num, - communicator, - nlocal, - nghost, - ) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 9d906c0874..32ac4e7dc5 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -38,6 +38,41 @@ ) +def _translate_energy_keys( + model_ret: dict[str, torch.Tensor], + *, + do_grad_r: bool, + do_grad_c: bool, + do_atomic_virial: bool, + local: bool, +) -> dict[str, torch.Tensor]: + """Map internal fitting keys -> public energy-model keys (shared by the + dense and graph ``forward_lower`` export traces). + + Operates on plain dicts (make_fx-safe). ``local=True`` is the GRAPH path + (per-node ``N == sum(n_node)`` local atoms, no ghost/extended region) and + emits ``force``/``atom_virial``; ``local=False`` is the DENSE extended-region + path and emits ``extended_force``/``extended_virial`` (folded to local by + ``communicate_extended_output`` at inference). + """ + out: dict[str, torch.Tensor] = {} + out["atom_energy"] = model_ret["energy"] + out["energy"] = model_ret["energy_redu"] + if do_grad_r: + out["force" if local else "extended_force"] = model_ret[ + "energy_derv_r" + ].squeeze(-2) + if do_grad_c: + out["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial: + out["atom_virial" if local else "extended_virial"] = model_ret[ + "energy_derv_c" + ].squeeze(-2) + if "mask" in model_ret: + out["mask"] = model_ret["mask"] + return out + + def _fused_energy_force_graph( model: Any, graph: Any, @@ -1206,6 +1241,191 @@ def fn_spin( spin, ) + def forward_lower_graph_exportable_with_comm( + self, + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + send_list: torch.Tensor, + send_proc: torch.Tensor, + recv_proc: torch.Tensor, + send_num: torch.Tensor, + recv_num: torch.Tensor, + communicator: torch.Tensor, + nlocal: torch.Tensor, + nghost: torch.Tensor, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + """Trace ``forward_common_lower_graph`` with comm_dict tensors as + additional positional inputs -- the with-comm counterpart of + :meth:`forward_lower_graph_exportable` for message-passing graph + descriptors (dpa2's repformer block drives cross-rank ghost refresh + via ``deepmd_export::border_op``, see + :meth:`~deepmd.pt_expt.descriptor.repformers. + DescrptBlockRepformers._exchange_ghosts_graph`). + + Mirrors the dense ``forward_common_lower_exportable_with_comm`` + (``pt_expt/model/make_model.py``): packs the 8 trailing positional + comm tensors into a ``comm_dict`` inside the traced function. Also + derives ``n_local`` (the per-frame OWNED node count, reshaped to + ``(1,)``; single-frame -- LAMMPS always drives inference with + ``nf=1``) from the scalar ``nlocal`` tensor, so the differentiated + reduction excludes ghost (not-owned) nodes (see + :meth:`forward_common_lower_graph`'s ``n_local`` parameter). Unlike + the plain-graph export path (which traces + ``forward_common_lower_graph_exportable`` and then wraps a SECOND + make_fx trace around the key-translation closure), this method + traces ONCE: the comm-dict packing, ``n_local`` derivation, the + ``forward_common_lower_graph`` call and the key translation all live + in a single traced ``fn`` -- following the dense with-comm + precedent, which is also a single trace. + + Parameters + ---------- + atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin, do_atomic_virial + As in :meth:`forward_lower_graph_exportable`. + send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost + The 8 comm tensors (see ``_make_comm_sample_inputs`` in + ``serialization.py``), packed into ``comm_dict`` inside the + traced function. + + Runtime device contract: ALL 8 stay on CPU, symmetric with + the dense with-comm artifact -- they are consumed only by the + opaque ``border_op`` whose HOST code dereferences their + ``data_ptr`` (``send_list`` carries raw host pointers) and + reads ``nlocal``/``nghost`` via cheap host ``.item()`` calls. + Deriving the in-graph owned count from a device-placed + ``nlocal`` instead (the previous design) made every per-layer + ``border_op`` forward AND custom backward pull the scalars + back with synchronizing D2H reads (``4 * nlayers`` per MD + step). The C++ ``run_model_graph_with_comm`` implements this + placement. + n_local + (1,) int64 ON THE MODEL DEVICE: the per-frame OWNED node + count consumed IN-GRAPH by the owned-node energy mask (it + becomes a device kernel operand after + ``move_to_device_pass``, like ``n_node``; a CPU tensor fed + there is read as a device pointer -- CUDA illegal memory + access). Carries the same value as the ``nlocal`` comm + tensor; the two inputs exist precisely to separate the + device-compute role from the host-MPI-control role. + **make_fx_kwargs + Extra keyword arguments forwarded to ``make_fx`` + (e.g. ``tracing_mode="symbolic"``). + + Returns + ------- + torch.nn.Module + A traced module whose ``forward`` accepts ``(atype, n_node, + n_local, edge_index, edge_vec, edge_mask, destination_order, + destination_row_ptr, source_order, source_row_ptr, fparam, + aparam, charge_spin, send_list, send_proc, recv_proc, send_num, + recv_num, communicator, nlocal, nghost)`` and returns a dict with the + SAME public keys as :meth:`forward_lower_graph_exportable` + (``atom_energy``, ``energy``, ``force``, ``virial``, + ``atom_virial`` when ``do_atomic_virial``). + """ + model = self + do_grad_r = self.do_grad_r("energy") + do_grad_c = self.do_grad_c("energy") + + def fn( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + send_list: torch.Tensor, + send_proc: torch.Tensor, + recv_proc: torch.Tensor, + send_num: torch.Tensor, + recv_num: torch.Tensor, + communicator: torch.Tensor, + nlocal: torch.Tensor, + nghost: torch.Tensor, + ) -> dict[str, torch.Tensor]: + comm_dict = { + "send_list": send_list, + "send_proc": send_proc, + "recv_proc": recv_proc, + "send_num": send_num, + "recv_num": recv_num, + "communicator": communicator, + "nlocal": nlocal, + "nghost": nghost, + } + # ``n_local`` (slot 2, DEVICE) is the owned-count input consumed + # by the in-graph owned-node mask; the CPU ``nlocal`` comm + # tensor is host control metadata for border_op only. + model_ret = model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=True, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + comm_dict=comm_dict, + ) + return _translate_energy_keys( + model_ret, + do_grad_r=do_grad_r, + do_grad_c=do_grad_c, + do_atomic_virial=do_atomic_virial, + local=True, + ) + + return make_fx(fn, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + send_list, + send_proc, + recv_proc, + send_num, + recv_num, + communicator, + nlocal, + nghost, + ) + def forward_common_lower_exportable_with_comm( self, extended_coord: torch.Tensor, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 1129696ed1..b0d63d45e3 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -645,7 +645,7 @@ def _trace_and_compile_graph( make_fx, ) - from deepmd.pt_expt.model.ener_model import ( + from deepmd.pt_expt.model.make_model import ( _translate_energy_keys, ) From e674c301a55ecf7e32c976249545e5e39f936e9e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 09:46:00 +0800 Subject: [PATCH 04/26] refactor(pt_expt): serialization consumes atomic-model capabilities _needs_with_comm_artifact / compact-edge-pairs / edge-dtype / graph-export answer via the atomic model; two-DPA2 linear compositions now get their with-comm artifact (issue #5906 Task 4). --- deepmd/pt_expt/utils/serialization.py | 91 ++++++------------- .../pt_expt/model/test_export_with_comm.py | 84 +++++++++++++++++ .../pt_expt/utils/test_graph_pt2_metadata.py | 84 +++++++++++++++-- 3 files changed, 188 insertions(+), 71 deletions(-) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 897d91cec3..3bdc9d676d 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -134,19 +134,19 @@ def _needs_with_comm_artifact( ``use_loc_mapping=True`` keep all per-layer messaging local to each rank's owned atoms; they need only the regular artifact. - Delegates to ``descriptor.has_message_passing_across_ranks()``, which - descriptor classes implement explicitly. Returns ``False`` defensively - when the model has no single descriptor (linear/zbl/frozen) or when - the method is somehow missing or raises. + Capabilities are answered by the atomic model + (``has_message_passing_across_ranks`` / ``supports_edge_parallel`` / + ``dense_lower_supports_comm``); compositions aggregate over children, + so linear/zbl models answer for themselves instead of being denied by + wrapper type (issue #5906 Task 4). Not every lower path that needs cross-rank exchange implements it: DPA4's graph lower carries a real per-layer ``border_op`` exchange, but its dense (nlist) lower's adapter raises on ``comm_dict``. ``lower_kind`` selects which lower is being traced so the gate can consult the per-lower capability instead of assuming both lowers agree. Non-graph - kinds additionally check ``descriptor.dense_lower_supports_comm()`` - (absent on descriptors, such as dpa2/dpa3, whose dense lower always - supports comm — treated as ``True``). + kinds additionally check ``dense_lower_supports_comm()`` (``True`` for + dpa2/dpa3, whose dense lower is the production multi-rank path). Native spin participates on the GRAPH lower, matching pt's ``SeZMModel.supports_edge_parallel`` (which ``SeZMNativeSpinModel`` does @@ -181,34 +181,17 @@ def _needs_with_comm_artifact( if isinstance(model, NativeSpinModelKind) and lower_kind != "graph": return False - # Analytical bridging models are single-rank only (pt's - # ``supports_edge_parallel() == False`` contract: ZBL + SFPG fold each - # node's full outgoing-edge set, which a single rank cannot observe for - # ghost owners) -- never compile a with-comm artifact for them. - from deepmd.dpmodel.atomic_model.linear_atomic_model import ( - LinearEnergyAtomicModel, - ) - - atomic_model = getattr(model, "atomic_model", None) - if isinstance(atomic_model, LinearEnergyAtomicModel): - # Compositions (e.g. analytical bridging: learned + InterPotential) - # are single-rank on the graph route: per-edge analytical terms fold - # each node's full edge set, which a single rank cannot observe for - # ghost owners (pt's supports_edge_parallel()==False rationale). - return False - - desc = getattr(getattr(model, "atomic_model", None), "descriptor", None) - if desc is None or not desc.has_message_passing_across_ranks(): + atomic_model = model.atomic_model + if not ( + atomic_model.has_message_passing_across_ranks() + and atomic_model.supports_edge_parallel() + ): return False if lower_kind == "graph": return True - # Non-graph kinds trace the DENSE with-comm wrapper; a descriptor whose - # dense lower has no comm implementation (DPA4: the dense adapter raises - # on comm_dict) must not emit a dead or untraceable dense artifact. - # Descriptors without the method (dpa2/dpa3/...) implement dense comm — - # it is their production multi-rank path. - dense_ok = getattr(desc, "dense_lower_supports_comm", None) - return True if dense_ok is None else bool(dense_ok()) + # Non-graph kinds trace the DENSE with-comm wrapper; a model whose dense + # lower has no comm implementation (DPA4) must not emit a dead artifact. + return bool(atomic_model.dense_lower_supports_comm()) def check_graph_trace_torch_version(model: torch.nn.Module) -> None: @@ -235,9 +218,10 @@ def check_graph_trace_torch_version(model: torch.nn.Module) -> None: ``attn_layer > 0`` and for dpa2 with ``update_g2_has_attn`` or ``update_h2`` (keying on ``get_numb_attn_layer()`` alone missed dpa2, whose repformer attention rides ``center_edge_pairs`` without - implementing that dpa1 accessor). Models without a single - descriptor (linear/zbl/frozen) pass the check (they take the dense - route anyway), as do descriptors without the capability method. + implementing that dpa1 accessor). Compositions answer by + aggregation (ANY child emitting compact pairs trips the guard); + since graph-capable models auto-resolve onto the graph route, + linear/zbl compositions are NOT exempt. Raises ------ @@ -245,8 +229,7 @@ def check_graph_trace_torch_version(model: torch.nn.Module) -> None: If the descriptor's graph lower traces compact edge pairs and the running torch is older than 2.6. """ - desc = getattr(getattr(model, "atomic_model", None), "descriptor", None) - if desc is None or not desc.uses_compact_edge_pairs(): + if not model.atomic_model.uses_compact_edge_pairs(): return version = torch.__version__.split("+")[0] major_minor = tuple(int(p) for p in version.split(".")[:2] if p.isdigit()) @@ -1018,37 +1001,17 @@ def _build_dynamic_shapes( def _graph_edge_dtype(model: torch.nn.Module, lower_kind: str) -> str: """Return the graph edge-vector dtype encoded by the deployment artifact. - Geometrically compressed DPA1 with float32 descriptor statistics evaluates - both descriptor directions in float32 and therefore accepts float32 - geometry directly. Other graph descriptors retain the model-agnostic - float64 geometry ABI. + The dtype itself is the atomic model's capability; only the lower-kind + applicability (which artifact kinds carry edge geometry) lives here. """ - atomic_model = getattr(model, "atomic_model", None) - descriptor = getattr(atomic_model, "descriptor", None) - descriptor_block = getattr(descriptor, "se_atten", None) - statistics = getattr(descriptor_block, "mean", None) - if ( - lower_kind in ("graph", "dpa1_canonical") - and bool(getattr(descriptor, "geo_compress", False)) - and isinstance(statistics, torch.Tensor) - and statistics.dtype == torch.float32 - ): - return "float32" - return "float64" + if lower_kind not in ("graph", "dpa1_canonical"): + return "float64" + return str(model.atomic_model.graph_edge_dtype()) def _supports_graph_export(model: torch.nn.Module) -> bool: - """Whether the model has an exportable graph-lower implementation. - - A compressed descriptor must use its opaque graph operator during export; - tracing through the reference tabulation kernel is unsupported. - """ - atomic_model = getattr(model, "atomic_model", None) - descriptor = getattr(atomic_model, "descriptor", None) - if not bool(getattr(descriptor, "geo_compress", False)): - return True - eligible = getattr(descriptor, "_fused_eligible", None) - return callable(eligible) and bool(eligible("cuda")) + """Whether the model has an exportable graph-lower implementation.""" + return bool(model.atomic_model.supports_graph_export()) def _collect_metadata( diff --git a/source/tests/pt_expt/model/test_export_with_comm.py b/source/tests/pt_expt/model/test_export_with_comm.py index f34b20932f..8ddcbaadcc 100644 --- a/source/tests/pt_expt/model/test_export_with_comm.py +++ b/source/tests/pt_expt/model/test_export_with_comm.py @@ -360,3 +360,87 @@ def test_pte_with_comm_dict_traces_and_loads(tmp_path) -> None: f"with-comm exported program must accept 15 positional inputs " f"(7 base + 8 comm); got {len(spec)}" ) + + +_DPA2_CHILD_CONFIG = { + "descriptor": { + "type": "dpa2", + "repinit": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": 10, + "neuron": [4, 8], + "axis_neuron": 2, + }, + "repformer": { + "rcut": 3.0, + "rcut_smth": 0.5, + "nsel": 6, + "nlayers": 1, + "g1_dim": 8, + "g2_dim": 4, + }, + }, + "fitting_net": {"neuron": [8, 8], "seed": 1}, +} + + +def _build_two_dpa2_linear_model() -> torch.nn.Module: + """Linear composition of two DPA2 children (OutisLi, #5884 review).""" + import copy + + from deepmd.pt_expt.model.get_model import ( + get_linear_model, + ) + + config = { + "type_map": ["O", "H"], + "models": [ + copy.deepcopy(_DPA2_CHILD_CONFIG), + copy.deepcopy(_DPA2_CHILD_CONFIG), + ], + "weights": "mean", + } + model = get_linear_model(config) + model.to("cpu") + model.eval() + return model + + +def test_two_dpa2_linear_composition_gets_with_comm_artifact() -> None: + """Issue #5906 Task 4 (OutisLi, #5884 review): a graph-eligible + linear_ener of two DPA2 children -- both needing cross-rank message + passing -- must request the with-comm artifact. Previously denied by + BOTH the isinstance(LinearEnergyAtomicModel) check and the + ``.descriptor is None`` fallthrough. + """ + from deepmd.pt_expt.utils.serialization import ( + _needs_with_comm_artifact, + ) + + model = _build_two_dpa2_linear_model() + assert _needs_with_comm_artifact(model, lower_kind="graph") is True + # dpa2's dense lower supports comm, so the nlist kind agrees. + assert _needs_with_comm_artifact(model, lower_kind="nlist") is True + + +def test_bridged_composition_still_denied_until_sfpg_exchange() -> None: + """Negative contract: bridged DPA4+ZBL stays single-rank in Phase 1 -- + the veto now comes from supports_edge_parallel aggregation, not from + the wrapper's type. + """ + import copy + + from deepmd.pt_expt.model.get_model import get_model as get_pt_expt_model + from deepmd.pt_expt.utils.serialization import ( + _needs_with_comm_artifact, + ) + + from .test_zbl_bridging import ( + ZBL_CONFIG, + ) + + model = get_pt_expt_model(copy.deepcopy(ZBL_CONFIG)) + model.to("cpu") + model.eval() + assert _needs_with_comm_artifact(model, lower_kind="graph") is False diff --git a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py index bcd0b3b0c8..bf05e5339a 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -137,7 +137,16 @@ def test_graph_pt2_small_sel_exports() -> None: def test_compressed_graph_uses_compute_precision_edge_geometry( statistics_dtype: torch.dtype, expected: str ) -> None: - """Compressed DPA1 graph geometry follows descriptor compute precision.""" + """Compressed DPA1 graph geometry follows descriptor compute precision. + + The stub borrows the REAL capability methods (dpmodel + ``DescrptDPA1.graph_edge_dtype``, pt_expt + ``DescrptDPA1.supports_graph_export``) and the real DPAtomicModel-style + delegation, so the helpers are tested through the capability seam they + consume in production (issue #5906 Task 4). + """ + from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as _DPDescrptDPA1 + from deepmd.pt_expt.descriptor.dpa1 import DescrptDPA1 as _PEDescrptDPA1 class _Descriptor: geo_compress = True @@ -150,9 +159,18 @@ class _Block: def _fused_eligible(self, backend: str) -> bool: return backend == "cuda" and self.se_atten.mean.dtype == torch.float32 + graph_edge_dtype = _DPDescrptDPA1.graph_edge_dtype + supports_graph_export = _PEDescrptDPA1.supports_graph_export + class _AtomicModel: descriptor = _Descriptor() + def graph_edge_dtype(self) -> str: + return str(self.descriptor.graph_edge_dtype()) + + def supports_graph_export(self) -> bool: + return bool(self.descriptor.supports_graph_export()) + class _Model: atomic_model = _AtomicModel() @@ -214,6 +232,10 @@ class _FakeAtomicModel: def __init__(self, n_attn: int) -> None: self.descriptor = _FakeDesc(n_attn) + def uses_compact_edge_pairs(self) -> bool: + # mirrors DPAtomicModel's capability delegation (issue #5906 Task 4) + return self.descriptor.uses_compact_edge_pairs() + class _FakeModel: def __init__(self, n_attn: int) -> None: @@ -258,19 +280,36 @@ def test_graph_trace_version_guard_passes(monkeypatch, version, n_attn) -> None: check_graph_trace_torch_version(_FakeModel(n_attn)) -def test_graph_trace_version_guard_tolerates_no_descriptor(monkeypatch) -> None: - """Composite models without a single descriptor pass (dense route anyway).""" +def test_graph_trace_version_guard_checks_compositions(monkeypatch) -> None: + """Compositions answer by aggregation and are NOT exempt (issue #5906). + + The old defensive fallthrough silently passed any model without a + single ``.descriptor``; a linear composition whose child emits compact + pairs must now trip the torch < 2.6 guard like the child itself would. + """ import torch from deepmd.pt_expt.utils.serialization import ( check_graph_trace_torch_version, ) - class _NoDesc: - pass + class _FakeLinearAtomicModel: + def __init__(self, children) -> None: + self.models = children + + def uses_compact_edge_pairs(self) -> bool: + return any(m.uses_compact_edge_pairs() for m in self.models) + + class _FakeLinearModel: + def __init__(self, n_attns) -> None: + self.atomic_model = _FakeLinearAtomicModel( + [_FakeAtomicModel(n) for n in n_attns] + ) monkeypatch.setattr(torch, "__version__", "2.5.1") - check_graph_trace_torch_version(_NoDesc()) + with pytest.raises(RuntimeError, match=r"torch >= 2\.6"): + check_graph_trace_torch_version(_FakeLinearModel([0, 2])) + check_graph_trace_torch_version(_FakeLinearModel([0, 0])) @pytest.mark.parametrize( @@ -365,7 +404,10 @@ def _build_model(model_kind: str) -> torch.nn.Module: ---------- model_kind : str ``"dpa4"`` (bridging-free SeZM, config shared with - ``test_dpa4_export.py``) or ``"dpa2"`` (``DPA2_GUARD_CONFIG`` above). + ``test_dpa4_export.py``), ``"dpa2"`` (``DPA2_GUARD_CONFIG`` above), + or ``"linear-two-dpa2"`` (a linear composition of two + ``DPA2_GUARD_CONFIG`` children -- capability aggregation, + issue #5906 Task 4). Returns ------- @@ -382,6 +424,24 @@ def _build_model(model_kind: str) -> torch.nn.Module: config = _DPA4_CONFIG elif model_kind == "dpa2": config = DPA2_GUARD_CONFIG + elif model_kind == "linear-two-dpa2": + from deepmd.pt_expt.model.get_model import ( + get_linear_model, + ) + + child = { + "descriptor": DPA2_GUARD_CONFIG["descriptor"], + "fitting_net": DPA2_GUARD_CONFIG["fitting_net"], + } + config = { + "type_map": DPA2_GUARD_CONFIG["type_map"], + "models": [copy.deepcopy(child), copy.deepcopy(child)], + "weights": "mean", + } + model = get_linear_model(config) + model.to("cpu") + model.eval() + return model else: raise ValueError(f"unknown model_kind {model_kind!r}") model = get_pt_expt_model(copy.deepcopy(config)) @@ -405,6 +465,16 @@ def _build_model(model_kind: str) -> torch.nn.Module: True, ), # dense with-comm is dpa2's production MP path — unchanged ("dpa2", "graph", True), # graph with-comm unchanged + ( + "linear-two-dpa2", + "graph", + True, + ), # composition aggregates children (issue #5906 Task 4) + ( + "linear-two-dpa2", + "nlist", + True, + ), # dpa2 children's dense lower supports comm -> composition does ], ) def test_needs_with_comm_artifact_kind_aware(model_kind, lower_kind, expected) -> None: From bd56fb37e466d5e4d5fe44821d86b6669cc13ed5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 10:02:40 +0800 Subject: [PATCH 05/26] fix(pt_expt): composition-safe ntypes/compile/bridging config paths .pt-checkpoint eval and enable_compile no longer reach for a single descriptor; get_standard_model honors bridging_method like its dpmodel twin, with _compose_bridging as the one composition owner (issue #5906 Task 4 audit findings). --- deepmd/pt_expt/infer/deep_eval.py | 4 +- deepmd/pt_expt/model/get_model.py | 118 ++++++++++++---- deepmd/pt_expt/model/native_spin_model.py | 10 +- deepmd/pt_expt/train/training.py | 66 +++++---- .../pt_expt/model/test_get_model_bridging.py | 129 ++++++++++++++++++ 5 files changed, 267 insertions(+), 60 deletions(-) create mode 100644 source/tests/pt_expt/model/test_get_model_bridging.py diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 999e7f69f1..10cc10b0a6 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -621,7 +621,9 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: # `_collect_metadata` writes into metadata.json. self.metadata = { "type_map": model.get_type_map(), - "ntypes": model.get_descriptor().get_ntypes(), + # via the model API, not the descriptor: compositions + # (LinearEnergyAtomicModel) have no single descriptor to reach for + "ntypes": len(model.get_type_map()), "rcut": model.get_rcut(), "sel": model.get_sel(), "dim_fparam": model.get_dim_fparam(), diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 6e779dd23e..3cc82a5327 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -8,6 +8,9 @@ import copy import logging +from typing import ( + Any, +) from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, @@ -65,7 +68,6 @@ pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, ) -get_standard_model = _model_factory.get_standard_model get_zbl_model = _model_factory.get_zbl_model @@ -164,38 +166,93 @@ def get_sezm_model(data: dict) -> EnergyModel: pair_exclude_types=pair_exclude_types, ) if bridging_enabled: - # Composition, not a flag (first-principles design): the analytical - # bridging term is its own atomic model, summed with the learned one - # by the existing linear composition machinery. - from deepmd.dpmodel.atomic_model.inter_potential import ( - InterPotentialAtomicModel, - ) - from deepmd.dpmodel.atomic_model.linear_atomic_model import ( - LinearEnergyAtomicModel, - ) - from deepmd.pt_expt.model.dp_linear_model import ( - LinearEnergyModel, - ) - - zbl_atomic = InterPotentialAtomicModel( - type_map=data["type_map"], - mode=bridging_method, - rcut=descriptor.get_rcut(), - sel=descriptor.get_sel(), - ) - composed = LinearEnergyAtomicModel( - models=[model.atomic_model, zbl_atomic], - type_map=data["type_map"], - weights="sum", - # Both exclusions belong to the composition: its children share one - # graph, so "excluded" must cover the analytical term too. - atom_exclude_types=data.get("atom_exclude_types", []), - pair_exclude_types=pair_exclude_types, - ) - return LinearEnergyModel(atomic_model_=composed) + return _compose_bridging(model, data, bridging_method) return model +def _compose_bridging(model: Any, data: dict, bridging_method: str) -> Any: + """Compose the learned model with its analytical bridging term. + + Composition, not a flag (first-principles design): the analytical + bridging term is its own atomic model, summed with the learned one by + the existing linear composition machinery. The ONE owner of the + composition build for this backend -- both :func:`get_sezm_model` + (``type: "dpa4"``) and :func:`get_standard_model` (``type: + "standard"``) route through here, mirroring the dpmodel twin + (``deepmd/dpmodel/model/model.py``). + + Parameters + ---------- + model + The learned backbone model (its descriptor already carries the + bridging radii injected by the caller). + data + The model config (``type_map`` and the exclusion lists are read). + bridging_method + The analytical bridging mode (e.g. ``"ZBL"``). + + Returns + ------- + Any + A :class:`LinearEnergyModel` over ``[learned, InterPotential]``. + """ + from deepmd.dpmodel.atomic_model.inter_potential import ( + InterPotentialAtomicModel, + ) + from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel, + ) + from deepmd.pt_expt.model.dp_linear_model import ( + LinearEnergyModel, + ) + + descriptor = model.atomic_model.descriptor + zbl_atomic = InterPotentialAtomicModel( + type_map=data["type_map"], + mode=bridging_method, + rcut=descriptor.get_rcut(), + sel=descriptor.get_sel(), + ) + composed = LinearEnergyAtomicModel( + models=[model.atomic_model, zbl_atomic], + type_map=data["type_map"], + weights="sum", + # Both exclusions belong to the composition: its children share one + # graph, so "excluded" must cover the analytical term too. + atom_exclude_types=data.get("atom_exclude_types", []), + pair_exclude_types=data.get("pair_exclude_types", []), + ) + return LinearEnergyModel(atomic_model_=composed) + + +def get_standard_model(data: dict) -> Any: + """Build a pt_expt standard model, honoring ``bridging_method``. + + pt_expt twin of :func:`deepmd.dpmodel.model.model.get_standard_model`: + the analytical-bridging radii feed the DESCRIPTOR's + InnerClamp/BridgingSwitch and the method composes the atomic model with + its InterPotential term. Without this wrapper a ``type: "standard"`` + config with ``bridging_method`` silently dropped the bridging term + (backend divergence from dpmodel -- issue #5906 Task 4 audit). + + Parameters + ---------- + data : dict + The data to construct the model. + """ + data = copy.deepcopy(data) + bridging_method = str(data.get("bridging_method", "none")) + bridging_enabled = bridging_method.lower() not in ("none", "") + if bridging_enabled: + data.setdefault("descriptor", {}) + data["descriptor"]["inner_clamp_r_inner"] = data.get("bridging_r_inner", 0.5) + data["descriptor"]["inner_clamp_r_outer"] = data.get("bridging_r_outer", 0.8) + model = _model_factory.get_standard_model(data) + if not bridging_enabled: + return model + return _compose_bridging(model, data, bridging_method) + + def get_native_spin_model(data: dict) -> NativeSpinEnergyModel: """Build a pt_expt native (virtual-atom-free) spin model. @@ -335,6 +392,7 @@ def get_model(data: dict) -> BaseModel: """ return _model_factory.get_model( data, + standard_model_factory=get_standard_model, spin_model_factory=get_spin_model, native_spin_model_factory=get_native_spin_model, model_factories={ diff --git a/deepmd/pt_expt/model/native_spin_model.py b/deepmd/pt_expt/model/native_spin_model.py index 0d63e7e530..6b3a87b49e 100644 --- a/deepmd/pt_expt/model/native_spin_model.py +++ b/deepmd/pt_expt/model/native_spin_model.py @@ -72,12 +72,10 @@ def forward( do_atomic_virial If calculate the atomic virial. charge_spin - Frame-level charge/spin conditioning, shape nf x 2. Accepted for - call-signature compatibility with ``ModelWrapper.forward`` (which - always forwards this keyword); charge-spin FiLM combined with - native spin is rejected at construction time - (``add_chg_spin_ebd`` on the descriptor), so this is always - ``None`` in practice for a model this class can build. + Frame-level charge/spin conditioning, shape nf x 2. Forwarded to + the backbone like every other conditioning input; native spin and + charge-spin FiLM combine freely (DPA4 declares both + capabilities). Returns ------- diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index b0d63d45e3..b8c332413a 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -161,6 +161,48 @@ def _detect_task_buffers( return result +def _warn_compiled_attention(model: torch.nn.Module, task_key: str) -> None: + """Warn when compiling DPA1/se_atten_v2 attention (hardware-sensitive). + + Compiled DPA1/se_atten_v2 attention is numerically more sensitive than + other descriptors: the inductor-fused and eager force/grad outputs can + diverge above 1e-10 on multi-threaded CPU hosts because parallel + reduction order is hardware-dependent. Warn but do not reject — + energies remain well within training tolerance and the user may accept + the trade-off for compile speed. + + Compositions (``LinearEnergyAtomicModel``) have no single descriptor to + probe; ``enable_compile`` must degrade gracefully for them instead of + crashing on the reach-in (issue #5906 Task 4 audit). + + Parameters + ---------- + model + The per-task model about to be compiled. + task_key + The task label used in the warning message. + """ + from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP + + try: + descriptor = model.get_descriptor() + except AttributeError: + return + if isinstance(descriptor, DescrptDPA1DP): + n_attn = descriptor.get_numb_attn_layer() + if n_attn > 0: + log.warning( + "Compiling DPA1/se_atten_v2 with %d attention " + "layer(s) (task=%s): the compiled forces/grads " + "are slightly hardware-sensitive (multi-thread " + "reduction order), and may not match the eager " + "path bit-for-bit. Use 'enable_compile: false' " + "or 'attn_layer: 0' for fully reproducible runs.", + n_attn, + task_key, + ) + + def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: """Return a key that is identical iff two tasks can safely share a compiled graph. @@ -1971,8 +2013,6 @@ def _compile_model(self, compile_opts: dict[str, Any]) -> None: defaultdict, ) - from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP - # Pre-pass: group tasks by structure key and auto-detect per-task buffers. # Grouping is needed so _detect_task_buffers can diff buffer identities # across all tasks that share the same compiled graph. @@ -2021,27 +2061,7 @@ def _compile_model(self, compile_opts: dict[str, Any]) -> None: for task_key in self.model_keys: model = wrapper_mod.model[task_key] - # Compiled DPA1/se_atten_v2 attention is numerically more - # sensitive than other descriptors: the inductor-fused and - # eager force/grad outputs can diverge above 1e-10 on - # multi-threaded CPU hosts because parallel reduction order - # is hardware-dependent. Warn but do not reject — energies - # remain well within training tolerance and the user may - # accept the trade-off for compile speed. - descriptor = model.get_descriptor() - if isinstance(descriptor, DescrptDPA1DP): - n_attn = descriptor.get_numb_attn_layer() - if n_attn > 0: - log.warning( - "Compiling DPA1/se_atten_v2 with %d attention " - "layer(s) (task=%s): the compiled forces/grads " - "are slightly hardware-sensitive (multi-thread " - "reduction order), and may not match the eager " - "path bit-for-bit. Use 'enable_compile: false' " - "or 'attn_layer: 0' for fully reproducible runs.", - n_attn, - task_key, - ) + _warn_compiled_attention(model, task_key) structure_key = _key_for[task_key] task_bufs = _task_bufs_for[task_key] diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py new file mode 100644 index 0000000000..e8ab81e920 --- /dev/null +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""pt_expt ``get_standard_model`` must honor ``bridging_method`` like its +dpmodel twin (``deepmd/dpmodel/model/model.py``) -- issue #5906 Task 4 +variant-alignment audit, gap 2: a ``type: "standard"`` config with bridging +silently dropped the InterPotential composition in pt_expt. +""" + +import copy + +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel, +) +from deepmd.pt_expt.model.get_model import ( + get_model, + get_standard_model, +) + + +def _dpa4_standard_config() -> dict: + """The ZBL config of test_zbl_bridging.py minus the 'dpa4' model type.""" + return { + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4", + "rcut": 4.0, + "sel": 8, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 7, + "random_gamma": False, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [8, 8], + "precision": "float64", + "seed": 7, + }, + } + + +def test_standard_model_type_builds_bridging_composition() -> None: + """pt_expt twin of dpmodel model.py's get_standard_model: a config with + bridging_method must compose [learned, InterPotential], not silently + drop the bridging term. + """ + data = _dpa4_standard_config() + data["bridging_method"] = "ZBL" + data["bridging_r_inner"] = 0.8 + data["bridging_r_outer"] = 1.2 + model = get_standard_model(copy.deepcopy(data)) + assert isinstance(model.atomic_model, LinearEnergyAtomicModel) + assert len(model.atomic_model.models) == 2 + # The descriptor radii injection must ride the same seam (a composition + # without the inner-clamp radii would be a half-applied bridging config): + desc = model.atomic_model.models[0].descriptor + assert desc.bridging_switch is not None + # And the get_model router (type omitted -> "standard") reaches the same + # composition: + routed = get_model(copy.deepcopy(data)) + assert isinstance(routed.atomic_model, LinearEnergyAtomicModel) + + +def test_standard_model_type_maps_dpa4_fitting() -> None: + """Dpmodel maps dpa4_ener/sezm_ener fitting under type:'standard' via the + model registry; pt_expt must not raise where dpmodel builds. + """ + model = get_standard_model(_dpa4_standard_config()) + assert model is not None + assert model.atomic_model.descriptor.bridging_switch is None + + +def test_pt_checkpoint_eval_works_for_composition(tmp_path) -> None: + """``DeepEval`` on a ``.pt`` checkpoint of a bridging composition. + + ``LinearEnergyAtomicModel`` has no single ``.descriptor``; + ``_load_pt``'s metadata block previously crashed with + ``AttributeError`` on ``model.get_descriptor()`` (issue #5906 Task 4 + audit, gap 1). ``ntypes`` now comes from the model API. + """ + import numpy as np + import torch + + from deepmd.infer import ( + DeepPot, + ) + from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, + ) + + config = _dpa4_standard_config() + config["type"] = "dpa4" + config["bridging_method"] = "ZBL" + config["bridging_r_inner"] = 0.8 + config["bridging_r_outer"] = 1.2 + model = get_model(copy.deepcopy(config)).to(torch.float64).eval() + ckpt = str(tmp_path / "dpa4_zbl.pt") + wrapper = ModelWrapper(model, model_params=copy.deepcopy(config)) + torch.save({"model": wrapper.state_dict()}, ckpt) + + dp = DeepPot(ckpt) + assert dp.get_ntypes() == 2 + rng = np.random.default_rng(5) + coord = rng.uniform(1.5, 5.5, size=(1, 6, 3)) + coord[0, 1] = coord[0, 0] + np.array([0.9, 0.0, 0.0]) + atype = np.array([[0, 0, 1, 0, 1, 1]], dtype=np.int64) + box = 8.0 * np.eye(3, dtype=np.float64).reshape(1, 9) + e, f, v = dp.eval(coord.reshape(1, -1), box, atype[0].tolist()) + assert np.all(np.isfinite(e)) + assert np.all(np.isfinite(f)) + + +def test_compile_attention_probe_tolerates_composition() -> None: + """``enable_compile``'s DPA1-attention warning probe must degrade + gracefully for compositions instead of crashing on + ``model.get_descriptor()`` (issue #5906 Task 4 audit, gap 1 twin). + """ + from deepmd.pt_expt.train.training import ( + _warn_compiled_attention, + ) + + data = _dpa4_standard_config() + data["bridging_method"] = "ZBL" + model = get_standard_model(data) + # must not raise + _warn_compiled_attention(model, "Default") From 4eeacf525d11ad4bd02489e89512c5b68b0b4e6d Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 10:07:55 +0800 Subject: [PATCH 06/26] feat(pt_expt): autograd for border_op_backward (transpose of border_op) R = B^T as linear maps over node rows, so the reverse-accumulate's vector-Jacobian product is the forward broadcast (issue #5906: gate gradients cross ranks). --- deepmd/pt_expt/utils/comm.py | 78 +++++++++++++++++++ .../pt_expt/utils/test_border_op_backward.py | 67 ++++++++++++++++ 2 files changed, 145 insertions(+) diff --git a/deepmd/pt_expt/utils/comm.py b/deepmd/pt_expt/utils/comm.py index af2be287ff..a0e70dec35 100644 --- a/deepmd/pt_expt/utils/comm.py +++ b/deepmd/pt_expt/utils/comm.py @@ -190,6 +190,79 @@ def _border_op_backward( ) +def _border_op_backward_grad_setup_context( + ctx: torch.autograd.function.FunctionCtx, + inputs: tuple, + output: torch.Tensor, +) -> None: + ( + sendlist, + sendproc, + recvproc, + sendnum, + recvnum, + _grad_g1, + communicator, + nlocal, + nghost, + ) = inputs + ctx.save_for_backward( + sendlist, + sendproc, + recvproc, + sendnum, + recvnum, + communicator, + nlocal, + nghost, + ) + + +def _border_op_backward_grad( + ctx: torch.autograd.function.FunctionCtx, + grad_output: torch.Tensor, +) -> tuple: + """Gradient of the reverse-accumulate: the forward broadcast. + + ``border_op_backward`` is the exact transpose of ``border_op`` as a + linear map over node rows, so its vector-Jacobian product IS + ``border_op``'s forward walk (issue #5906: this closes the autograd + loop for the SFPG completion so gate gradients cross ranks). + """ + ( + sendlist, + sendproc, + recvproc, + sendnum, + recvnum, + communicator, + nlocal, + nghost, + ) = ctx.saved_tensors + grad_in = torch.ops.deepmd_export.border_op( + sendlist, + sendproc, + recvproc, + sendnum, + recvnum, + grad_output.contiguous(), + communicator, + nlocal, + nghost, + ) + return ( + None, + None, + None, + None, + None, # sendlist..recvnum + grad_in, # grad_g1 + None, + None, + None, # communicator, nlocal, nghost + ) + + def ensure_comm_registered() -> None: """Load libdeepmd_op_pt.so and register fake/autograd metadata for border_op. @@ -223,4 +296,9 @@ def ensure_comm_registered() -> None: _border_op_backward, setup_context=_border_op_setup_context, ) + torch.library.register_autograd( + "deepmd_export::border_op_backward", + _border_op_backward_grad, + setup_context=_border_op_backward_grad_setup_context, + ) _registered = True diff --git a/source/tests/pt_expt/utils/test_border_op_backward.py b/source/tests/pt_expt/utils/test_border_op_backward.py index 07a5ac67ab..2d89b8527c 100644 --- a/source/tests/pt_expt/utils/test_border_op_backward.py +++ b/source/tests/pt_expt/utils/test_border_op_backward.py @@ -256,3 +256,70 @@ def test_border_op_export_autograd(dtype: torch.dtype) -> None: atol=atol, rtol=rtol, ) + + +# --------------------------------------------------------------------------- +# 3. deepmd_export::border_op_backward as a differentiable op (issue #5906): +# its gradient is border_op's forward (the two are exact transposes). +# --------------------------------------------------------------------------- + + +def test_export_backward_op_semantics_self_comm() -> None: + """R then B on a self-comm dict: owner rows hold own+ghost sums, ghost + rows hold the completed owner value (the SFPG completion sequence). + """ + from deepmd.pt_expt.utils.comm import ( + ensure_comm_registered, + ) + + ensure_comm_registered() + nloc, nghost = 3, 2 + keepalive: list = [] + # ghosts [3,4] are images of owners [0,1] + comm = _build_self_swap( + nloc, + nghost, + np.array([0, 1], dtype=np.int32), + keepalive, + torch.float64, + ) + x = torch.arange(10, dtype=torch.float64).reshape(5, 2).contiguous() + # the underlying op accumulates INTO its input buffer -- snapshot first + x_orig = x.clone() + r = torch.ops.deepmd_export.border_op_backward(*comm[:5], x.clone(), *comm[5:]) + expected_owner = x_orig[:3].clone() + expected_owner[0] += x_orig[3] + expected_owner[1] += x_orig[4] + torch.testing.assert_close(r[:3], expected_owner, rtol=0, atol=0) + torch.testing.assert_close( + r[3:], torch.zeros(2, 2, dtype=torch.float64), rtol=0, atol=0 + ) + b = torch.ops.deepmd_export.border_op(*comm[:5], r, *comm[5:]) + torch.testing.assert_close(b[3], b[0], rtol=0, atol=0) + torch.testing.assert_close(b[4], b[1], rtol=0, atol=0) + + +def test_export_backward_op_gradient_is_border_op() -> None: + """d/dx of border_op_backward(x) applied to a cotangent v equals + border_op(v) -- the ops are transposes of each other. + """ + from deepmd.pt_expt.utils.comm import ( + ensure_comm_registered, + ) + + ensure_comm_registered() + nloc, nghost = 3, 2 + keepalive: list = [] + comm = _build_self_swap( + nloc, + nghost, + np.array([0, 1], dtype=np.int32), + keepalive, + torch.float64, + ) + x = torch.randn(5, 2, dtype=torch.float64, requires_grad=True) + y = torch.ops.deepmd_export.border_op_backward(*comm[:5], x, *comm[5:]) + v = torch.randn_like(y) + (grad,) = torch.autograd.grad(y, x, v) + expected = torch.ops.deepmd_export.border_op(*comm[:5], v.contiguous(), *comm[5:]) + torch.testing.assert_close(grad, expected, rtol=0, atol=0) From a0c24223b15d9c9e4d29e8885ecffac7f75e5045 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 10:19:11 +0800 Subject: [PATCH 07/26] feat(dpmodel): SFPG gate accepts a cross-rank partial-completion hook compute_edge_src_gate packs (N,2) [log_eta, zero_count] through the optional node_partial_exchange hook; zero_count becomes float so both partials ride one border exchange. dpmodel's _gate_partial_exchange raises (single-process reference); pt_expt overrides it (issue #5906). --- deepmd/dpmodel/descriptor/dpa4.py | 44 +++++++++++++++++++ .../dpmodel/descriptor/dpa4_nn/edge_cache.py | 44 ++++++++++++++++--- .../tests/common/dpmodel/test_descrpt_dpa4.py | 39 ++++++++++++++++ 3 files changed, 120 insertions(+), 7 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 95f43e09ca..acb3306228 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -35,6 +35,7 @@ annotations, ) +import functools import math from typing import ( TYPE_CHECKING, @@ -1541,6 +1542,14 @@ def _run_graph( type_ebed, spin, atype_flat, n_nodes=n_nodes ) + # Cross-rank SFPG completion (issue #5906): only a bridged model + # under domain decomposition needs it -- the gate's src-keyed + # per-node partials are rank-incomplete then. + node_partial_exchange = None + if comm_dict is not None and self.bridging_switch is not None: + node_partial_exchange = functools.partial( + self._gate_partial_exchange, comm_dict=comm_dict + ) # === Step 3. Build edge cache once (sparse edges) === edge_cache = _edge_cache_from_arrays( type_ebed=type_ebed, @@ -1562,6 +1571,7 @@ def _run_graph( random_gamma=self.random_gamma and self._in_training_mode(), wigner_calc=self.wigner_calc, build_wigner=self._need_full_wigner, + node_partial_exchange=node_partial_exchange, ) ebed_dim_0 = self.node_init_dim # (node_init_lmax+1)^2 @@ -2196,6 +2206,40 @@ def _canonicalize_charge_spin( raise ValueError("`charge_spin` first dimension must match nframes.") return charge_spin + def _gate_partial_exchange( + self, + partials: Array, + comm_dict: dict[str, Array], + ) -> Array: + """Complete the SFPG per-node partials across ranks. + + Reverse-accumulate ghost rows into their owners, then broadcast the + completed owner values back — the backend-specific pt_expt subclass + implements it on ``border_op_backward``/``border_op``; dpmodel is + the single-process reference and rejects comm outright. + + Parameters + ---------- + partials + (n_nodes, 2) float tensor of [log_eta, zero_count] partials. + comm_dict + The border-exchange control tensors. + + Returns + ------- + Array + The globally completed (n_nodes, 2) tensor. + + Raises + ------ + NotImplementedError + Always, in the dpmodel backend. + """ + raise NotImplementedError( + "Multi-rank SFPG partial exchange (comm_dict) is not supported " + "in the dpmodel backend." + ) + def _block_comm( self, block_idx: int, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py index 99f20c4285..f924dc5802 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py @@ -131,6 +131,7 @@ def compute_edge_src_gate( n_nodes: int, bridging_switch: Callable[[Any], Any], edge_keep_f: Any = None, + node_partial_exchange: Callable[[Any], Any] | None = None, ) -> Any: """ Compute the per-edge source gate for SFPG from edge lengths. @@ -182,6 +183,14 @@ def compute_edge_src_gate( Optional per-edge keep weights with shape (E, 1), with ``0`` on masked edges and ``1`` on kept edges. If provided, masked edges are rewritten to ``w = 1`` before the product reduction. + node_partial_exchange + Optional cross-rank completion hook for the per-node partials + (issue #5906). Receives the ``(n_nodes, 2)`` float tensor + ``[log_eta, zero_count]`` and returns the globally completed one + (reverse-accumulate ghost rows into owners, then broadcast the + completed owner values back onto ghosts). ``None`` (the default) + is the single-process path where the local partials are already + complete. Returns ------- @@ -209,18 +218,33 @@ def compute_edge_src_gate( log_eta = xp_add_at( xp.zeros((n_nodes,), dtype=edge_w.dtype, device=device), src, log_safe ) - eta_nonzero_path = xp.exp(log_eta) # === Step 3. Exact-zero indicator per source node === - # ``scatter_add`` over an ``int64`` cast of the zero mask counts how - # many frozen edges each source node owns. A strictly positive count - # means the product is 0 by the hard-freeze rule. + # ``scatter_add`` over the zero mask counts how many frozen edges each + # source node owns. A strictly positive count means the product is 0 by + # the hard-freeze rule. Float count (values are small integers, exact + # in fp) so both partials ride ONE border-exchange tensor when + # completing across ranks. zero_count = xp_add_at( - xp.zeros((n_nodes,), dtype=xp.int64, device=device), + xp.zeros((n_nodes,), dtype=edge_w.dtype, device=device), src, - xp.astype(is_zero, xp.int64), + xp.astype(is_zero, edge_w.dtype), ) - any_zero = zero_count > 0 + + # === Step 3b. Cross-rank completion of the per-node partials === + # A rank only holds edges whose dst is owned, so the src-keyed sums + # above are PARTIAL for every node under domain decomposition. The hook + # (reverse-accumulate ghost->owner, then forward-broadcast owner->ghost) + # completes them; log-products are additive and each edge lives on + # exactly one rank, so nothing double-counts (issue #5906). + if node_partial_exchange is not None: + packed = xp.stack([log_eta, zero_count], axis=-1) # (n_nodes, 2) + packed = node_partial_exchange(packed) + log_eta = packed[..., 0] + zero_count = packed[..., 1] + + eta_nonzero_path = xp.exp(log_eta) + any_zero = zero_count > 0.5 # === Step 4. Combine and broadcast back to edges via source === eta = xp.where(any_zero, xp.zeros_like(eta_nonzero_path), eta_nonzero_path) @@ -244,6 +268,7 @@ def _edge_cache_from_arrays( wigner_calc: WignerCalculatorFn, build_wigner: bool = True, gamma: Any = None, + node_partial_exchange: Callable[[Any], Any] | None = None, ) -> EdgeCache: """ Build the global edge cache from a sparse edge list. @@ -295,6 +320,10 @@ def _edge_cache_from_arrays( ``random_gamma`` is True. When None, drawn with the backend's RNG (:func:`~deepmd.dpmodel.array_api.xp_uniform`) uniformly in ``[0, 2*pi)``; callers may inject angles to pin a draw. + node_partial_exchange + Optional cross-rank completion hook forwarded to + :func:`compute_edge_src_gate` (issue #5906); only meaningful when + ``bridging_switch`` is provided. Returns ------- @@ -358,6 +387,7 @@ def _edge_cache_from_arrays( n_nodes=n_nodes, bridging_switch=bridging_switch, edge_keep_f=edge_keep_f, + node_partial_exchange=node_partial_exchange, ) return _finalize_edge_cache( diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py index 51e717124a..4294026004 100644 --- a/source/tests/common/dpmodel/test_descrpt_dpa4.py +++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py @@ -122,6 +122,45 @@ def test_capability_split_needs_vs_supports(self) -> None: assert dd_plain.supports_edge_parallel() is True assert dd_bridged.supports_edge_parallel() is False + def test_gate_partial_exchange_dpmodel_raises(self) -> None: + """The dpmodel backend is the single-process reference; comm on a + bridged model must raise, never silently compute a partial gate. + """ + dd = make_descriptor(inner_clamp_r_inner=0.5, inner_clamp_r_outer=1.0) + with pytest.raises(NotImplementedError, match="dpmodel"): + dd._gate_partial_exchange(np.zeros((4, 2)), {"nlocal": 2}) + + def test_edge_src_gate_identity_exchange_is_noop(self) -> None: + """The hook seam: an identity exchange reproduces the no-hook gate + bit-exactly (pins the pack/unpack layout [log_eta, zero_count]). + """ + from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( + compute_edge_src_gate, + ) + + dd = make_descriptor(inner_clamp_r_inner=0.5, inner_clamp_r_outer=1.0) + sw = dd.bridging_switch + # 4 nodes, 5 edges; one edge inside r_inner (w=0, hard-freezes its + # src node), the rest spread across the transition zone and beyond. + el = np.array([[0.3], [0.7], [0.9], [1.5], [0.75]], dtype=np.float64) + src = np.array([0, 0, 1, 2, 3], dtype=np.int64) + gate_ref = compute_edge_src_gate( + edge_len=el, src=src, n_nodes=4, bridging_switch=sw + ) + gate_hook = compute_edge_src_gate( + edge_len=el, + src=src, + n_nodes=4, + bridging_switch=sw, + node_partial_exchange=lambda p: p, + ) + np.testing.assert_array_equal(gate_ref, gate_hook) + # geometry sanity: node 0 is hard-frozen, node 3 is inside the + # transition zone (0 < gate < 1) -- the test data actually + # exercises both gate branches. + assert gate_ref[0, 0] == 0.0 + assert 0.0 < gate_ref[4, 0] < 1.0 + def test_serialize_roundtrip_exact(self) -> None: dd = make_descriptor() data = dd.serialize() From 78cbe68c962ef3d3204020611fc428414932ffc8 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 10:26:35 +0800 Subject: [PATCH 08/26] feat(pt_expt): complete the SFPG gate across ranks via border-op pair _gate_partial_exchange = reverse-accumulate (border_op_backward) then forward-broadcast (border_op) of the (N,2) [log_eta, zero_count] partials. Eager self-comm parity vs the folded reference at 1e-12 with a cross-boundary close pair in both bridging channels (hard-freeze and transition zone); identity-exchange ablation diverges (issue #5906). --- deepmd/pt_expt/descriptor/dpa4.py | 43 +++ .../pt_expt/model/test_dpa4_zbl_parallel.py | 253 ++++++++++++++++++ 2 files changed, 296 insertions(+) create mode 100644 source/tests/pt_expt/model/test_dpa4_zbl_parallel.py diff --git a/deepmd/pt_expt/descriptor/dpa4.py b/deepmd/pt_expt/descriptor/dpa4.py index 5060de1a55..de9e4e3824 100644 --- a/deepmd/pt_expt/descriptor/dpa4.py +++ b/deepmd/pt_expt/descriptor/dpa4.py @@ -226,6 +226,49 @@ def _in_training_mode(self) -> bool: """ return bool(self.training) + def _gate_partial_exchange( + self, + partials: torch.Tensor, + comm_dict: dict[str, torch.Tensor], + ) -> torch.Tensor: + """Reverse-accumulate ghost partials to owners, then broadcast back. + + ``border_op_backward`` sums each ghost row into its owner across + ranks and zeroes the ghost rows; ``border_op`` refills them with the + completed owner values. Both ops carry autograd (the two are + transposes), so gate gradients cross ranks (issue #5906). + + Parameters + ---------- + partials + (n_nodes, 2) float tensor of [log_eta, zero_count] partials. + comm_dict + The border-exchange control tensors. + + Returns + ------- + torch.Tensor + The globally completed (n_nodes, 2) tensor. + """ + # border_op exchanges rows by raw pointer arithmetic; a strided + # view would corrupt the exchange. + p = partials.contiguous() + comm_args = ( + comm_dict["send_list"], + comm_dict["send_proc"], + comm_dict["recv_proc"], + comm_dict["send_num"], + comm_dict["recv_num"], + ) + tail = ( + comm_dict["communicator"], + comm_dict["nlocal"], + comm_dict["nghost"], + ) + p = torch.ops.deepmd_export.border_op_backward(*comm_args, p, *tail) + p = torch.ops.deepmd_export.border_op(*comm_args, p, *tail) + return p + def disable_graph_lower(self) -> None: """Persisted variant of the dpmodel escape hatch (see base class). diff --git a/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py b/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py new file mode 100644 index 0000000000..450f8de720 --- /dev/null +++ b/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Issue #5906 Task 2: SFPG completion across ranks, eager self-comm rung. + +Reference = the folded single-rank graph forward (periodic carry-all graph, +src folded onto owners, gate complete by construction). Candidate = the +UNFOLDED extended-layout forward (owners + ghost images, edges restricted to +owned destinations -- the LAMMPS layout) with a 1-swap self-send comm_dict. + +Geometry is the load-bearing part: the cell places a close pair ACROSS the +periodic x boundary so the close contact's src is a ghost image row and its +``log w``/``zero_count`` partial lives on a different row than the owner -- +without that, every edge contributes ``log w = 0`` and the exchange is +untestable (the issue's vacuous-test warning). Parametrized over a +hard-freeze pair (r < r_inner, exercises the ``zero_count`` channel) and a +transition-zone pair (r_inner < r < r_outer, exercises ``log_eta``). +""" + +import copy +import ctypes + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, +) +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.pt_expt.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.utils.comm import ( + ensure_comm_registered, +) + +from ...dpa4_fixtures import ( + jitter_zero_arrays, +) +from .test_zbl_bridging import ( + ZBL_CONFIG, +) + +# bridging window of ZBL_CONFIG: r_inner=0.8, r_outer=1.2 (Angstrom) +_L = 6.0 # cubic box edge; > rcut so self-images never bond + + +def _close_pair_coords(gap: float) -> np.ndarray: + """4 atoms; atoms 0/1 straddle the x periodic boundary at ``gap``. + + Atom 0 sits at ``x = gap/2`` and atom 1 at ``x = L - gap/2``, so their + minimum-image distance is exactly ``gap``. Atoms 2/3 give the + descriptor a normal environment (their 1.118 A contact sits inside the + transition zone but does NOT straddle the boundary -- its partial is + rank-local and pins the no-exchange-needed case alongside). + """ + half = gap / 2.0 + return np.array( + [ + [half, 3.0, 3.0], + [_L - half, 3.0, 3.0], + [2.5, 3.0, 3.0], + [3.6, 3.2, 3.0], + ], + dtype=np.float64, + ).reshape(1, 4, 3) + + +def _addr_of(np_arr: np.ndarray) -> int: + return np_arr.ctypes.data_as(ctypes.c_void_p).value + + +def _build_self_comm_dict( + *, + nloc: int, + nghost: int, + sendlist_indices: np.ndarray, + keepalive: list, +) -> dict: + """Single-rank self-exchange comm_dict (per-file copy, repo precedent).""" + sendlist_indices = np.ascontiguousarray(sendlist_indices, dtype=np.int32) + keepalive.append(sendlist_indices) + addr = _addr_of(sendlist_indices) + return { + "send_list": torch.tensor([addr], dtype=torch.int64, device="cpu"), + "send_proc": torch.zeros(1, dtype=torch.int32, device="cpu"), + "recv_proc": torch.zeros(1, dtype=torch.int32, device="cpu"), + "send_num": torch.tensor([nghost], dtype=torch.int32, device="cpu"), + "recv_num": torch.tensor([nghost], dtype=torch.int32, device="cpu"), + "communicator": torch.zeros(1, dtype=torch.int64, device="cpu"), + "nlocal": torch.tensor(nloc, dtype=torch.int32, device="cpu"), + "nghost": torch.tensor(nghost, dtype=torch.int32, device="cpu"), + } + + +def _make_bridged_model(): + """Bridged ZBL model with jittered zero-init residuals. + + A fresh DPA4 is architecturally edge-independent (zero-init residual + projections -- see ``jitter_zero_arrays``), so an un-jittered model + makes BOTH the parity and the ablation vacuous: the SFPG gate would + multiply edge messages that never reach the output. + """ + model = get_model(copy.deepcopy(ZBL_CONFIG)) + learned = model.atomic_model.models[0] + data = jitter_zero_arrays(learned.descriptor.serialize(), np.random.default_rng(99)) + learned.descriptor = DescrptDPA4.deserialize(data) + return model.to(torch.float64).to("cpu").eval() + + +def _extended_quartet(coord: np.ndarray): + """LAMMPS-layout quartet: ghosts materialized, nlist over owned dst.""" + atype = np.array([[0, 0, 1, 1]], dtype=np.int64) + box = (_L * np.eye(3, dtype=np.float64)).reshape(1, 3, 3) + rcut = 4.0 + ext_coord, ext_atype, mapping = extend_coord_with_ghosts( + coord, atype, box.reshape(1, 9), rcut + ) + # sel=128 is deliberately non-binding: this dense nlist only supplies + # the owned-dst edge list; truncation would silently desynchronize the + # two routes' neighbor sets. + nlist = build_neighbor_list( + ext_coord, + ext_atype, + 4, + rcut, + [128], + distinguish_types=False, + ) + return ( + np.asarray(ext_coord).reshape(1, -1, 3), + np.asarray(ext_atype), + np.asarray(nlist), + np.asarray(mapping), + ) + + +def _unfolded_graph_inputs(ext_coord, ext_atype, nlist): + """Unfolded graph over owners+ghosts, edges for OWNED dst only.""" + nloc = nlist.shape[1] + nall = ext_coord.shape[1] + idx = nlist[0] # (nloc, nnei), extended indices, -1 padding + valid = idx >= 0 + src = idx[valid].astype(np.int64) + dst = np.repeat(np.arange(nloc, dtype=np.int64), idx.shape[1])[valid.ravel()] + edge_vec = ext_coord[0, src] - ext_coord[0, dst] + return { + "atype": torch.tensor(ext_atype[0], dtype=torch.int64), + "n_node": torch.tensor([nall], dtype=torch.int64), + "n_local": torch.tensor([nloc], dtype=torch.int64), + "edge_index": torch.tensor(np.stack([src, dst]), dtype=torch.int64), + "edge_vec": torch.tensor(edge_vec, dtype=torch.float64), + "edge_mask": torch.ones(len(src), dtype=torch.bool), + } + + +def _fold_forces(per_node_force: np.ndarray, mapping: np.ndarray) -> np.ndarray: + """Sum ghost-image force rows onto their owners (LAMMPS reverse comm).""" + nloc = 4 + out = np.zeros((nloc, 3), dtype=np.float64) + np.add.at(out, mapping[0], per_node_force) + return out + + +class TestBridgedGraphSelfComm: + @pytest.fixture(autouse=True) + def _setup(self): + ensure_comm_registered() + self.model = _make_bridged_model() + + def _run_folded(self, coord: np.ndarray): + atype = np.array([[0, 0, 1, 1]], dtype=np.int64) + box = (_L * np.eye(3, dtype=np.float64)).reshape(1, 3, 3) + graph = build_neighbor_graph(coord, atype, box, 4.0, canonicalize=True) + out = self.model.forward_common_lower_graph( + torch.tensor(atype.reshape(-1), dtype=torch.int64), + torch.as_tensor(np.asarray(graph.n_node), dtype=torch.int64), + torch.as_tensor(np.asarray(graph.n_node), dtype=torch.int64), + torch.as_tensor(np.asarray(graph.edge_index), dtype=torch.int64), + torch.as_tensor(np.asarray(graph.edge_vec), dtype=torch.float64), + torch.as_tensor(np.asarray(graph.edge_mask), dtype=torch.bool), + ) + e = out["energy_redu"].detach().numpy().reshape(-1) + f = -out["energy_derv_r"].detach().numpy().reshape(-1, 3) + return e, f + + def _run_self_comm(self, coord: np.ndarray): + ext_coord, ext_atype, nlist, mapping = _extended_quartet(coord) + gi = _unfolded_graph_inputs(ext_coord, ext_atype, nlist) + nall = ext_coord.shape[1] + keepalive: list = [] + comm_dict = _build_self_comm_dict( + nloc=4, + nghost=nall - 4, + sendlist_indices=mapping[0, 4:].astype(np.int32), + keepalive=keepalive, + ) + out = self.model.forward_common_lower_graph( + gi["atype"], + gi["n_node"], + gi["n_local"], + gi["edge_index"], + gi["edge_vec"], + gi["edge_mask"], + comm_dict=comm_dict, + ) + e = out["energy_redu"].detach().numpy().reshape(-1) + f_ext = -out["energy_derv_r"].detach().numpy().reshape(-1, 3) + return e, _fold_forces(f_ext, mapping) + + @pytest.mark.parametrize( + "gap", + [ + 0.4, # r < r_inner=0.8: hard freeze, exercises the zero_count channel + 1.0, # r_inner < r < r_outer=1.2: exercises the log_eta channel + ], + ) + def test_self_comm_matches_folded_reference(self, gap: float) -> None: + coord = _close_pair_coords(gap) + e_ref, f_ref = self._run_folded(coord) + e_par, f_par = self._run_self_comm(coord) + np.testing.assert_allclose(e_par, e_ref, rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(f_par, f_ref, rtol=1e-12, atol=1e-12) + + @pytest.mark.parametrize( + "gap", + [ + 0.4, # zero_count channel: ablation loses the hard freeze + 1.0, # log_eta channel: ablation loses the transition attenuation + ], + ) + def test_ablation_identity_exchange_diverges( + self, gap: float, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Negative contract: with the exchange stubbed to identity the + parity BREAKS -- proves the geometry actually exercises the gate. + """ + from deepmd.pt_expt.descriptor import dpa4 as pe_dpa4 + + monkeypatch.setattr( + pe_dpa4.DescrptDPA4, + "_gate_partial_exchange", + lambda self, partials, comm_dict: partials, + ) + coord = _close_pair_coords(gap) + e_ref, _ = self._run_folded(coord) + e_par, _ = self._run_self_comm(coord) + assert np.abs(e_par - e_ref).max() > 1e-6 From 9695d2728abd093b7b552057a775b091062a3e68 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 10:49:46 +0800 Subject: [PATCH 09/26] feat(pt_expt): bridged DPA4 compositions export the with-comm artifact supports_edge_parallel flips to True (the SFPG exchange completes the gate partials across ranks); the graph freeze embeds the nested forward_lower_with_comm.pt2 for bridged compositions and gen_dpa4_zbl asserts it (issue #5906 Task 2). --- deepmd/dpmodel/descriptor/dpa4.py | 11 ++-- .../dpmodel/test_atomic_model_capabilities.py | 12 ++-- .../tests/common/dpmodel/test_descrpt_dpa4.py | 6 +- source/tests/infer/gen_dpa4_zbl.py | 21 ++++--- .../pt_expt/model/test_dpa4_zbl_parallel.py | 58 +++++++++++++++++++ .../pt_expt/model/test_export_with_comm.py | 12 ++-- .../tests/pt_expt/model/test_zbl_bridging.py | 10 ++-- 7 files changed, 99 insertions(+), 31 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index acb3306228..2f48a1487f 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -2327,7 +2327,7 @@ def has_message_passing_across_ranks(self) -> bool: The GRAPH lower implements the exchange via per-block ``border_op`` (pt_expt ``exchange_ghost_features``), so multi-rank inference always needs the with-comm artifact. Whether multi-rank is POSSIBLE at all - is :meth:`supports_edge_parallel` (bridging vetoes it there). + is :meth:`supports_edge_parallel`. The DENSE (nlist) lower remains comm-less — see :meth:`dense_lower_supports_comm`; the freeze machinery consults both @@ -2336,12 +2336,13 @@ def has_message_passing_across_ranks(self) -> bool: return True def supports_edge_parallel(self) -> bool: - """Bridging vetoes multi-rank until the SFPG exchange lands. + """Bridging included: multi-rank is supported for every SeZM config. - The Source Freeze Propagation Gate folds each node's full - outgoing-edge set, which no single rank observes (issue #5906). + The SFPG per-node partials are completed across ranks by + ``_gate_partial_exchange`` (reverse-accumulate + broadcast) before + the gate is applied (issue #5906). """ - return self.bridging_switch is None + return True def dense_lower_supports_comm(self) -> bool: """The DPA4 dense (nlist) lower has no comm_dict implementation. diff --git a/source/tests/common/dpmodel/test_atomic_model_capabilities.py b/source/tests/common/dpmodel/test_atomic_model_capabilities.py index 145e03bf9c..0d1bf8cdc0 100644 --- a/source/tests/common/dpmodel/test_atomic_model_capabilities.py +++ b/source/tests/common/dpmodel/test_atomic_model_capabilities.py @@ -100,7 +100,10 @@ def test_dp_atomic_model_delegates_to_descriptor() -> None: plain = _dp_atomic_model(_dpa4_descriptor(bridging=False)) local_only = _dp_atomic_model(DescrptSeA(rcut=4.0, rcut_smth=3.5, sel=[8, 8])) assert bridged.has_message_passing_across_ranks() is True - assert bridged.supports_edge_parallel() is False + # bridged is True too since the SFPG cross-rank completion (issue + # #5906); the ALL-aggregation's False branch is pinned by the stub + # children in test_linear_aggregation_mixed_children. + assert bridged.supports_edge_parallel() is True assert plain.has_message_passing_across_ranks() is True assert plain.supports_edge_parallel() is True assert local_only.has_message_passing_across_ranks() is False @@ -161,13 +164,14 @@ def _stub_linear(**child_kwargs_pair) -> LinearEnergyAtomicModel: def test_linear_aggregation_any_all() -> None: - """Real ZBL composition: the bridged DP child sets needs=True (any) and - vetoes edge-parallel (all). + """Real ZBL composition: the bridged DP child sets needs=True (any); + edge-parallel aggregates to True since the SFPG cross-rank completion + (issue #5906). """ am = get_model(copy.deepcopy(ZBL_CONFIG)).atomic_model assert isinstance(am, LinearEnergyAtomicModel) assert am.has_message_passing_across_ranks() is True - assert am.supports_edge_parallel() is False + assert am.supports_edge_parallel() is True # ZBL rides the graph route with the learned child assert am.supports_graph_export() is True assert am.graph_edge_dtype() == "float64" diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py index 4294026004..93b94f76e7 100644 --- a/source/tests/common/dpmodel/test_descrpt_dpa4.py +++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py @@ -112,15 +112,15 @@ def test_message_passing_semantics(self) -> None: def test_capability_split_needs_vs_supports(self) -> None: """has_message_passing_across_ranks = NEEDS exchange (always True for - SeZM); supports_edge_parallel = CAN run multi-rank (bridging vetoes, - until the SFPG exchange lands -- issue #5906). + SeZM); supports_edge_parallel = CAN run multi-rank (True for bridged + models too since the SFPG cross-rank completion -- issue #5906). """ dd_plain = make_descriptor() dd_bridged = make_descriptor(inner_clamp_r_inner=0.5, inner_clamp_r_outer=1.0) assert dd_plain.has_message_passing_across_ranks() is True assert dd_bridged.has_message_passing_across_ranks() is True assert dd_plain.supports_edge_parallel() is True - assert dd_bridged.supports_edge_parallel() is False + assert dd_bridged.supports_edge_parallel() is True def test_gate_partial_exchange_dpmodel_raises(self) -> None: """The dpmodel backend is the single-process reference; comm on a diff --git a/source/tests/infer/gen_dpa4_zbl.py b/source/tests/infer/gen_dpa4_zbl.py index a8f5f6ade0..2c0a768505 100644 --- a/source/tests/infer/gen_dpa4_zbl.py +++ b/source/tests/infer/gen_dpa4_zbl.py @@ -10,11 +10,14 @@ only end-to-end test drove the ``.pt2`` through the PYTHON ``DeepPot``, which never touches ``DeepPotPTExpt``. -Single-rank only: bridging enables the descriptor's Source Freeze -Propagation Gate, whose per-node ``eta_j = prod_{e: src_e = j} w_e`` folds a -node's FULL outgoing-edge set. Edges exist only for owned centres, so eta is -incomplete on every rank and a with-comm artifact is not exported -(``has_comm_artifact=false``, asserted below). +Multi-rank capable (issue #5906): bridging enables the descriptor's Source +Freeze Propagation Gate, whose per-node ``eta_j = prod_{e: src_e = j} w_e`` +folds a node's FULL outgoing-edge set. Edges exist only for owned centres, +so the per-node ``[log_eta, zero_count]`` partials are rank-incomplete; the +with-comm artifact completes them via one reverse-accumulate +(``border_op_backward``) + forward-broadcast (``border_op``) exchange before +the gate is applied, so the nested ``forward_lower_with_comm.pt2`` is +embedded (``has_comm_artifact=true``, asserted below). Generation mirrors ``gen_dpa4_spin.py``: the dpmodel is built in-process from the inline config with a fixed weight-init seed, its zero-initialised @@ -150,10 +153,10 @@ def main(): ) assert md["type_map"] == ZBL_CONFIG["type_map"] assert md["lower_input_kind"] == "graph" - # Single-rank only -- see the module docstring (SFPG eta is incomplete - # per rank), so no nested with-comm artifact may be present. - assert md["has_comm_artifact"] is False - assert "model/extra/forward_lower_with_comm.pt2" not in names + # Multi-rank capable (issue #5906): the SFPG per-node partials are + # completed across ranks, so the nested with-comm artifact is embedded. + assert md["has_comm_artifact"] is True + assert "model/extra/forward_lower_with_comm.pt2" in names # ---- 4. Evaluate (PBC + NoPbc) ---- dp = DeepPot(pt2_path) diff --git a/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py b/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py index 450f8de720..73e1a2f39d 100644 --- a/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py +++ b/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py @@ -251,3 +251,61 @@ def test_ablation_identity_exchange_diverges( e_ref, _ = self._run_folded(coord) e_par, _ = self._run_self_comm(coord) assert np.abs(e_par - e_ref).max() > 1e-6 + + +class TestBridgedGraphWithCommExport: + """Issue #5906 Task 2: the bridged composition's with-comm export rungs.""" + + @pytest.fixture(autouse=True) + def _setup(self): + ensure_comm_registered() + self.model = _make_bridged_model() + + def test_make_fx_traces_with_comm(self) -> None: + """The exchange (border_op_backward + border_op) traces symbolically + into the with-comm graph forward. + """ + from deepmd.pt_expt.utils.serialization import ( + _trace_and_export, + ) + + data = {"model": self.model.serialize()} + exported, _meta, _dj, _keys = _trace_and_export( + copy.deepcopy(data), + model_json_override=None, + with_comm_dict=True, + lower_kind="graph", + ) + loaded = exported.module() + placeholders = loaded.graph.find_nodes(op="placeholder") + assert len(placeholders) == 21, ( + f"graph with-comm program must accept 21 positional inputs " + f"(13 graph-base incl. n_local + 8 comm); got {len(placeholders)}" + ) + gm_code = str(loaded.code) + assert "deepmd_export.border_op_backward" in gm_code + assert ( + "deepmd_export.border_op." in gm_code + or "deepmd_export.border_op(" in gm_code + ) + + def test_freeze_embeds_with_comm_artifact(self, tmp_path) -> None: + """Freezing the bridged model to a graph ``.pt2`` embeds the nested + with-comm artifact (mirrors the plain-DPA4 embed test). + """ + import json + import zipfile + + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + data = {"model": self.model.serialize()} + p = str(tmp_path / "m_dpa4_zbl_graph.pt2") + deserialize_to_file(p, copy.deepcopy(data), lower_kind="graph") + with zipfile.ZipFile(p, "r") as zf: + names = zf.namelist() + meta = json.loads(zf.read("model/extra/metadata.json")) + assert "model/extra/forward_lower_with_comm.pt2" in names + assert meta["has_comm_artifact"] is True + assert meta["lower_input_kind"] == "graph" diff --git a/source/tests/pt_expt/model/test_export_with_comm.py b/source/tests/pt_expt/model/test_export_with_comm.py index 8ddcbaadcc..886d7aaed3 100644 --- a/source/tests/pt_expt/model/test_export_with_comm.py +++ b/source/tests/pt_expt/model/test_export_with_comm.py @@ -424,10 +424,11 @@ def test_two_dpa2_linear_composition_gets_with_comm_artifact() -> None: assert _needs_with_comm_artifact(model, lower_kind="nlist") is True -def test_bridged_composition_still_denied_until_sfpg_exchange() -> None: - """Negative contract: bridged DPA4+ZBL stays single-rank in Phase 1 -- - the veto now comes from supports_edge_parallel aggregation, not from - the wrapper's type. +def test_bridged_composition_gets_with_comm_artifact() -> None: + """Bridged DPA4+ZBL is multi-rank once the SFPG exchange lands + (issue #5906 Task 2): supports_edge_parallel aggregation admits it, + so the graph with-comm artifact is requested. The dense (nlist) kind + stays denied -- DPA4's dense lower has no comm implementation. """ import copy @@ -443,4 +444,5 @@ def test_bridged_composition_still_denied_until_sfpg_exchange() -> None: model = get_pt_expt_model(copy.deepcopy(ZBL_CONFIG)) model.to("cpu") model.eval() - assert _needs_with_comm_artifact(model, lower_kind="graph") is False + assert _needs_with_comm_artifact(model, lower_kind="graph") is True + assert _needs_with_comm_artifact(model, lower_kind="nlist") is False diff --git a/source/tests/pt_expt/model/test_zbl_bridging.py b/source/tests/pt_expt/model/test_zbl_bridging.py index 8151c88f7f..a1239b984d 100644 --- a/source/tests/pt_expt/model/test_zbl_bridging.py +++ b/source/tests/pt_expt/model/test_zbl_bridging.py @@ -220,15 +220,15 @@ def test_serialize_roundtrip(self) -> None: out["energy"], out2["energy"], rtol=1e-12, atol=1e-12 ) - def test_with_comm_gate_off_for_composition(self) -> None: - """Compositions never compile a with-comm artifact (single-rank).""" + def test_with_comm_gate_on_for_composition(self) -> None: + """The SFPG exchange makes bridged compositions multi-rank + (issue #5906 Task 2): the graph with-comm artifact is compiled. + """ from deepmd.pt_expt.utils.serialization import ( _needs_with_comm_artifact, ) - assert ( - _needs_with_comm_artifact(self.pt_expt_model, lower_kind="graph") is False - ) + assert _needs_with_comm_artifact(self.pt_expt_model, lower_kind="graph") is True def test_pt_bridging_checkpoint_rejected(self) -> None: """Reject pt's flag-serialized bridging checkpoints. From 28588e6992161854f7477466224dc2d4b164fadb Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 11:02:25 +0800 Subject: [PATCH 10/26] feat(pt): SFPG cross-rank completion on the pt edge path Closes the #5906 pt gap: the gate's per-node partials are completed via border_op_backward + border_op before the gate is applied, and the model-level ZBL injection reads EXTENDED types (the parallel path used to crash on ghost src indices). Bridged SeZM reports supports_edge_parallel=True; close-pair parity pinned on both bridging channels with an identity-exchange ablation. --- deepmd/pt/model/descriptor/sezm.py | 62 ++++- .../pt/model/descriptor/sezm_nn/edge_cache.py | 40 +++- deepmd/pt/model/model/sezm_model.py | 24 +- source/tests/pt/model/test_sezm_parallel.py | 13 +- .../test_sezm_parallel_bridging_parity.py | 211 ++++++++++++++++++ 5 files changed, 324 insertions(+), 26 deletions(-) create mode 100644 source/tests/pt/model/test_sezm_parallel_bridging_parity.py diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 9b0dabd052..394c3a2eea 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -32,6 +32,7 @@ annotations, ) +import functools import math from contextlib import ( contextmanager, @@ -1509,6 +1510,14 @@ def forward_with_edges( type_ebed, spin, atype_flat, n_nodes=n_nodes ) + # Cross-rank SFPG completion (issue #5906): only a bridged model + # under domain decomposition needs it -- the gate's src-keyed + # per-node partials are rank-incomplete then. + node_partial_exchange = None + if parallel and self.bridging_switch is not None: + node_partial_exchange = functools.partial( + self._gate_partial_exchange, comm_dict=comm_dict + ) # === Step 3. Build edge cache once (sparse edges) === with nvtx_range("build_edge_cache"): edge_cache = build_edge_cache_from_edges( @@ -1533,6 +1542,7 @@ def forward_with_edges( random_gamma=self.random_gamma and self.training, wigner_calc=self.wigner_calc, build_wigner=self._need_full_wigner, + node_partial_exchange=node_partial_exchange, ) ebed_dim_0 = self.node_init_dim # (node_init_lmax+1)^2 @@ -2290,12 +2300,56 @@ def has_message_passing_across_ranks(self) -> bool: return True def supports_edge_parallel(self) -> bool: - """Bridging vetoes multi-rank until the SFPG exchange lands. + """Bridging included: multi-rank is supported for every SeZM config. - The Source Freeze Propagation Gate folds each node's full - outgoing-edge set, which no single rank observes (issue #5906). + The SFPG per-node partials are completed across ranks by + ``_gate_partial_exchange`` (reverse-accumulate + broadcast) before + the gate is applied (issue #5906). """ - return self.bridging_switch is None + return True + + def _gate_partial_exchange( + self, + partials: torch.Tensor, + comm_dict: dict[str, torch.Tensor], + ) -> torch.Tensor: + """Reverse-accumulate ghost partials to owners, then broadcast back. + + ``border_op_backward`` sums each ghost row into its owner across + ranks and zeroes the ghost rows; ``border_op`` refills them with the + completed owner values. Both ops carry autograd (the two are + transposes), so gate gradients cross ranks (issue #5906). + + Parameters + ---------- + partials + (n_nodes, 2) float tensor of [log_eta, zero_count] partials. + comm_dict + The border-exchange control tensors. + + Returns + ------- + torch.Tensor + The globally completed (n_nodes, 2) tensor. + """ + # border_op exchanges rows by raw pointer arithmetic; a strided + # view would corrupt the exchange. + p = partials.contiguous() + comm_args = ( + comm_dict["send_list"], + comm_dict["send_proc"], + comm_dict["recv_proc"], + comm_dict["send_num"], + comm_dict["recv_num"], + ) + tail = ( + comm_dict["communicator"], + comm_dict["nlocal"], + comm_dict["nghost"], + ) + p = torch.ops.deepmd_export.border_op_backward(*comm_args, p, *tail) + p = torch.ops.deepmd_export.border_op(*comm_args, p, *tail) + return p def need_sorted_nlist_for_lower(self) -> bool: return False diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py index 383c4d3e5d..ff1d497d8b 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py +++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py @@ -117,6 +117,7 @@ def compute_edge_src_gate( n_nodes: int, bridging_switch: Callable[[torch.Tensor], torch.Tensor], edge_keep_f: torch.Tensor | None = None, + node_partial_exchange: Callable[[torch.Tensor], torch.Tensor] | None = None, ) -> torch.Tensor: """ Compute the per-edge source gate for SFPG from edge lengths. @@ -168,6 +169,14 @@ def compute_edge_src_gate( Optional per-edge keep weights with shape (E, 1), with ``0`` on masked edges and ``1`` on kept edges. If provided, masked edges are rewritten to ``w = 1`` before the product reduction. + node_partial_exchange + Optional cross-rank completion hook for the per-node partials + (issue #5906). Receives the ``(n_nodes, 2)`` float tensor + ``[log_eta, zero_count]`` and returns the globally completed one + (reverse-accumulate ghost rows into owners, then broadcast the + completed owner values back onto ghosts). ``None`` (the default) + is the single-process path where the local partials are already + complete. Returns ------- @@ -193,16 +202,31 @@ def compute_edge_src_gate( log_eta = torch.zeros( n_nodes, dtype=edge_w.dtype, device=edge_w.device ).scatter_add(0, src, log_safe) - eta_nonzero_path = torch.exp(log_eta) # === Step 3. Exact-zero indicator per source node === - # ``scatter_add`` over an ``int64`` cast of the zero mask counts how - # many frozen edges each source node owns. A strictly positive count - # means the product is 0 by the hard-freeze rule. + # ``scatter_add`` over the zero mask counts how many frozen edges each + # source node owns. A strictly positive count means the product is 0 by + # the hard-freeze rule. Float count (values are small integers, exact + # in fp) so both partials ride ONE border-exchange tensor when + # completing across ranks. zero_count = torch.zeros( - n_nodes, dtype=torch.int64, device=edge_w.device - ).scatter_add(0, src, is_zero.to(torch.int64)) - any_zero = zero_count > 0 + n_nodes, dtype=edge_w.dtype, device=edge_w.device + ).scatter_add(0, src, is_zero.to(edge_w.dtype)) + + # === Step 3b. Cross-rank completion of the per-node partials === + # A rank only holds edges whose dst is owned, so the src-keyed sums + # above are PARTIAL for every node under domain decomposition. The hook + # (reverse-accumulate ghost->owner, then forward-broadcast owner->ghost) + # completes them; log-products are additive and each edge lives on + # exactly one rank, so nothing double-counts (issue #5906). + if node_partial_exchange is not None: + packed = torch.stack([log_eta, zero_count], dim=-1) # (n_nodes, 2) + packed = node_partial_exchange(packed) + log_eta = packed[..., 0] + zero_count = packed[..., 1] + + eta_nonzero_path = torch.exp(log_eta) + any_zero = zero_count > 0.5 # === Step 4. Combine and broadcast back to edges via source === eta = torch.where(any_zero, torch.zeros_like(eta_nonzero_path), eta_nonzero_path) @@ -389,6 +413,7 @@ def build_edge_cache_from_edges( random_gamma: bool, wigner_calc: WignerCalculatorFn, build_wigner: bool = True, + node_partial_exchange: Callable[[torch.Tensor], torch.Tensor] | None = None, ) -> EdgeFeatureCache: """ Build the global edge cache from a sparse edge list. @@ -499,6 +524,7 @@ def build_edge_cache_from_edges( n_nodes=n_nodes, bridging_switch=bridging_switch, edge_keep_f=edge_keep_f, + node_partial_exchange=node_partial_exchange, ) return _finalize_edge_cache( diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 98780899db..3e8e931479 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -1564,10 +1564,15 @@ def core_compute( # ZBL is evaluated from ``edge_vec`` (the autograd leaf) so its force # and virial flow through the same edge backward as the learned energy. if self.inter_potential is not None and "energy" in fit_ret: + # The parallel (comm_dict) path indexes edges over the EXTENDED + # node set (ghost srcs), so the Z lookup needs the extended + # types; ``index_add_`` targets stay owned (dst < nloc) and the + # per-atom read-out is unchanged (issue #5906). + zbl_atype = atype if extended_atype is None else extended_atype fit_ret["energy"] = fit_ret["energy"] + self.inter_potential( edge_vec=edge_vec, edge_index=edge_index, - atype_flat=atype.reshape(-1), + atype_flat=zbl_atype.reshape(-1), edge_mask=edge_mask, n_node=nf * nloc, real_type_count=self._get_inter_potential_real_type_count(), @@ -3030,16 +3035,15 @@ def has_message_passing(self) -> bool: def supports_edge_parallel(self) -> bool: """Whether the edge-based LAMMPS multi-rank with-comm artifact applies. - Cross-rank ghost-feature exchange is well-defined only for the - conservative non-bridging path: analytical ZBL bridging and its Source - Freeze Propagation gate fold each node's full outgoing-edge set, which a - single rank cannot observe for ghost owners. The native spin scheme - reuses the edge_vec interface and therefore participates; only the - deepspin (virtual-atom) scheme uses the nlist interface and is excluded - by the freeze entry point's edge_vec gate. + Delegates to the descriptor capability. Bridging participates since + issue #5906: the Source Freeze Propagation Gate's per-node partials + are completed across ranks (``_gate_partial_exchange``) and the + analytical ZBL injection reads extended types, so the parallel path + reproduces the folded one. The native spin scheme reuses the + edge_vec interface and therefore participates; only the deepspin + (virtual-atom) scheme uses the nlist interface and is excluded by + the freeze entry point's edge_vec gate. """ - if self.inter_potential is not None: - return False descriptor = self.atomic_model.descriptor return bool(descriptor.supports_edge_parallel()) diff --git a/source/tests/pt/model/test_sezm_parallel.py b/source/tests/pt/model/test_sezm_parallel.py index ceb53de9f5..be3b093fc1 100644 --- a/source/tests/pt/model/test_sezm_parallel.py +++ b/source/tests/pt/model/test_sezm_parallel.py @@ -431,9 +431,9 @@ def test_native_spin_mag_force_fold_parity_cpu(self) -> None: class TestSeZMEdgeParallelCapability(unittest.TestCase): """The with-comm export predicate admits the edge_vec contract. - Plain energy and native spin both use the edge_vec lower interface and are - rank-decomposable, so they support the with-comm artifact; only analytical - bridging (Source Freeze Propagation) is gated out. + Plain energy, native spin, and analytical bridging all use the edge_vec + lower interface and are rank-decomposable; bridging participates since + the SFPG cross-rank completion (issue #5906). """ def test_plain_model_supports_edge_parallel(self) -> None: @@ -452,7 +452,10 @@ def test_native_spin_supports_edge_parallel(self) -> None: self.assertTrue(model.supports_edge_parallel()) self.assertEqual(model.export_lower_input_kind(), "edge_vec") - def test_bridging_model_fails_fast(self) -> None: + def test_bridging_model_supports_edge_parallel(self) -> None: + """Bridging is multi-rank since the SFPG exchange (issue #5906); + parity is pinned in test_sezm_parallel_bridging_parity.py. + """ # ZBL needs real element symbols for its analytical pair potential. model = _build_model( torch.device("cpu"), @@ -461,7 +464,7 @@ def test_bridging_model_fails_fast(self) -> None: bridging_r_inner=0.5, bridging_r_outer=1.0, ) - self.assertFalse(model.supports_edge_parallel()) + self.assertTrue(model.supports_edge_parallel()) class TestSeZMExchangeSchedule(unittest.TestCase): diff --git a/source/tests/pt/model/test_sezm_parallel_bridging_parity.py b/source/tests/pt/model/test_sezm_parallel_bridging_parity.py new file mode 100644 index 0000000000..eb318e97a7 --- /dev/null +++ b/source/tests/pt/model/test_sezm_parallel_bridging_parity.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Issue #5906 Task 2 (pt backend): SFPG completion across ranks. + +A BRIDGED SeZM model's parallel (comm_dict) path must reproduce the folded +single-domain path. The Source Freeze Propagation Gate folds each node's +full outgoing-edge set; under domain decomposition a rank only holds edges +with owned destinations, so the src-keyed per-node partials are incomplete +and must be completed by one reverse-accumulate + forward-broadcast border +exchange before the gate is applied. + +Geometry is the load-bearing part: a sub-``r_outer`` pair STRADDLES the +periodic x boundary so the close contact's src is a ghost image row. +Before issue #5906's fix, pt's bridged parallel path had no guard: the +descriptor computed the partial (rank-incomplete) gate silently, and the +model-level ZBL injection crashed outright on extended src indices +(``z_all[src]`` IndexError) -- this test's red run demonstrated both. +""" + +import unittest + +import numpy as np +import torch + +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.utils import ( + env, # noqa: F401 - imports pt test env side effects +) + +from .test_sezm_parallel import ( + _perturb_descriptor, + _self_comm_dict, + _tiny_parallel_model_params, +) + +_L = 10.0 # box edge (= rcut * 2.5, matching the non-bridged parity file) + + +def _bridged_model_params() -> dict: + params = _tiny_parallel_model_params() + # ZBL bridging needs REAL element symbols (nuclear charges) + params["type_map"] = ["Ni", "O"] + # bridging window: r_inner=0.8, r_outer=1.2 (matches the pt_expt twin + # test, source/tests/pt_expt/model/test_dpa4_zbl_parallel.py) + params["bridging_method"] = "ZBL" + params["bridging_r_inner"] = 0.8 + params["bridging_r_outer"] = 1.2 + return params + + +def _close_pair_coords(gap: float, nloc: int = 4) -> np.ndarray: + """Atoms 0/1 straddle the x periodic boundary at distance ``gap``.""" + half = gap / 2.0 + coords = np.array( + [ + [half, 5.0, 5.0], + [_L - half, 5.0, 5.0], + [4.0, 5.0, 5.0], + [5.1, 5.2, 5.0], + ], + dtype=np.float64, + ) + assert coords.shape[0] == nloc + return coords.reshape(1, nloc, 3) + + +def _build_close_pair_system( + model: torch.nn.Module, device: torch.device, gap: float +) -> dict[str, torch.Tensor]: + """The non-bridged file's ``_build_extended_system`` with fixed coords.""" + from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, + ) + from deepmd.pt_expt.utils.edge_schema import ( + edge_schema_from_extended, + ) + + nloc = 4 + rcut = float(model.get_rcut()) + sel = list(model.get_sel()) + coord_np = _close_pair_coords(gap, nloc) + ntypes = len(model.get_type_map()) + atype_np = (np.arange(nloc, dtype=np.int32) % ntypes).reshape(1, nloc) + box = np.eye(3, dtype=np.float64) * _L + + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord_np, atype_np, box.reshape(1, 9), rcut + ) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + nloc, + rcut, + sel, + distinguish_types=not model.mixed_types(), + ) + extended_coord = np.asarray(extended_coord).reshape(1, -1, 3) + + ext_coord = torch.tensor(extended_coord, dtype=torch.float64, device=device) + ext_atype = torch.tensor( + np.asarray(extended_atype), dtype=torch.int64, device=device + ) + nlist_t = torch.tensor(np.asarray(nlist), dtype=torch.int64, device=device) + mapping_t = torch.tensor(np.asarray(mapping), dtype=torch.int64, device=device) + + formatted = model.format_nlist(ext_coord, ext_atype, nlist_t) + schema = edge_schema_from_extended( + ext_coord, ext_atype[:, :nloc], formatted, mapping_t + ) + return { + "coord": schema.coord, + "atype": schema.atype, + "extended_atype": ext_atype, + "edge_index": schema.edge_index, + "edge_vec": schema.edge_vec, + "edge_scatter_index": schema.edge_scatter_index, + "edge_mask": schema.edge_mask, + "mapping": mapping_t, + "nloc": nloc, + "nall": ext_coord.shape[1], + } + + +class TestSeZMBridgingSelfCommParity(unittest.TestCase): + """Bridged parallel path == folded path once the SFPG exchange lands.""" + + @classmethod + def setUpClass(cls) -> None: + from deepmd.pt_expt.utils.comm import ( + ensure_comm_registered, + ) + + ensure_comm_registered() + + def _run_pair(self, gap: float, device: torch.device): + model = get_model(_bridged_model_params()) + model.eval() + model.to(device) + _perturb_descriptor(model.atomic_model.descriptor) + sysm = _build_close_pair_system(model, device, gap) + comm = _self_comm_dict(sysm["mapping"], sysm["nloc"], sysm["nall"]) + + ref = model.forward_lower( + sysm["coord"], + sysm["atype"], + sysm["edge_index"], + sysm["edge_vec"], + sysm["edge_scatter_index"], + sysm["edge_mask"], + do_atomic_virial=True, + ) + par = model.forward_lower( + sysm["coord"], + sysm["atype"], + sysm["edge_scatter_index"], + sysm["edge_vec"], + sysm["edge_scatter_index"], + sysm["edge_mask"], + do_atomic_virial=True, + comm_dict=comm, + extended_atype=sysm["extended_atype"], + ) + return ref, par + + def _assert_parity(self, gap: float) -> None: + device = torch.device("cpu") + ref, par = self._run_pair(gap, device) + self.assertGreater( + ref["extended_force"].abs().max().item(), + 1e-6, + msg="reference forces are ~0; the parity check would be vacuous", + ) + for key in ("energy", "extended_force", "virial"): + torch.testing.assert_close( + par[key], ref[key], rtol=1e-8, atol=1e-9, msg=f"mismatch in {key}" + ) + + def test_parity_hard_freeze_pair_cpu(self) -> None: + """A gap < r_inner pair: the zero_count channel crosses the boundary.""" + self._assert_parity(0.4) + + def test_parity_transition_pair_cpu(self) -> None: + """r_inner < gap < r_outer: the log_eta channel crosses the boundary.""" + self._assert_parity(1.0) + + def test_ablation_identity_exchange_diverges(self) -> None: + """Negative contract: stubbing the exchange to identity breaks the + parity -- proves the geometry actually exercises the gate. + """ + from unittest import ( + mock, + ) + + from deepmd.pt.model.descriptor import ( + sezm as pt_sezm, + ) + + with mock.patch.object( + pt_sezm.DescrptSeZM, + "_gate_partial_exchange", + lambda self, partials, comm_dict: partials, + ): + ref, par = self._run_pair(0.4, torch.device("cpu")) + diff = (par["energy"] - ref["energy"]).abs().max().item() + self.assertGreater(diff, 1e-6) + + +if __name__ == "__main__": + unittest.main() From 40b384095f6c0cc942476aeaa4807f03414c9d7a Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 13:14:35 +0800 Subject: [PATCH 11/26] test(lmp): bridged DPA4 2-rank parity with a cross-boundary close pair Replaces the fails-fast test: the with-comm artifact now completes the SFPG partials across ranks, so a 2-rank run with the 0.9 A pair straddling the processors 2 1 1 boundary must (and does) match the 1-rank reference (issue #5906 Task 2 E2E). --- source/lmp/tests/test_lammps_dpa4_zbl_pt2.py | 152 +++++++++++-------- 1 file changed, 86 insertions(+), 66 deletions(-) diff --git a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py index 135bf04426..18e9cf75fe 100644 --- a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py @@ -12,18 +12,19 @@ handling, per-atom virial accumulation and unit conversion), so a regression confined to it was invisible. -Single-rank only, deliberately ------------------------------- +Multi-rank correctness (issue #5906) +------------------------------------ Bridging enables the descriptor's Source Freeze Propagation Gate, whose per-node ``eta_j = prod_{e: src_e = j} w_e`` folds a node's FULL outgoing-edge -set. Edges exist only for owned centres, so eta is incomplete on every rank -and the freeze exports NO with-comm artifact (``gen_dpa4_zbl.py`` asserts -``has_comm_artifact is False`` and that no nested -``forward_lower_with_comm.pt2`` entry exists). There is therefore no -correct multi-rank answer to compare against; what this file pins instead is -that a multi-rank run FAILS LOUDLY rather than silently returning -wrong-but-plausible numbers -- see -``test_pair_deepmd_mpi_dpa4_zbl_fails_fast``. +set. Edges exist only for owned centres, so the per-node partials are +rank-incomplete; the with-comm artifact completes them with one +reverse-accumulate + forward-broadcast border exchange before the gate is +applied (``gen_dpa4_zbl.py`` asserts ``has_comm_artifact is True`` and that +the nested ``forward_lower_with_comm.pt2`` entry exists). This file pins the +end-to-end contract with a 2-rank vs 1-rank parity run whose sub-``r_outer`` +close pair STRADDLES the ``processors 2 1 1`` x-boundary -- without that +geometry every cross-rank gate contribution is ``log w = 0`` and the parity +holds vacuously -- see ``test_pair_deepmd_mpi_dpa4_zbl_close_pair_parity``. Reference values are computed LIVE at test-setup time via ``deepmd.infer.DeepPot.eval`` on the archive itself, mirroring @@ -72,13 +73,13 @@ data_file = Path(__file__).parent / "data_dpa4_zbl_pt2.lmp" # The MPI runner is backend-agnostic (DATAFILE PB_FILE OUTPUT + flags); reuse # the DPA3 driver verbatim rather than duplicate it (same pattern as -# test_lammps_dpa4_graph_pt2.py). Only the fail-fast test below uses it. +# test_lammps_dpa4_graph_pt2.py). Only the multi-rank tests below use it. mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py" -# Ceiling for the mpirun invocation. The fail-fast under test is expected to -# throw on EVERY rank before any collective, so a timeout means the guard -# regressed into a deadlock -- which is a test failure, not a slow machine. -_MPI_DEFAULT_TIMEOUT = 300.0 +# Ceiling for every mpirun invocation: a with-comm desync hangs the +# collective forever, so an unbounded should-succeed regression would hang +# the whole suite -- a timeout is a test failure, not a slow machine. +_MPI_DEFAULT_TIMEOUT = 600.0 # 6-atom NiO system, coordinates verbatim from # ``source/tests/infer/gen_dpa4_zbl.py``'s ``_COORDS``: atoms 0 and 1 sit @@ -101,6 +102,16 @@ # -> LAMMPS types [1,1,1,2,2,2] under identity ``pair_coeff * *``. type_NiO = np.array([1, 1, 1, 2, 2, 2]) +# Close-pair variant for the 2-rank parity test (issue #5906): the SAME +# 6-atom geometry shifted along x so the 0.9 A Ni-Ni pair (inside the +# bridging window, r_inner=0.8 < 0.9 < r_outer=1.2) STRADDLES the +# ``processors 2 1 1`` boundary at lx/2 = 6.5. Atom 0 lands at x = 6.3 +# (rank 0) and atom 1 at x = 7.2 (rank 1), so the pair's SFPG gate +# contribution crosses the rank boundary -- the load-bearing geometry. +_CLOSE_PAIR_X_SHIFT = 5.3 +coord_close_pair = coord + np.array([_CLOSE_PAIR_X_SHIFT, 0.0, 0.0]) +data_file_close_pair = Path(__file__).parent / "data_dpa4_zbl_close_pair_pt2.lmp" + # Reference values, populated by ``_compute_expected`` in ``setup_module``. expected_e = None expected_ae = None @@ -192,11 +203,13 @@ def setup_module() -> None: ) _compute_expected() write_lmp_data(box, coord, type_NiO, data_file) + write_lmp_data(box, coord_close_pair, type_NiO, data_file_close_pair) def teardown_module() -> None: - if data_file.exists(): - os.remove(data_file) + for f in (data_file, data_file_close_pair): + if f.exists(): + os.remove(f) def _lammps(data_file, units="metal") -> PyLammps: @@ -295,17 +308,24 @@ def test_pair_deepmd_atom_energy_and_virial(lammps) -> None: # --------------------------------------------------------------------------- -# Multi-rank: NOT a correctness test -- a fail-fast test. +# Multi-rank: 2-rank vs 1-rank close-pair parity (issue #5906 Task 2 E2E). # --------------------------------------------------------------------------- -def _run_mpi_subprocess(nprocs: int, processors: str, timeout: float) -> dict: +def _run_mpi_subprocess( + nprocs: int, + processors: str, + timeout: float = _MPI_DEFAULT_TIMEOUT, + data_path: Path | None = None, +) -> dict: """Run the (backend-agnostic) DPA3 MPI runner against the bridged archive - and return ``{"returncode", "stdout", "stderr", "timed_out"}``. + and return the parsed ``{"pe", "forces", "virials"}`` output. Always bounded: on expiry the WHOLE mpirun process group is SIGKILLed (killing only mpirun can leave orphaned ranks blocking in a collective). """ + if data_path is None: + data_path = data_file with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: out_path = f.name try: @@ -315,7 +335,7 @@ def _run_mpi_subprocess(nprocs: int, processors: str, timeout: float) -> dict: str(nprocs), sys.executable, str(mpi_runner), - str(data_file.resolve()), + str(data_path.resolve()), str(pb_file.resolve()), out_path, "--processors", @@ -329,18 +349,25 @@ def _run_mpi_subprocess(nprocs: int, processors: str, timeout: float) -> dict: except sp.TimeoutExpired: os.killpg(os.getpgid(proc.pid), signal.SIGKILL) stdout, stderr = proc.communicate() - return { - "returncode": None, - "stdout": stdout or "", - "stderr": stderr or "", - "timed_out": True, - } - return { - "returncode": proc.returncode, - "stdout": stdout, - "stderr": stderr, - "timed_out": False, - } + raise RuntimeError( + f"mpirun timed out after {timeout}s (process group killed); " + "a should-succeed MPI regression is deadlocked.\n" + f"stdout:\n{(stdout or '')[-2000:]}\n" + f"stderr:\n{(stderr or '')[-2000:]}" + ) from None + if proc.returncode != 0: + raise RuntimeError( + f"mpirun exited {proc.returncode}.\n" + f"stdout:\n{stdout[-2000:]}\nstderr:\n{stderr[-2000:]}" + ) + with open(out_path) as fh: + lines = fh.read().strip().splitlines() + pe = float(lines[0]) + rows = np.array( + [list(map(float, line.split())) for line in lines[1:]], + dtype=np.float64, + ) + return {"pe": pe, "forces": rows[:, :3], "virials": rows[:, 3:]} finally: if os.path.exists(out_path): os.remove(out_path) @@ -352,40 +379,33 @@ def _run_mpi_subprocess(nprocs: int, processors: str, timeout: float) -> dict: @pytest.mark.skipif( importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" ) -def test_pair_deepmd_mpi_dpa4_zbl_fails_fast() -> None: - """A multi-rank run of a BRIDGED archive must fail loudly, not answer. - - The bridged model is single-rank only by construction (see the module - docstring), so its freeze exports no with-comm artifact while still - declaring ``has_message_passing``. ``DeepPotPTExpt::compute_inner``'s - dispatch reads exactly that combination -- graph lower + ``nprocs > 1`` + - message passing + no with-comm artifact -- and throws before building any - tensors. Without the guard the run would fall through to the plain - single-rank artifact on a per-rank subdomain, where the bridging gate's - per-node eta is incomplete: wrong, finite, plausible numbers. - - The failure is uniform across ranks (every rank evaluates the same - metadata-only predicate before any collective), so a TIMEOUT is a failure - of this test: it would mean the guard regressed into a deadlock. - - This is deliberately the ONLY multi-rank test in this file; there is no - correct multi-rank reference for a bridged model to compare against. +def test_pair_deepmd_mpi_dpa4_zbl_close_pair_parity() -> None: + """Issue #5906 Task 2 E2E: a bridged model, 2 ranks, with a 0.9 A + contact STRADDLING the ``processors 2 1 1`` x-boundary. + + Without the cross-boundary close pair this test is vacuous (edges at + ``r >= r_outer`` contribute ``log w = 0``); the geometry guard below + pins that the pair actually straddles the split. The 2-rank run + exercises the with-comm artifact's SFPG completion (reverse-accumulate + + broadcast of the per-node ``[log_eta, zero_count]`` partials), the + per-block ghost exchange, and the reverse-comm force fold; the 1-rank + run is the plain-artifact reference on the same trajectory. """ - out = _run_mpi_subprocess( - nprocs=2, processors="2 1 1", timeout=_MPI_DEFAULT_TIMEOUT - ) - assert not out["timed_out"], ( - "Multi-rank run of the bridged archive timed out instead of failing " - "promptly; the dispatch guard must throw on every rank BEFORE any " - "collective." + lx = float(box[1] - box[0]) + x_lo, x_hi = float(coord_close_pair[0, 0]), float(coord_close_pair[1, 0]) + assert x_lo < lx / 2.0 < x_hi, ( + "the close pair no longer straddles the processors 2 1 1 boundary; " + "the parity below would be vacuous for the SFPG exchange" ) - assert out["returncode"] != 0, ( - "Expected the multi-rank run of a bridged (no with-comm artifact) " - "archive to fail loudly, but it exited 0.\n" - f"stdout:\n{out['stdout'][-2000:]}\nstderr:\n{out['stderr'][-2000:]}" + ref = _run_mpi_subprocess( + nprocs=1, processors="1 1 1", data_path=data_file_close_pair ) - combined = out["stdout"] + out["stderr"] - assert "with-comm artifact" in combined, ( - "Expected the documented fail-loud message (mentioning the missing " - f"'with-comm artifact'), got:\n{combined[-2000:]}" + par = _run_mpi_subprocess( + nprocs=2, processors="2 1 1", data_path=data_file_close_pair ) + assert par["pe"] == pytest.approx(ref["pe"], rel=1e-8, abs=1e-10) + np.testing.assert_allclose(par["forces"], ref["forces"], atol=1e-8, rtol=0) + # Same tolerance rationale as test_lammps_dpa4_graph_pt2.py's twin: the + # relative component absorbs CUDA atomic-scatter ordering noise without + # loosening the CPU-exact case. + np.testing.assert_allclose(par["virials"], ref["virials"], atol=1e-8, rtol=1e-8) From 88ce0d8a58dc41ecd98a02bd826ec5d5893ad31c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 14:06:24 +0800 Subject: [PATCH 12/26] test(spin): native-spin + ZBL multi-rank enabled and pinned (#5906 Task 3) The spin wrapper needed NO production change: the gate exchange sits below it. gen_dpa4_spin_zbl asserts the with-comm artifact; python self-comm parity (energy/force/force_mag, both bridging channels) at 1e-12; first-ever LAMMPS file for the spin+ZBL variant, incl. 2-rank close-pair parity via pair_style deepspin. --- .../tests/test_lammps_dpa4_spin_zbl_pt2.py | 566 ++++++++++++++++++ source/tests/infer/gen_dpa4_spin_zbl.py | 42 +- .../pt_expt/model/test_dpa4_zbl_parallel.py | 116 ++++ 3 files changed, 701 insertions(+), 23 deletions(-) create mode 100644 source/lmp/tests/test_lammps_dpa4_spin_zbl_pt2.py diff --git a/source/lmp/tests/test_lammps_dpa4_spin_zbl_pt2.py b/source/lmp/tests/test_lammps_dpa4_spin_zbl_pt2.py new file mode 100644 index 0000000000..a287553564 --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4_spin_zbl_pt2.py @@ -0,0 +1,566 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""LAMMPS ``pair_style deepspin`` on the COMBINED native-spin + ZBL-bridged +DPA4 NeighborGraph (graph-schema) ``.pt2`` +(``deeppot_dpa4_spin_zbl_graph.pt2``, generated by +``source/tests/infer/gen_dpa4_spin_zbl.py``). + +This pins the spin+ZBL variant end-to-end, which previously had NO LAMMPS +coverage at all: single-rank pair/virial against a LIVE ``DeepPot`` reference +computed at test-setup time on this module's own geometry, AND 2-rank vs +1-rank MPI parity -- including a close-pair geometry whose ~0.978 A Ni-Ni +pair straddles the ``processors 2 1 1`` domain boundary INSIDE the bridging +transition zone (the generator sets ``bridging_r_inner = 0.8`` / +``bridging_r_outer = 1.2``, verbatim from ``gen_dpa4_zbl.py``; 0.8 < 0.978 < +1.2). Bridging enables the descriptor's Source Freeze Propagation Gate, +whose per-node ``[log_eta, zero_count]`` partials fold a node's FULL +outgoing-edge set; edges exist only for owned centres, so the partials are +rank-incomplete. The SFPG cross-rank completion -- one reverse-accumulate +(``border_op_backward``) + forward-broadcast (``border_op``) exchange of the +per-node partials inside the with-comm artifact -- completes them before the +gate is applied, which is what makes bridged models multi-rank capable +(issue #5906). The close-pair test is the geometry that actually exercises +that exchange: without a bridging-zone pair split across ranks the parity +comparison would be vacuous. + +Reference (energy / force / force_mag / virial) values are computed LIVE at +test-setup time via ``deepmd.infer.DeepPot.eval`` on the archive for the +fixed 4-atom NiO system reused from ``test_lammps_dpa4_spin_graph_pt2.py`` +(box 13x13x13, 2 spin-active Ni + 2 non-magnetic O) -- deliberately NOT +hardcoded (a hardcoded reference goes stale the moment DPA4 numerics shift) +and NOT read from the generator's ``.expected`` sidecar (its own 6-atom +system sits in a 6x6x6 box whose edge exactly equals DPA4's ghost cutoff -- +not a safe geometry for a LAMMPS periodic run). ``_compute_expected`` runs +in a subprocess so importing ``deepmd``'s Python package does not share a +process with the LAMMPS plugin's own loaded ``libdeepmd_op_pt.so`` (see +``test_lammps_model_devi_pt2.py``). +""" + +import importlib.util +import json +import os +import shutil +import signal +import subprocess as sp +import sys +import tempfile +import textwrap +from pathlib import ( + Path, +) + +import constants +import numpy as np +import pytest +from lammps import ( + PyLammps, +) +from write_lmp_data import ( + write_lmp_data_spin, +) + +pb_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4_spin_zbl_graph.pt2" +) +data_file = Path(__file__).parent / "data_dpa4_spin_zbl_pt2.lmp" +data_file_close_pair = Path(__file__).parent / "data_dpa4_spin_zbl_close_pair_pt2.lmp" +# The MPI runner is graph-spin-specific (no aparam / no NULL-type extras, +# unlike run_mpi_pair_deepmd_spin_dpa3_pt2.py's virtual-atom-scheme runner): +# the native-spin (+ZBL) DPA4 fixtures take no fparam/aparam, so the plain +# native-spin runner is reused verbatim -- only the archive differs. +mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_spin_graph_dpa4_pt2.py" + +_MPI_DEFAULT_TIMEOUT = 120.0 + +# Same 4-atom NiO system as test_lammps_dpa4_spin_graph_pt2.py (box, +# coordinates, and LAMMPS type ordering all reused verbatim): 2 Ni atoms +# (LAMMPS type 1, deepmd atype 0, spin-active) + 2 O atoms (LAMMPS type 2, +# deepmd atype 1, non-magnetic) -- matches the archive's +# ``type_map=["Ni", "O"]`` and ``use_spin=[True, False]`` +# (gen_dpa4_spin_zbl.py inherits both from gen_dpa4_spin.py). The Ni-Ni +# pair (atoms 0 and 1) sits ~0.978 A apart -- inside the bridging +# transition zone (0.8, 1.2), so the ZBL channel is active even in the +# single-rank tests. +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +spin = np.array( + [ + [0, 0, 1.2737], + [0, 0, 1.2737], + [0, 0, 0], + [0, 0, 0], + ] +) +type_NiO = np.array([1, 1, 2, 2]) + +# Close-pair geometry for the MPI-parity test: the SAME 4-atom system +# shifted along x by -5.9, so the ~0.978 A Ni-Ni pair lands at x = 6.93 / +# 6.19 -- straddling the lx/2 = 6.5 domain boundary of ``processors 2 1 1``. +# The shift pushes atom 2 (x = 3.51) below zero, so the x column is wrapped +# back into the periodic box with ``np.mod(x, 13)`` (LAMMPS is periodic, so +# wrapping changes nothing physically; it just keeps the data file's +# coordinates inside [0, lx) as write_lmp_data_spin expects). +_CLOSE_PAIR_SHIFT_X = -5.9 +coord_close_pair = coord.copy() +coord_close_pair[:, 0] = np.mod(coord[:, 0] + _CLOSE_PAIR_SHIFT_X, box[1]) + +# LAMMPS's ``fm`` (what ``compute property/atom fmx fmy fmz`` reports) is +# NOT the raw DeepEval force_mag: pair_deepspin.cpp scales it by +# ``spin_norm / hbar`` per atom (metal-units ``hbar = 6.5821191e-04``, see +# ``source/lmp/pair_deepspin.cpp:531,535``). ``spin_norm`` is 0 for the two +# non-magnetic O atoms, so the scaling is a no-op there (0 stays 0). +_HBAR_METAL = 6.5821191e-04 + +# Reference values (energy / atom-energy / force / force_mag / virial), +# populated by ``_compute_expected`` in ``setup_module`` -- see the module +# docstring for why these are computed live via a DeepPot subprocess call +# rather than hardcoded or read from a sidecar file. +expected_e = None +expected_ae = None +expected_f = None +expected_fm = None +expected_v = None + + +def _cell_from_lammps_box(lmp_box: np.ndarray) -> np.ndarray: + """Convert a LAMMPS ``xlo xhi ylo yhi zlo zhi xy xz yz`` box spec to a + flat, row-major 3x3 cell matrix (deepmd's ``box`` convention). + """ + xlo, xhi, ylo, yhi, zlo, zhi, xy, xz, yz = lmp_box + return np.array( + [ + xhi - xlo, + 0.0, + 0.0, + xy, + yhi - ylo, + 0.0, + xz, + yz, + zhi - zlo, + ] + ) + + +def _compute_expected() -> None: + """Load ``deeppot_dpa4_spin_zbl_graph.pt2`` via ``DeepPot`` and evaluate + the module's fixed 4-atom NiO system to obtain the Python reference. + + Runs in a subprocess to avoid importing ``deepmd`` in the LAMMPS test + process (see ``test_lammps_model_devi_pt2.py``'s ``_compute_expected`` + for the same precaution: the LAMMPS plugin already loads + ``libdeepmd_op_pt.so`` at the C++ level, and importing the Python + package on top of that can segfault). + """ + global expected_e, expected_ae, expected_f, expected_fm, expected_v + + cell = _cell_from_lammps_box(box) + atype = (type_NiO - 1).tolist() # LAMMPS 1-based -> deepmd 0-based (Ni=0, O=1) + + # The archive lives in ``source/tests/infer`` next to ``gen_common.py``, + # whose ``load_custom_ops()`` loads the build-tree ``libdeepmd_op_pt.so`` + # (registering ``deepmd::edge_force_virial``, which the graph ``.pt2`` + # inference needs). ``import deepmd.pt`` alone only loads the op library + # from SHARED_LIB_DIR, which the build-test env does not populate -- so + # the subprocess reuses that fallback (after importing ``deepmd.pt``, per + # its docstring) before constructing ``DeepPot``. + infer_dir = str(pb_file.resolve().parent) + script = textwrap.dedent(f"""\ + import json + import sys + import numpy as np + + sys.path.insert(0, {infer_dir!r}) + import deepmd.pt # noqa: F401 (triggers the base op-library load) + from gen_common import load_custom_ops + + load_custom_ops() + from deepmd.infer import DeepPot + + dp = DeepPot({str(pb_file.resolve())!r}) + e, f, v, ae, av, fm, mm = dp.eval( + np.array({coord.tolist()!r}).reshape(1, -1, 3), + np.array({cell.tolist()!r}).reshape(1, 9), + {atype!r}, + atomic=True, + spin=np.array({spin.tolist()!r}).reshape(1, -1, 3), + ) + print(json.dumps({{ + "e": float(e[0, 0]), + "ae": np.asarray(ae[0]).reshape(-1).tolist(), + "f": np.asarray(f[0]).tolist(), + "fm": np.asarray(fm[0]).tolist(), + "av": np.asarray(av[0]).tolist(), + }})) + """) + proc = sp.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + ) + if proc.returncode != 0: + raise RuntimeError(f"Failed to compute expected values:\n{proc.stderr}") + result = json.loads(proc.stdout.strip()) + + expected_e = result["e"] + expected_ae = np.array(result["ae"]) + expected_f = np.array(result["f"]) + # Raw DeepEval force_mag (dE/dspin), scaled by LAMMPS's own + # spin_norm / hbar unit convention (see the comment on ``_HBAR_METAL`` + # above) before comparison. + fm_raw = np.array(result["fm"]) + spin_norm = np.linalg.norm(spin, axis=1) + expected_fm = fm_raw * (spin_norm / _HBAR_METAL)[:, None] + # Per-atom virial, sign-flipped (LAMMPS convention) relative to DeepPot's + # atomic virial output (mirrors test_lammps_spin_pt2.py's convention). + expected_v = -np.array(result["av"]) + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) + if not pb_file.exists(): + pytest.skip("deeppot_dpa4_spin_zbl_graph.pt2 not found") + _compute_expected() + write_lmp_data_spin(box, coord, spin, type_NiO, data_file) + write_lmp_data_spin(box, coord_close_pair, spin, type_NiO, data_file_close_pair) + + +def teardown_module() -> None: + for path in (data_file, data_file_close_pair): + if path.exists(): + os.remove(path) + + +def _lammps(data_file, units="metal") -> PyLammps: + """Standard DeepSpin LAMMPS system, plus ``atom_modify map yes``. + + Mirrors ``lammps_test_utils.make_spin_lammps`` (not reused directly: it + does not set ``atom_modify``), with the map turned on -- the native-spin + DPA4 GRAPH ``.pt2`` needs the LAMMPS atom-map to resolve ghost-atom + indices to local owners for single-rank inference (same requirement as + the energy graph route; see ``pair_deepspin.cpp``'s + ``DeePMD-kit Error: Single-rank LAMMPS .pt2 inference requires + `atom_modify map yes``` check). + """ + if units != "metal": + raise ValueError("units for spin should be metal") + + lammps = PyLammps() + lammps.units(units) + lammps.boundary("p p p") + lammps.atom_style("spin") + lammps.atom_modify("map yes") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + lammps.mass("1 58") + lammps.mass("2 16") + lammps.timestep(0.0005) + lammps.fix("1 all nve") + return lammps + + +@pytest.fixture +def lammps(): + lmp = _lammps(data_file=data_file) + yield lmp + lmp.close() + + +def _gather_force_mag(lammps: PyLammps, natoms: int) -> np.ndarray: + """Extract per-atom force_mag in atom-id order. + + LAMMPS does not expose ``fm`` through the legacy ``extract``/ + ``gather_atoms`` registry (see ``run_mpi_pair_deepmd_spin_dpa3_pt2.py``'s + module docstring), so go via ``compute property/atom fmx fmy fmz`` + + ``gather`` (id-ordered on every rank, single-rank included). + """ + fm_global = lammps.lmp.gather("c_fmprop", 1, 3) + return np.array(fm_global, dtype=np.float64).reshape(natoms, 3) + + +def test_pair_deepspin(lammps) -> None: + """Single-rank LAMMPS energy + force + force_mag vs the Python DeepEval + reference on the spin+ZBL archive, on the same 4-atom NiO system. + """ + lammps.pair_style(f"deepspin {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("fmprop all property/atom fmx fmy fmz") + lammps.run(0) + + assert lammps.eval("pe") == pytest.approx(expected_e) + + forces = np.array([lammps.atoms[ii].force for ii in range(4)], dtype=np.float64) + ids = np.array([lammps.atoms[ii].id for ii in range(4)]) + order = np.argsort(ids) + forces = forces[order] + np.testing.assert_allclose(forces, expected_f, atol=1e-8, rtol=0) + + force_mag = _gather_force_mag(lammps, coord.shape[0]) + np.testing.assert_allclose(force_mag, expected_fm, atol=1e-8, rtol=0) + # Anti-vacuity / native-spin design invariant: force_mag on the two + # non-spin (O) atoms must be exactly zero, both in the Python reference + # (baked into expected_fm above) and as produced by LAMMPS. The + # analytical ZBL child, which knows nothing about spin, must not leak + # into this channel either. + np.testing.assert_array_equal(force_mag[2:], np.zeros((2, 3))) + + lammps.run(1) + + +def test_pair_deepspin_virial(lammps) -> None: + """Single-rank per-atom pe/pressure/virial via + ``pe/atom`` / ``pressure`` / ``centroid/stress/atom``, atol=1e-8, + rtol=1e-8. + """ + lammps.pair_style(f"deepspin {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("peatom all pe/atom pair") + lammps.compute("pressure all pressure NULL pair") + lammps.compute("virial all centroid/stress/atom NULL pair") + lammps.variable("eatom atom c_peatom") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps.variable(f"pressure{jj} equal c_pressure[{ii + 1}]") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") + lammps.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps.run(0) + + assert lammps.eval("pe") == pytest.approx(expected_e) + + forces = np.array([lammps.atoms[ii].force for ii in range(4)], dtype=np.float64) + ids = np.array([lammps.atoms[ii].id for ii in range(4)]) + order = np.argsort(ids) + forces = forces[order] + np.testing.assert_allclose(forces, expected_f, atol=1e-8, rtol=0) + + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 + np.testing.assert_allclose( + np.array(lammps.variables["eatom"].value), + expected_ae[idx_map], + atol=1e-8, + rtol=1e-8, + ) + + vol = box[1] * box[3] * box[5] + for ii in range(6): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + pressure_jj = np.array(lammps.variables[f"pressure{jj}"].value) / ( + constants.nktv2p + ) + expected_pressure_jj = -expected_v[idx_map, jj].sum(axis=0) / vol + np.testing.assert_allclose( + pressure_jj, expected_pressure_jj, atol=1e-8, rtol=1e-8 + ) + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + virial_jj = np.array(lammps.variables[f"virial{jj}"].value) / (constants.nktv2p) + np.testing.assert_allclose( + virial_jj, expected_v[idx_map, jj], atol=1e-8, rtol=1e-8 + ) + + +# --------------------------------------------------------------------------- +# Multi-rank: the spin+ZBL graph .pt2 carries the nested with-comm artifact +# (which includes the SFPG reverse-accumulate + broadcast exchange, issue +# #5906), so a 2-rank run must REPRODUCE the 1-rank result (energy, force +# and force_mag) rather than fail fast. +# --------------------------------------------------------------------------- + + +def _run_mpi_subprocess( + extra_args: list[str] | None = None, + nprocs: int = 2, + data_path: Path | None = None, + processors: str | None = None, + capture: bool = False, + timeout: float | None = None, +) -> dict: + """Invoke the graph-spin MPI runner under ``mpirun -n `` against + the spin+ZBL DPA4 graph ``.pt2``. + + Copied (module-global closure, not imported) from + ``test_lammps_dpa4_spin_graph_pt2.py``'s twin. With ``capture=True``, + return raw subprocess info (``returncode``, ``stdout``, ``stderr``, + ``timed_out``); every invocation is bounded by ``timeout`` (default + ``_MPI_DEFAULT_TIMEOUT``) so a should-fail-but-doesn't run cannot hang + the suite, and on expiry the WHOLE mpirun process group is SIGKILLed. + """ + if data_path is None: + data_path = data_file + if timeout is None: + timeout = _MPI_DEFAULT_TIMEOUT + with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: + out_path = f.name + try: + argv = [ + "mpirun", + "-n", + str(nprocs), + sys.executable, + str(mpi_runner), + str(data_path.resolve()), + str(pb_file.resolve()), + out_path, + ] + if processors is not None: + argv.extend(["--processors", processors]) + elif nprocs == 1: + argv.extend(["--processors", "1 1 1"]) + if extra_args: + argv.extend(extra_args) + proc = sp.Popen( + argv, + stdout=sp.PIPE if capture else None, + stderr=sp.PIPE if capture else None, + text=True, + start_new_session=True, + ) + try: + stdout, stderr = proc.communicate(timeout=timeout) + except sp.TimeoutExpired: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + stdout, stderr = proc.communicate() + if capture: + return { + "returncode": None, + "stdout": stdout or "", + "stderr": stderr or "", + "timed_out": True, + } + raise RuntimeError( + f"mpirun timed out after {timeout}s (process group killed); " + "a should-succeed MPI regression is deadlocked." + ) from None + if capture: + return { + "returncode": proc.returncode, + "stdout": stdout, + "stderr": stderr, + "timed_out": False, + } + if proc.returncode != 0: + raise sp.CalledProcessError(proc.returncode, argv) + with open(out_path) as fh: + lines = fh.read().strip().splitlines() + pe = float(lines[0]) + rows = np.array( + [list(map(float, line.split())) for line in lines[1:]], + dtype=np.float64, + ) + return {"pe": pe, "rows": rows} + finally: + if os.path.exists(out_path): + os.remove(out_path) + + +def _assert_mpi_matches_single_rank(single: dict, multi: dict) -> None: + """Compare a 2-rank MPI result against the 1-rank one: pe, per-atom + force and per-atom force_mag, all at rtol=atol=1e-10 (the sibling's MPI + tolerances). Both runs report LAMMPS's own ``fm`` (already scaled by + ``spin_norm / _HBAR_METAL`` inside pair_deepspin.cpp), so the scaling + cancels and the rows compare directly. + """ + # anti-vacuity: a degenerate fixture (all-zero forces) would make the + # comparison pass for the wrong reason. + assert np.abs(multi["rows"][:, :3]).max() > 1e-6, "forces are trivially zero" + assert np.abs(multi["rows"][:, 3:6]).max() > 1e-6, "force_mag is trivially zero" + + np.testing.assert_allclose( + multi["pe"], single["pe"], rtol=1e-10, atol=1e-10, err_msg="energy" + ) + np.testing.assert_allclose( + multi["rows"][:, :3], + single["rows"][:, :3], + rtol=1e-10, + atol=1e-10, + err_msg="force", + ) + np.testing.assert_allclose( + multi["rows"][:, 3:6], + single["rows"][:, 3:6], + rtol=1e-10, + atol=1e-10, + err_msg="force_mag", + ) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepspin_mpi_matches_single_rank() -> None: + """A 2-rank MPI run must reproduce the 1-rank result on the SAME archive. + + The spin+ZBL DPA4 graph ``.pt2`` carries the nested with-comm artifact, + so ``DeepSpinPTExpt::compute_inner`` drives a real domain-decomposed run + through ``run_model_graph_with_comm``: the per-block ghost FEATURE + refresh (including the SFPG partials exchange) rides ``border_op`` / + ``border_op_backward`` while ghost SPINS arrive via the LAMMPS ``sp`` + forward-comm. Both the conservative force and the MAGNETIC force must + be rank-count invariant. + + With this (unshifted) geometry the close Ni-Ni pair sits entirely on + one rank's subdomain, so this test pins the baseline multi-rank route; + the boundary-straddling bridging-zone pair is exercised by + ``test_pair_deepspin_mpi_close_pair_across_ranks`` below. + """ + single = _run_mpi_subprocess(nprocs=1, processors="1 1 1") + multi = _run_mpi_subprocess(nprocs=2, processors="2 1 1") + _assert_mpi_matches_single_rank(single, multi) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepspin_mpi_close_pair_across_ranks() -> None: + """2-rank vs 1-rank parity with the bridging-zone Ni-Ni pair straddling + the rank boundary. + + The shifted geometry places the ~0.978 A Ni-Ni pair at x = 6.93 / 6.19, + on OPPOSITE sides of the ``processors 2 1 1`` boundary at lx/2 = 6.5, + and 0.8 < 0.978 < 1.2 puts the pair inside the bridging transition zone + -- so each rank's per-node ``[log_eta, zero_count]`` SFPG partial for + the pair is rank-incomplete and correctness depends on the + reverse-accumulate + broadcast exchange inside the with-comm artifact + (issue #5906). + """ + # Static straddle guard: the two Ni atoms must sit on opposite sides of + # the lx/2 boundary. Without this the SFPG-exchange parity below is + # vacuous (the whole bridging-zone pair would be owned by one rank and + # the cross-rank completion would never be exercised). + half_lx = box[1] / 2.0 + x_lo, x_hi = sorted(coord_close_pair[:2, 0]) + assert x_lo < half_lx < x_hi, ( + f"close-pair geometry no longer straddles the processors 2 1 1 " + f"boundary: Ni x = {x_lo}, {x_hi}, lx/2 = {half_lx}" + ) + + single = _run_mpi_subprocess( + nprocs=1, processors="1 1 1", data_path=data_file_close_pair + ) + multi = _run_mpi_subprocess( + nprocs=2, processors="2 1 1", data_path=data_file_close_pair + ) + _assert_mpi_matches_single_rank(single, multi) diff --git a/source/tests/infer/gen_dpa4_spin_zbl.py b/source/tests/infer/gen_dpa4_spin_zbl.py index 09d1d4121d..5813ff0f20 100644 --- a/source/tests/infer/gen_dpa4_spin_zbl.py +++ b/source/tests/infer/gen_dpa4_spin_zbl.py @@ -25,14 +25,14 @@ ------------------------------------------------------- Bridging enables the descriptor's Source Freeze Propagation Gate, whose per-node ``eta_j = prod_{e: src_e = j} w_e`` folds a node's FULL outgoing-edge -set. Edges exist only for owned centres, so eta is incomplete on every rank -and no with-comm artifact may be exported. ``_check_metadata`` asserts BOTH -``has_comm_artifact is False`` AND that the nested -``model/extra/forward_lower_with_comm.pt2`` entry is absent (mirroring -``gen_dpa4_zbl.py``) -- an artifact appearing there would silently promise a -multi-rank capability the model cannot honour. Note this is the opposite of -the UNbridged native-spin fixture (``gen_dpa4_spin.py``), which does carry -the with-comm twin: bridging is what removes it. +set. Edges exist only for owned centres, so the per-node partials are +rank-incomplete; the with-comm artifact completes them with one +reverse-accumulate + forward-broadcast border exchange before the gate is +applied (issue #5906). ``_check_metadata`` asserts BOTH +``has_comm_artifact is True`` AND that the nested +``model/extra/forward_lower_with_comm.pt2`` entry is present (mirroring +``gen_dpa4_zbl.py``) -- the same contract as the UNbridged native-spin +fixture (``gen_dpa4_spin.py``). Generation mirrors ``gen_dpa4_spin_chgspin.py`` (the closest precedent): the dpmodel is built in-process from ``NATIVE_SPIN_CONFIG`` imported from @@ -354,24 +354,20 @@ def _check_metadata(pt2_path: str) -> None: f"the composition dropped the spin flag on the way to the freeze." ) assert md["use_spin"] == [True, False] - # Single-rank only -- see the module docstring (the bridging gate's eta is - # incomplete per rank). BOTH halves matter: the flag is what the C++ - # dispatch reads, the archive entry is what it would load. - assert md["has_comm_artifact"] is False, ( + # Multi-rank capable (issue #5906): the SFPG per-node partials are + # completed across ranks, so the with-comm artifact is embedded. BOTH + # halves matter: the flag is what the C++ dispatch reads, the archive + # entry is what it loads. + assert md["has_comm_artifact"] is True, ( f"{pt2_path}: metadata has_comm_artifact = " - f"{md.get('has_comm_artifact')!r}, expected False; a bridged model " - f"cannot support multi-rank message passing (its Source Freeze " - f"Propagation Gate folds each node's full outgoing-edge set, which no " - f"rank owns), so advertising one would promise a capability the model " - f"cannot honour." + f"{md.get('has_comm_artifact')!r}, expected True; the SFPG cross-rank " + f"completion (issue #5906) makes bridged native-spin models " + f"multi-rank capable." ) - assert "model/extra/forward_lower_with_comm.pt2" not in names, ( - f"{pt2_path}: a nested forward_lower_with_comm.pt2 was exported for a " - f"bridged model; see above -- the archive must not carry one." + assert "model/extra/forward_lower_with_comm.pt2" in names, ( + f"{pt2_path}: the nested forward_lower_with_comm.pt2 is missing from " + f"a bridged native-spin archive; multi-rank dispatch would fail." ) - # The descriptor still message-passes WITHIN a rank; it is only the - # cross-rank exchange that bridging forbids. This is the flag that makes - # the C++ side fail fast under mpirun instead of answering wrongly. assert md["has_message_passing"] is True for key in ("atom_energy", "energy", "force", "force_mag", "virial"): assert key in md["output_keys"] diff --git a/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py b/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py index 73e1a2f39d..b5442c4b99 100644 --- a/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py +++ b/source/tests/pt_expt/model/test_dpa4_zbl_parallel.py @@ -309,3 +309,119 @@ def test_freeze_embeds_with_comm_artifact(self, tmp_path) -> None: assert "model/extra/forward_lower_with_comm.pt2" in names assert meta["has_comm_artifact"] is True assert meta["lower_input_kind"] == "graph" + + +SPIN_ZBL_CONFIG = { + **copy.deepcopy(ZBL_CONFIG), + "spin": {"use_spin": [True, False], "scheme": "native"}, +} + + +def _make_bridged_spin_model(): + """Native-spin + ZBL composition with jittered residuals (see above).""" + model = get_model(copy.deepcopy(SPIN_ZBL_CONFIG)) + learned = model.atomic_model.models[0] + data = jitter_zero_arrays(learned.descriptor.serialize(), np.random.default_rng(99)) + learned.descriptor = DescrptDPA4.deserialize(data) + return model.to(torch.float64).to("cpu").eval() + + +class TestBridgedSpinGraphSelfComm: + """Issue #5906 Task 3: native spin + ZBL, same ladder as the non-spin + class. The gate exchange sits below the spin wrapper + (``NativeSpinEnergyModel`` re-classes the SAME composed atomic model), + so no spin-specific production change is expected -- these tests pin + that the machinery composes. + """ + + @pytest.fixture(autouse=True) + def _setup(self): + ensure_comm_registered() + self.model = _make_bridged_spin_model() + + def _spins(self, n: int) -> np.ndarray: + rng = np.random.default_rng(11) + sp = rng.normal(size=(n, 3)) + return sp / np.linalg.norm(sp, axis=-1, keepdims=True) + + def _run_folded(self, coord: np.ndarray): + atype = np.array([[0, 0, 1, 1]], dtype=np.int64) + box = (_L * np.eye(3, dtype=np.float64)).reshape(1, 3, 3) + graph = build_neighbor_graph(coord, atype, box, 4.0, canonicalize=True) + spin = torch.tensor(self._spins(4), dtype=torch.float64) + out = self.model.forward_common_lower_graph( + torch.tensor(atype.reshape(-1), dtype=torch.int64), + torch.as_tensor(np.asarray(graph.n_node), dtype=torch.int64), + torch.as_tensor(np.asarray(graph.n_node), dtype=torch.int64), + torch.as_tensor(np.asarray(graph.edge_index), dtype=torch.int64), + torch.as_tensor(np.asarray(graph.edge_vec), dtype=torch.float64), + torch.as_tensor(np.asarray(graph.edge_mask), dtype=torch.bool), + spin=spin, + ) + e = out["energy_redu"].detach().numpy().reshape(-1) + f = -out["energy_derv_r"].detach().numpy().reshape(-1, 3) + fm = -out["energy_derv_r_mag"].detach().numpy().reshape(-1, 3) + return e, f, fm + + def _run_self_comm(self, coord: np.ndarray): + ext_coord, ext_atype, nlist, mapping = _extended_quartet(coord) + gi = _unfolded_graph_inputs(ext_coord, ext_atype, nlist) + nall = ext_coord.shape[1] + # ghost spins mirror their owners (LAMMPS forwards ``sp``) + spin_ext = torch.tensor(self._spins(4)[mapping[0]], dtype=torch.float64) + keepalive: list = [] + comm_dict = _build_self_comm_dict( + nloc=4, + nghost=nall - 4, + sendlist_indices=mapping[0, 4:].astype(np.int32), + keepalive=keepalive, + ) + out = self.model.forward_common_lower_graph( + gi["atype"], + gi["n_node"], + gi["n_local"], + gi["edge_index"], + gi["edge_vec"], + gi["edge_mask"], + spin=spin_ext, + comm_dict=comm_dict, + ) + e = out["energy_redu"].detach().numpy().reshape(-1) + f = _fold_forces(-out["energy_derv_r"].detach().numpy().reshape(-1, 3), mapping) + fm = _fold_forces( + -out["energy_derv_r_mag"].detach().numpy().reshape(-1, 3), mapping + ) + return e, f, fm + + @pytest.mark.parametrize( + "gap", + [ + 0.4, # zero_count channel across the boundary + 1.0, # log_eta channel across the boundary + ], + ) + def test_self_comm_matches_folded_reference(self, gap: float) -> None: + coord = _close_pair_coords(gap) + e_ref, f_ref, fm_ref = self._run_folded(coord) + e_par, f_par, fm_par = self._run_self_comm(coord) + np.testing.assert_allclose(e_par, e_ref, rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(f_par, f_ref, rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(fm_par, fm_ref, rtol=1e-12, atol=1e-12) + + def test_freeze_embeds_with_comm_artifact(self, tmp_path) -> None: + """Freezing the bridged spin model embeds the nested artifact.""" + import json + import zipfile + + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + data = {"model": self.model.serialize()} + p = str(tmp_path / "m_dpa4_spin_zbl_graph.pt2") + deserialize_to_file(p, copy.deepcopy(data), lower_kind="graph") + with zipfile.ZipFile(p, "r") as zf: + names = zf.namelist() + meta = json.loads(zf.read("model/extra/metadata.json")) + assert "model/extra/forward_lower_with_comm.pt2" in names + assert meta["has_comm_artifact"] is True From d361cf340b86e2b7c415721339dcb8b4db2a011c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 15:37:34 +0800 Subject: [PATCH 13/26] test: align variant coverage with standard dpa4 Empty-rank behavior (zbl fail-fast twin; first test of the DeepSpin phantom path -- owned-empty ranks with ghosts phantom-pad and match), charge_spin via pair_deepspin, and dp-freeze default-CLI resolution for zbl and native-spin compositions (issue #5906 Task 12b). --- .../test_lammps_dpa4_chg_spin_deepspin_pt2.py | 341 ++++++++++++++++++ .../tests/test_lammps_dpa4_spin_graph_pt2.py | 106 +++++- source/lmp/tests/test_lammps_dpa4_zbl_pt2.py | 96 ++++- source/tests/pt_expt/test_dp_freeze.py | 56 +++ 4 files changed, 596 insertions(+), 3 deletions(-) create mode 100644 source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py diff --git a/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py b/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py new file mode 100644 index 0000000000..95b012f819 --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py @@ -0,0 +1,341 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test the LAMMPS ``charge_spin`` keyword through ``pair_style deepspin``. + +The sibling ``test_lammps_chg_spin_pt2.py`` exercises the keyword through +``pair_style deepmd`` (DPA3, non-spin); the DeepSpin ingestion seam is a +SEPARATE code path (``pair_deepspin.cpp``'s own ``charge_spin`` keyword +parse feeding ``DeepSpin::compute(..., charge_spin)``, see +``source/lmp/pair_deepspin.cpp``), so a regression confined to it was +invisible to that file -- every pre-existing spin fixture has +``dim_chg_spin == 0``, making the whole argument inert. This file mirrors +the sibling's three-test structure (default run, explicit run, sensitivity) +against the COMBINED native-spin + charge-spin FiLM DPA4 archive +``deeppot_dpa4_spin_chgspin.pt2`` (issue #5906 Task 12b: DPA4-variant +behaviour must align with what the C++ gtest +``source/api_cc/tests/test_deepspin_dpa4_chgspin_ptexpt.cc`` already proves +for the C API -- the LAMMPS pair style is a distinct consumer). + +The archive and its stored ``default_chg_spin = [0.0, 1.0]`` come from +``source/tests/infer/gen_dpa4_spin_chgspin.py``; the explicit probe +``charge_spin 1.0 2.0`` matches the generator's ``_EXPLICIT_CHG_SPIN`` +(the embedding is CATEGORICAL, so the two probes land on distinct rows in +BOTH components -- see the generator's module docstring). + +Reference values are computed LIVE at test-setup time via +``deepmd.infer.DeepPot.eval`` on the archive itself (mirroring +``test_lammps_dpa4_spin_graph_pt2.py``'s ``_compute_expected``, which +explains the reasoning in full) rather than read from the generator's +``.expected`` sidecar: the sidecar's evaluation uses a 6x6x6 A cell whose +edge length equals DPA4's LAMMPS ghost cutoff exactly +(rcut(4.0)+skin(2.0)=6.0), which is not a safe geometry for a periodic +LAMMPS run. This module keeps the generator's 6-atom NiO geometry and +spins verbatim, in a 13x13x13 A box instead (the same box-swap the ZBL twin +``test_lammps_dpa4_zbl_pt2.py`` applies to its generator geometry). +""" + +import json +import os +import subprocess as sp +import sys +import textwrap +from pathlib import ( + Path, +) + +import numpy as np +import pytest +from lammps import ( + PyLammps, +) +from write_lmp_data import ( + write_lmp_data_spin, +) + +pb_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4_spin_chgspin.pt2" +) +data_file = Path(__file__).parent / "data_dpa4_chg_spin_deepspin_pt2.lmp" + +# 6-atom NiO system: coordinates, types and spins verbatim from +# ``gen_dpa4_spin_chgspin.py`` (3 spin-active Ni + 3 O; the O spins are +# deliberately nonzero there -- the model's own descriptor gating must zero +# the non-spin rows internally), in a 13x13x13 A box instead of the +# generator's 6x6x6 (see the module docstring). +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [1.0, 1.0, 1.0], + [3.2, 1.4, 1.1], + [1.3, 1.8, 1.0], + [0.4, 1.2, 1.6], + [3.6, 2.0, 1.3], + [3.4, 0.7, 1.7], + ] +) +spin = np.array( + [ + [0.11, 0.05, -0.02], + [-0.07, 0.09, 0.03], + [0.02, -0.06, 0.08], + [0.01, -0.01, 0.02], + [-0.02, 0.03, -0.01], + [0.015, 0.02, -0.03], + ] +) +# Model ``type_map`` is ["Ni", "O"]; the generator's atype [0,0,0,1,1,1] +# -> LAMMPS types [1,1,1,2,2,2] under identity ``pair_coeff * *``. +type_NiO = np.array([1, 1, 1, 2, 2, 2]) + +# Explicit runtime probe, matching the generator's ``_EXPLICIT_CHG_SPIN`` +# (distinct from the stored default [0.0, 1.0] in BOTH categorical +# components, so neither component alone can explain a response). +_EXPLICIT_CHG_SPIN = [1.0, 2.0] + +# LAMMPS's ``fm`` (what ``compute property/atom fmx fmy fmz`` reports) is +# NOT the raw DeepEval force_mag: pair_deepspin.cpp scales it by +# ``spin_norm / hbar`` per atom (metal-units ``hbar = 6.5821191e-04``; same +# convention as test_lammps_dpa4_spin_graph_pt2.py, which documents it). +_HBAR_METAL = 6.5821191e-04 + +# Reference values (energy / force / force_mag, default and explicit +# charge_spin), populated by ``_compute_expected`` in ``setup_module``. +expected_e_default = None +expected_f_default = None +expected_fm_default = None +expected_e_explicit = None +expected_f_explicit = None +expected_fm_explicit = None + + +def _cell_from_lammps_box(lmp_box: np.ndarray) -> np.ndarray: + """Convert a LAMMPS ``xlo xhi ylo yhi zlo zhi xy xz yz`` box spec to a + flat, row-major 3x3 cell matrix (deepmd's ``box`` convention). + """ + xlo, xhi, ylo, yhi, zlo, zhi, xy, xz, yz = lmp_box + return np.array([xhi - xlo, 0.0, 0.0, xy, yhi - ylo, 0.0, xz, yz, zhi - zlo]) + + +def _compute_expected() -> None: + """Load ``deeppot_dpa4_spin_chgspin.pt2`` via ``DeepPot`` and evaluate + the module's fixed 6-atom NiO system, once with NO ``charge_spin`` + (stored default) and once with the explicit probe. + + Runs in a subprocess to avoid importing ``deepmd`` in the LAMMPS test + process (the LAMMPS plugin already loads ``libdeepmd_op_pt.so`` at the + C++ level, and importing the Python package on top of that can + segfault) -- the same precaution as ``test_lammps_dpa4_spin_graph_pt2.py``. + """ + global expected_e_default, expected_f_default, expected_fm_default + global expected_e_explicit, expected_f_explicit, expected_fm_explicit + + cell = _cell_from_lammps_box(box) + atype = (type_NiO - 1).tolist() # LAMMPS 1-based -> deepmd 0-based (Ni=0, O=1) + + # The archive lives in ``source/tests/infer`` next to ``gen_common.py``, + # whose ``load_custom_ops()`` loads the build-tree ``libdeepmd_op_pt.so`` + # (registering ``deepmd::edge_force_virial``, which graph ``.pt2`` + # inference needs); ``import deepmd.pt`` alone only loads the op library + # from SHARED_LIB_DIR, which the build-test env does not populate. + infer_dir = str(pb_file.resolve().parent) + script = textwrap.dedent(f"""\ + import json + import sys + import numpy as np + + sys.path.insert(0, {infer_dir!r}) + import deepmd.pt # noqa: F401 (triggers the base op-library load) + from gen_common import load_custom_ops + + load_custom_ops() + from deepmd.infer import DeepPot + + dp = DeepPot({str(pb_file.resolve())!r}) + assert dp.deep_eval.get_dim_chg_spin() == 2 + out = {{}} + for label, chg_spin in ( + ("default", None), + ("explicit", {_EXPLICIT_CHG_SPIN!r}), + ): + kwargs = {{}} + if chg_spin is not None: + kwargs["charge_spin"] = np.array([chg_spin], dtype=np.float64) + e, f, v, ae, av, fm, mm = dp.eval( + np.array({coord.tolist()!r}).reshape(1, -1, 3), + np.array({cell.tolist()!r}).reshape(1, 9), + {atype!r}, + atomic=True, + spin=np.array({spin.tolist()!r}).reshape(1, -1, 3), + **kwargs, + ) + out[label] = {{ + "e": float(e[0, 0]), + "f": np.asarray(f[0]).tolist(), + "fm": np.asarray(fm[0]).tolist(), + }} + print(json.dumps(out)) + """) + proc = sp.run([sys.executable, "-c", script], capture_output=True, text=True) + if proc.returncode != 0: + raise RuntimeError(f"Failed to compute expected values:\n{proc.stderr}") + result = json.loads(proc.stdout.strip()) + + # Raw DeepEval force_mag (dE/dspin), scaled by LAMMPS's own + # spin_norm / hbar unit convention (see ``_HBAR_METAL`` above) before + # comparison. + spin_norm_scale = (np.linalg.norm(spin, axis=1) / _HBAR_METAL)[:, None] + + expected_e_default = result["default"]["e"] + expected_f_default = np.array(result["default"]["f"]) + expected_fm_default = np.array(result["default"]["fm"]) * spin_norm_scale + expected_e_explicit = result["explicit"]["e"] + expected_f_explicit = np.array(result["explicit"]["f"]) + expected_fm_explicit = np.array(result["explicit"]["fm"]) * spin_norm_scale + + # Anti-vacuity, checked once here so every test below is known to compare + # against a non-degenerate reference: the explicit probe must MOVE the + # energy (the generator asserts the same at generation time; re-asserting + # on THIS geometry keeps the sensitivity test below meaningful), and the + # charge-spin FiLM must not have killed the spin response. + assert abs(expected_e_explicit - expected_e_default) > 1e-6, ( + f"charge_spin={_EXPLICIT_CHG_SPIN} left the energy unchanged vs the " + f"stored default ({expected_e_default:.18e} vs " + f"{expected_e_explicit:.18e}); the FiLM conditioning is not reaching " + f"the forward on this geometry, so the sensitivity test is vacuous." + ) + assert np.max(np.abs(expected_fm_default[:3])) > 1e-6, ( + "expected non-trivial force_mag on the spin-active (Ni) atoms; the " + "fixture would be vacuous for the spin leaf." + ) + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip("Skip test because PyTorch support is not enabled.") + if not pb_file.exists(): + pytest.skip( + "deeppot_dpa4_spin_chgspin.pt2 not found (run " + "source/tests/infer/gen_dpa4_spin_chgspin.py)." + ) + _compute_expected() + write_lmp_data_spin(box, coord, spin, type_NiO, data_file) + + +def teardown_module() -> None: + if data_file.exists(): + os.remove(data_file) + + +def _lammps(data_file, units="metal") -> PyLammps: + """Standard DeepSpin LAMMPS system, plus ``atom_modify map yes``. + + Same setup as ``test_lammps_dpa4_spin_graph_pt2.py``: the native-spin + DPA4 GRAPH ``.pt2`` needs the LAMMPS atom-map to resolve ghost-atom + indices to local owners for single-rank inference. + """ + if units != "metal": + raise ValueError("units for spin should be metal") + + lammps = PyLammps() + lammps.units(units) + lammps.boundary("p p p") + lammps.atom_style("spin") + lammps.atom_modify("map yes") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + lammps.mass("1 58") # Ni + lammps.mass("2 16") # O + lammps.timestep(0.0005) + lammps.fix("1 all nve") + return lammps + + +@pytest.fixture +def lammps(): + lmp = _lammps(data_file=data_file) + yield lmp + lmp.close() + + +def _gather_force_mag(lammps: PyLammps, natoms: int) -> np.ndarray: + """Extract per-atom force_mag in atom-id order via + ``compute property/atom fmx fmy fmz`` + ``gather`` (LAMMPS does not + expose ``fm`` through the legacy ``extract``/``gather_atoms`` registry; + see ``test_lammps_dpa4_spin_graph_pt2.py``). + """ + fm_global = lammps.lmp.gather("c_fmprop", 1, 3) + return np.array(fm_global, dtype=np.float64).reshape(natoms, 3) + + +def _assert_run0_matches( + lammps: PyLammps, + e_ref: float, + f_ref: np.ndarray, + fm_ref: np.ndarray, +) -> None: + """Run 0 steps and compare pe / force / force_mag to the given reference + (both sides run the SAME compiled artifact, so ``atol=1e-8`` is a + cross-consumer bound, not a cross-backend one -- same rationale as the + sibling DPA4 LAMMPS tests). + """ + natoms = coord.shape[0] + lammps.compute("fmprop all property/atom fmx fmy fmz") + lammps.run(0) + + assert lammps.eval("pe") == pytest.approx(e_ref, rel=1e-10) + + forces = np.array( + [lammps.atoms[ii].force for ii in range(natoms)], dtype=np.float64 + ) + ids = np.array([lammps.atoms[ii].id for ii in range(natoms)]) + np.testing.assert_allclose(forces, f_ref[ids - 1], atol=1e-8, rtol=0) + + force_mag = _gather_force_mag(lammps, natoms) + np.testing.assert_allclose(force_mag, fm_ref, atol=1e-8, rtol=0) + # Native-spin design invariant: force_mag on the non-spin (O) atoms must + # be exactly zero -- the model's own type gating, not the spin values, + # decides (the O spins in this fixture are deliberately nonzero). + np.testing.assert_array_equal(force_mag[3:], np.zeros((3, 3))) + + +def test_pair_deepspin_charge_spin_default(lammps) -> None: + """No charge_spin keyword -> the model's stored default_chg_spin is used + (the DeepSpin twin of the backward-compatibility contract the C++ gtest + pins for an EMPTY runtime charge_spin). + """ + lammps.pair_style(f"deepspin {pb_file.resolve()}") + lammps.pair_coeff("* *") + _assert_run0_matches( + lammps, expected_e_default, expected_f_default, expected_fm_default + ) + lammps.run(1) + + +def test_pair_deepspin_charge_spin_explicit(lammps) -> None: + """Explicit ``charge_spin`` keyword is parsed by pair_deepspin and + threaded through DeepSpin to the model (energy, force AND force_mag -- + the spin-only output -- must all follow the explicit conditioning). + """ + cs = " ".join(str(v) for v in _EXPLICIT_CHG_SPIN) + lammps.pair_style(f"deepspin {pb_file.resolve()} charge_spin {cs}") + lammps.pair_coeff("* *") + _assert_run0_matches( + lammps, expected_e_explicit, expected_f_explicit, expected_fm_explicit + ) + lammps.run(1) + + +def test_charge_spin_changes_result(lammps) -> None: + """Different charge_spin must give a different energy (keyword takes + effect through pair_deepspin; ``_compute_expected`` already pinned that + the two references differ, so this catches the keyword being silently + dropped on the LAMMPS side). + """ + cs = " ".join(str(v) for v in _EXPLICIT_CHG_SPIN) + lammps.pair_style(f"deepspin {pb_file.resolve()} charge_spin {cs}") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") != pytest.approx(expected_e_default) diff --git a/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py b/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py index 96081f5f0d..2603e1afb8 100644 --- a/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py @@ -69,6 +69,27 @@ / "deeppot_dpa4_spin_graph.pt2" ) data_file = Path(__file__).parent / "data_dpa4_spin_graph_pt2.lmp" +# Wide-box, 3-way x-split variant for the OWNED-EMPTY-rank MPI corner +# (``processors 3 1 1``), adapted from ``test_lammps_dpa4_graph_pt2.py``'s +# ``data_file_empty_rank`` construction to this file's 4-atom spin system +# (issue #5906 Task 12b) -- with one load-bearing difference. The DeepSpin +# phantom path (PR #5485, ``DeepSpinPTExpt.cc``) engages when a rank owns +# ZERO local atoms but still holds ghosts (``nloc_real == 0 && +# nall_real > 0``); a GENUINELY empty rank (zero owned AND zero ghost) is +# instead rejected by the same collective fail-fast preflight as the energy +# route (``DeepSpinPTExpt.cc``'s "zero owned+ghost atoms" throw), so the +# genuinely-empty construction of the energy twin cannot exercise the +# phantom logic. This fixture therefore shifts the Ni pair to +# x ~= 26.8/26.1, within DPA4's ghost cutoff (rcut(4.0)+skin(2.0)=6.0) of +# the x=30 slab boundary of a [0, 90] box. With 3 even x-slabs of width +# 30: rank 0 owns [0, 30) (all 4 atoms), rank 1 ([30, 60)) owns NOTHING but +# receives ghosts of the two Ni atoms (3.2 and 3.9 from its lower +# boundary), and rank 2 ([60, 90)) owns nothing but receives periodic +# ghosts of the x < 6 O atoms (x=3.51 and x=4.27, wrapped around the box's +# x=90/x=0 seam) -- BOTH atom-less ranks carry ghosts, so both take the +# phantom path and none trips the genuinely-empty fail-fast. The shifted +# coordinates (``coord_empty_rank``) are defined below, after ``coord``. +data_file_empty_rank = Path(__file__).parent / "data_dpa4_spin_graph_pt2_empty_rank.lmp" # The MPI runner is graph-spin-specific (no aparam / no NULL-type # extras, unlike run_mpi_pair_deepmd_spin_dpa3_pt2.py's virtual-atom-scheme # runner): the native-spin DPA4 fixture takes no fparam/aparam. @@ -100,6 +121,19 @@ ) type_NiO = np.array([1, 1, 2, 2]) +# Owned-empty-rank variant of ``coord`` (see the comment above +# ``data_file_empty_rank``): the two Ni atoms shift to x ~= 26.8/26.1 so +# rank 1 of the ``processors 3 1 1`` split owns nothing but holds their +# ghosts; the O atoms stay at x < 6 so rank 2 holds their periodic-seam +# ghosts. The Ni-Ni and O-O pair geometries are internally unchanged +# (rigid x-shift of the Ni pair only), and in the wide [0, 90] box the two +# pairs sit far beyond rcut(4.0) of each other with or without the shift, +# so each pair still interacts internally and the system stays +# non-degenerate for the anti-vacuity checks below. +_EMPTY_RANK_NI_X_SHIFT = 14.0 +coord_empty_rank = coord.copy() +coord_empty_rank[:2, 0] += _EMPTY_RANK_NI_X_SHIFT + # LAMMPS's ``fm`` (what ``compute property/atom fmx fmy fmz`` reports) is # NOT the raw DeepEval force_mag: pair_deepspin.cpp scales it by # ``spin_norm / hbar`` per atom (metal-units ``hbar = 6.5821191e-04``, see @@ -223,11 +257,16 @@ def setup_module() -> None: pytest.skip("deeppot_dpa4_spin_graph.pt2 not found") _compute_expected() write_lmp_data_spin(box, coord, spin, type_NiO, data_file) + box_empty_rank = np.array([0, 90, 0, 13, 0, 13, 0, 0, 0]) + write_lmp_data_spin( + box_empty_rank, coord_empty_rank, spin, type_NiO, data_file_empty_rank + ) def teardown_module() -> None: - if data_file.exists(): - os.remove(data_file) + for f in (data_file, data_file_empty_rank): + if f.exists(): + os.remove(f) def _lammps(data_file, units="metal") -> PyLammps: @@ -501,3 +540,66 @@ def test_pair_deepspin_mpi_matches_single_rank() -> None: atol=1e-10, err_msg="force_mag", ) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepspin_mpi_empty_rank_phantom_pads_and_matches() -> None: + """A rank that owns ZERO local atoms (but holds ghosts) must SUCCEED + through the DeepSpin phantom path -- the first coverage of that path + (issue #5906 Task 12b). + + This is the deliberate divergence from the energy route: where + ``DeepPotPTExpt`` has no owned-empty special case, + ``DeepSpinPTExpt::compute_inner`` phantom-pads the owned-empty rank + route-agnostically -- it prepends 2 phantom atoms with an empty nlist + row (contributing exactly zero energy/force/virial), because the + inductor specialization assumes ``nloc >= 2`` (PR #5485). A crash, + fail-fast exit, or hang here would therefore be a real regression of + the phantom logic, not the expected behaviour. (A GENUINELY empty + rank -- zero owned AND zero ghost -- is a different corner: it is + rejected by the same collective "zero owned+ghost atoms" preflight as + the energy route, because the phantom path requires ``nall_real > 0``; + see the fixture comment.) + + ``data_file_empty_rank`` (3-way x-split, ``processors 3 1 1``) was + verified (see the module-level comment above the fixture) to put BOTH + non-first ranks in the owned-empty-with-ghosts state; the 1-rank run on + the SAME wide-box data file is the same-archive reference. Energy, + conservative force AND magnetic force must all be rank-count invariant + at the file's MPI tolerances -- force_mag only exists on this route, so + comparing it is what proves the spin leaf survives phantom padding. + """ + single = _run_mpi_subprocess( + nprocs=1, processors="1 1 1", data_path=data_file_empty_rank + ) + multi = _run_mpi_subprocess( + nprocs=3, processors="3 1 1", data_path=data_file_empty_rank + ) + + # anti-vacuity: a degenerate fixture (all-zero forces) would make the + # comparison pass for the wrong reason. + assert np.abs(multi["rows"][:, :3]).max() > 1e-6, "forces are trivially zero" + assert np.abs(multi["rows"][:, 3:6]).max() > 1e-6, "force_mag is trivially zero" + + np.testing.assert_allclose( + multi["pe"], single["pe"], rtol=1e-10, atol=1e-10, err_msg="energy" + ) + np.testing.assert_allclose( + multi["rows"][:, :3], + single["rows"][:, :3], + rtol=1e-10, + atol=1e-10, + err_msg="force", + ) + np.testing.assert_allclose( + multi["rows"][:, 3:6], + single["rows"][:, 3:6], + rtol=1e-10, + atol=1e-10, + err_msg="force_mag", + ) diff --git a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py index 18e9cf75fe..9ff483986e 100644 --- a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py @@ -112,6 +112,21 @@ coord_close_pair = coord + np.array([_CLOSE_PAIR_X_SHIFT, 0.0, 0.0]) data_file_close_pair = Path(__file__).parent / "data_dpa4_zbl_close_pair_pt2.lmp" +# Wide-box, 3-way x-split variant for the genuinely-empty-rank MPI corner +# (``processors 3 1 1``), same construction as +# ``test_lammps_dpa4_graph_pt2.py``'s ``data_file_empty_rank`` fixture (issue +# #5906 Task 12b: the ZBL composition must fail-fast exactly like standard +# DPA4). The 6 NiO atoms stay at x in [0.4, 3.6] near the left edge of a +# [0, 90] box. With 3 even x-slabs of width 30, rank 0 owns [0, 30) (all +# atoms), rank 2 owns [60, 90) (empty of local atoms but picks up periodic +# ghosts of the x < 6 atoms wrapped around the box's x=90/x=0 seam, all +# within DPA4's ghost cutoff rcut(4.0)+skin(2.0)=6.0), and rank 1 (the +# MIDDLE slab, [30, 60)) borders neither the real atoms directly (nearest +# real atom at distance 30-3.6 = 26.4 > 6) nor the periodic seam -- so +# rank 1 is the genuinely empty rank (zero owned AND zero ghost atoms) this +# fixture is built to produce. +data_file_empty_rank = Path(__file__).parent / "data_dpa4_zbl_pt2_empty_rank.lmp" + # Reference values, populated by ``_compute_expected`` in ``setup_module``. expected_e = None expected_ae = None @@ -204,10 +219,12 @@ def setup_module() -> None: _compute_expected() write_lmp_data(box, coord, type_NiO, data_file) write_lmp_data(box, coord_close_pair, type_NiO, data_file_close_pair) + box_empty_rank = np.array([0, 90, 0, 13, 0, 13, 0, 0, 0]) + write_lmp_data(box_empty_rank, coord, type_NiO, data_file_empty_rank) def teardown_module() -> None: - for f in (data_file, data_file_close_pair): + for f in (data_file, data_file_close_pair, data_file_empty_rank): if f.exists(): os.remove(f) @@ -317,12 +334,20 @@ def _run_mpi_subprocess( processors: str, timeout: float = _MPI_DEFAULT_TIMEOUT, data_path: Path | None = None, + capture: bool = False, ) -> dict: """Run the (backend-agnostic) DPA3 MPI runner against the bridged archive and return the parsed ``{"pe", "forces", "virials"}`` output. Always bounded: on expiry the WHOLE mpirun process group is SIGKILLed (killing only mpirun can leave orphaned ranks blocking in a collective). + + With ``capture=True`` (mirroring ``test_lammps_dpa4_graph_pt2.py``'s + helper), return raw subprocess info (``returncode``, ``stdout``, + ``stderr``, ``timed_out``) instead of parsed output -- used by the + fail-fast test below; a timeout there returns ``timed_out=True`` with + ``returncode=None`` for the caller to assert on, and a nonzero exit is + returned rather than raised. """ if data_path is None: data_path = data_file @@ -349,12 +374,26 @@ def _run_mpi_subprocess( except sp.TimeoutExpired: os.killpg(os.getpgid(proc.pid), signal.SIGKILL) stdout, stderr = proc.communicate() + if capture: + return { + "returncode": None, + "stdout": stdout or "", + "stderr": stderr or "", + "timed_out": True, + } raise RuntimeError( f"mpirun timed out after {timeout}s (process group killed); " "a should-succeed MPI regression is deadlocked.\n" f"stdout:\n{(stdout or '')[-2000:]}\n" f"stderr:\n{(stderr or '')[-2000:]}" ) from None + if capture: + return { + "returncode": proc.returncode, + "stdout": stdout, + "stderr": stderr, + "timed_out": False, + } if proc.returncode != 0: raise RuntimeError( f"mpirun exited {proc.returncode}.\n" @@ -409,3 +448,58 @@ def test_pair_deepmd_mpi_dpa4_zbl_close_pair_parity() -> None: # relative component absorbs CUDA atomic-scatter ordering noise without # loosening the CPU-exact case. np.testing.assert_allclose(par["virials"], ref["virials"], atol=1e-8, rtol=1e-8) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepmd_mpi_dpa4_zbl_empty_rank_does_not_silently_succeed() -> None: + """A genuinely empty rank (zero owned AND zero ghost atoms) under the + message-passing with-comm graph route must NOT silently produce + wrong-but-plausible numbers -- for the ZBL COMPOSITION exactly as for + standard DPA4 (issue #5906 Task 12b: variant/standard alignment). + + The bridged archive routes through the SAME model-agnostic C++ guard as + the plain-DPA4 twin (``test_lammps_dpa4_graph_pt2.py``'s + ``test_pair_deepmd_mpi_dpa4_graph_empty_rank_does_not_silently_succeed``, + which documents the mechanism in full): every rank preflights a + communicator-wide min-reduction of its node count + (``deepmd_export::allreduce_min_int``) BEFORE entering the per-layer + ``border_op`` collectives, so the non-empty peers detect the empty rank + and throw the documented error instead of blocking forever. A timeout + is therefore a FAILURE of this test, and the documented error message + must appear on a nonzero exit. What this pins for the composition + specifically: the linear/ZBL wrapping must not swallow or bypass the + guard on its way into the graph forward. + + ``data_file_empty_rank`` (3-way x-split, ``processors 3 1 1``) was + verified (see the module-level comment above the fixture) to put the + MIDDLE rank in a genuinely empty state, using DPA4's own ghost cutoff + (rcut(4.0)+skin(2.0)=6.0). + """ + out = _run_mpi_subprocess( + nprocs=3, + processors="3 1 1", + data_path=data_file_empty_rank, + capture=True, + timeout=120, + ) + assert not out["timed_out"], ( + "Multi-rank ZBL-bridged graph run with an empty rank timed out " + "instead of failing promptly: the collective empty-rank preflight " + "(allreduce_min_int) must make every rank throw BEFORE the " + "per-layer border_op collectives." + ) + assert out["returncode"] != 0, ( + "Expected the multi-rank message-passing ZBL-bridged run to fail " + "loudly on a genuinely empty rank, but it exited 0.\n" + f"stdout:\n{out['stdout'][-2000:]}\nstderr:\n{out['stderr'][-2000:]}" + ) + combined = out["stdout"] + out["stderr"] + assert "zero owned+ghost atoms" in combined, ( + "Expected the documented fail-loud message ('zero owned+ghost " + f"atoms'), got:\n{combined[-2000:]}" + ) diff --git a/source/tests/pt_expt/test_dp_freeze.py b/source/tests/pt_expt/test_dp_freeze.py index b10e22d51a..cc7a0bd265 100644 --- a/source/tests/pt_expt/test_dp_freeze.py +++ b/source/tests/pt_expt/test_dp_freeze.py @@ -8,6 +8,7 @@ deepcopy, ) +import pytest import torch from deepmd.pt_expt.entrypoints.main import ( @@ -334,5 +335,60 @@ def test_nonspin_model_rejects_spin(self) -> None: dp.eval(coord, box, atype, spin=spin) +@pytest.mark.skipif( + os.environ.get("CI") == "true", + reason="AOTInductor compile is slow (minutes); run locally only by default.", +) +@pytest.mark.parametrize( + "add_native_spin", + [ + pytest.param(False, id="zbl"), # ZBL composition (Linear over [DPA4, ZBL]) + pytest.param(True, id="native_spin"), # native-spin variant of the same + ], +) +def test_dpa4_variant_default_freeze_graph_with_comm(tmp_path, add_native_spin) -> None: + """The default ``dp freeze`` invocation on a DPA4 VARIANT yields a + with-comm graph ``.pt2`` -- aligned with standard DPA4 (issue #5906 + Task 12b). + + No ``--lower-kind`` flag means ``freeze()``'s ``lower_kind="nlist"`` + default, which the nlist->graph auto-override in + ``deepmd/pt_expt/entrypoints/main.py`` must resolve to the graph lower + for every graph-capable model -- for the ZBL composition and the + native-spin variant exactly as for standard DPA4 (the plain native-spin + case is already pinned by ``test_native_spin_default_freeze_routes_to_ + graph`` in ``model/test_dpa4_export.py``; these two cases pin the + variants). Beyond the routing, the frozen archive must carry the nested + with-comm artifact (``has_comm_artifact is True``): a variant that + silently dropped it would freeze fine, pass every single-rank test, and + only fail (or worse, silently mis-answer) on multi-rank LAMMPS. + """ + import json + import zipfile + + from .model.test_zbl_bridging import ( + ZBL_CONFIG, + ) + + config = deepcopy(ZBL_CONFIG) + if add_native_spin: + config["spin"] = {"use_spin": [True, False], "scheme": "native"} + + model = get_model(deepcopy(config)) + wrapper = ModelWrapper(model, model_params=deepcopy(config)) + ckpt = tmp_path / "model.pt" + torch.save({"model": wrapper.state_dict()}, ckpt) + + output = tmp_path / "frozen_dpa4_variant" # suffixless: default CLI form + freeze(model=str(ckpt), output=str(output)) + + pt2 = output.with_suffix(".pt2") + assert pt2.exists(), "default suffix must follow the resolved graph kind" + with zipfile.ZipFile(pt2) as zf: + metadata = json.loads(zf.read("model/extra/metadata.json")) + assert metadata["lower_input_kind"] == "graph" + assert metadata["has_comm_artifact"] is True + + if __name__ == "__main__": unittest.main() From 36d6c082f326e419087bc9125ea857a071727e09 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 30 Jul 2026 15:51:06 +0800 Subject: [PATCH 14/26] docs(dpa4): ZBL bridging and native-spin combinations are multi-rank The SFPG per-node partials are completed across ranks by one reverse-accumulate + forward-broadcast border exchange of an (N,2) tensor per forward pass (issue #5906); also sweeps the stale pre-#5884 native-spin single-rank claims. --- doc/model/dpa4.md | 77 ++++++++++++++++++++--------------------------- 1 file changed, 32 insertions(+), 45 deletions(-) diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 66ff3363da..7958fd2caa 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -454,18 +454,20 @@ route, and `--lower-kind` selects between the pt_expt lowers only. Which models lose multi-rank, on either route: -- **ZBL zone bridging.** The Source Freeze Propagation gate folds each node's - full *outgoing*-edge set, which no single rank observes for ghost owners, so - `supports_edge_parallel()` is `False` and no with-comm artifact is emitted. - Bridged archives are single-rank; a multi-rank run fails with a clear error - rather than silently dropping the exchange. - **The deepspin (virtual-atom) spin scheme**, which overrides the export ABI to `nlist` because it expands virtual atoms inside the graph. -Native spin (`scheme: "native"`) and charge/spin conditioning are *not* in -that list: native spin reuses the `edge_vec` interface on PT and the -NeighborGraph lower on pt_expt, and both support multi-rank. The remainder of -this subsection describes the multi-GPU launch recipe. +**ZBL zone bridging is multi-rank capable** (including combined with native +spin). The Source Freeze Propagation gate folds each node's full +*outgoing*-edge set, which no single rank observes for ghost owners; the +with-comm artifact completes the gate's per-node `[log eta, zero count]` +partials across ranks with one reverse-accumulate plus one forward-broadcast +border exchange of an `(N, 2)` tensor per forward pass — negligible next to +the per-block feature exchanges that already run. Native spin +(`scheme: "native"`) and charge/spin conditioning likewise support +multi-rank: native spin reuses the `edge_vec` interface on PT and the +NeighborGraph lower on pt_expt. The remainder of this subsection describes +the multi-GPU launch recipe. ::: The exported `.pt2` runs across multiple GPUs in LAMMPS using MPI domain @@ -536,8 +538,8 @@ graph-capable model is always frozen to the graph lower in any case, since the dense lower is deprecated in the pt_expt backend. See [Native spin (magnetic)](#native-spin-magnetic) below. Unlike the dense route (see [Multi-GPU (MPI) -inference](#multi-gpu-mpi-inference) above), a graph-frozen `.pt2` **of a -plain-energy (non-spin) model** embeds a with-comm AOTInductor artifact and +inference](#multi-gpu-mpi-inference) above), a graph-frozen `.pt2` embeds a +with-comm AOTInductor artifact and supports multi-rank LAMMPS: each block's cross-rank ghost-feature exchange runs through the `border_op` MPI path once per interaction block, the same mechanism used by DPA-2's graph route (see the "Graph-native inference route @@ -549,8 +551,8 @@ exchange. Pick a domain decomposition that keeps every rank non-empty, or use the dense route, which has no such restriction (but is single-rank only, as noted above). -**Native-spin graph `.pt2` archives are the exception: they carry no -with-comm artifact and are single-rank only** -- see [Native spin +Native-spin graph `.pt2` archives participate too: they carry the with-comm +artifact and support multi-rank LAMMPS -- see [Native spin (magnetic)](#native-spin-magnetic) below. ### Native spin (magnetic) @@ -573,17 +575,12 @@ inference](#multi-gpu-mpi-inference). a native-spin descriptor has only the graph lower. `--lower-kind auto` (the default) resolves to `graph`; `--lower-kind nlist` is not a valid option for a native-spin model. -- **Single-rank only.** The frozen archive's `has_comm_artifact` metadata is - `false` for native-spin models (no ghost-spin cross-rank exchange is - implemented), so a multi-rank LAMMPS run fails fast with an explicit error - at the first force evaluation, mirroring the dense route's single-rank - restriction described in [Multi-GPU (MPI) - inference](#multi-gpu-mpi-inference). Run native-spin models on a single - MPI rank (a single GPU, or CPU without `mpirun`). -- **Spin is per local atom.** The `spin` input is `(nframes, nloc, 3)` -- - one vector per *local* atom, not per ghost/extended atom (`nall`); there is - no ghost-spin exchange to populate ghost spins across a rank boundary, - consistent with the single-rank restriction above. +- **Multi-rank capable.** The frozen archive embeds the with-comm artifact + (`has_comm_artifact` is `true`), so multi-rank LAMMPS works exactly as + described in [Multi-GPU (MPI) inference](#multi-gpu-mpi-inference): the + per-block ghost node features ride `border_op`, and ghost spins arrive + through the LAMMPS `sp` forward communication -- spin itself needs no + extra cross-rank exchange. - **The magnetic force is a second energy gradient.** As in the general native-scheme convention (see [Spin](#spin) above), `force_mag = -\partial E/\partial\mathbf{s}`, computed by pt_expt as a @@ -595,18 +592,9 @@ inference](#multi-gpu-mpi-inference). route. The following combinations are **not yet supported** on the native-spin -graph route (follow-up work): - -- **Multi-rank inference.** Ghost-spin cross-rank exchange (analogous to the - plain-energy graph route's `border_op`-based ghost-feature exchange) is not - implemented. -- **Charge-spin FiLM conditioning.** Combining `add_chg_spin_ebd` with - `spin.scheme: native` is rejected at model-construction time; use one or - the other. -- **ZBL zone bridging.** Combining `bridging_method: ZBL` with - `spin.scheme: native` is not supported on the pt_expt backend (`bridging_method` - is rejected there independently of spin -- see [Zone bridging - (ZBL)](#zone-bridging-zbl)). +graph route: none of the earlier restrictions remain. Multi-rank inference, +charge-spin FiLM conditioning (`add_chg_spin_ebd`), and ZBL zone bridging +(`bridging_method: ZBL`) all combine freely with `spin.scheme: native`. ## Embedding extraction @@ -723,12 +711,12 @@ closed over the one-hop neighbor shell. - DPA4/SeZM is implemented for the PyTorch backend only. - Export uses `.pt2` (AOTInductor); the TorchScript freeze path is not used. - Model compression is not supported. -- Multi-rank (multi-GPU/MPI) LAMMPS inference works for a plain energy model - on both export routes: the PT `edge_vec` archive and the pt_expt - NeighborGraph archive each embed a with-comm artifact. ZBL zone bridging is - single-rank (its Source Freeze Propagation gate folds each node's full - outgoing-edge set, which no single rank observes), and a multi-rank run of - such an archive fails fast rather than dropping the exchange. See +- Multi-rank (multi-GPU/MPI) LAMMPS inference works on both export routes: + the PT `edge_vec` archive and the pt_expt NeighborGraph archive each embed + a with-comm artifact. ZBL zone bridging (and its native-spin combination) + participates: the Source Freeze Propagation gate's per-node partials are + completed across ranks by one reverse-accumulate plus one + forward-broadcast border exchange. See [Multi-GPU (MPI) inference](#multi-gpu-mpi-inference). - The pt_expt graph-native inference route is unavailable only for `deepspin`-scheme spin, which stays on the dense route. Charge/spin @@ -737,9 +725,8 @@ closed over the one-hop neighbor shell. - `spin.scheme: native` is graph-only (it has no dense route) and supports multi-rank LAMMPS: ghost node features ride `border_op` per interaction block and ghost spins arrive through the LAMMPS `sp` forward-comm. - Charge-spin FiLM conditioning combines with it. Combining it with ZBL zone - bridging works single-rank; that combination is single-rank for the same - bridging reason as above. See [Native spin + Charge-spin FiLM conditioning and ZBL zone bridging both combine with it, + multi-rank included. See [Native spin (magnetic)](#native-spin-magnetic). ## Citation From 57c25fefcdde73fcc81c887067b239551f6c8ac6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 30 Jul 2026 08:59:56 +0000 Subject: [PATCH 15/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../tests/pt/model/test_sezm_parallel_bridging_parity.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/source/tests/pt/model/test_sezm_parallel_bridging_parity.py b/source/tests/pt/model/test_sezm_parallel_bridging_parity.py index eb318e97a7..46e5dbe61f 100644 --- a/source/tests/pt/model/test_sezm_parallel_bridging_parity.py +++ b/source/tests/pt/model/test_sezm_parallel_bridging_parity.py @@ -24,9 +24,7 @@ from deepmd.pt.model.model import ( get_model, ) -from deepmd.pt.utils import ( - env, # noqa: F401 - imports pt test env side effects -) +from deepmd.pt.utils import env # noqa: F401 - imports pt test env side effects from .test_sezm_parallel import ( _perturb_descriptor, @@ -193,9 +191,7 @@ def test_ablation_identity_exchange_diverges(self) -> None: mock, ) - from deepmd.pt.model.descriptor import ( - sezm as pt_sezm, - ) + from deepmd.pt.model.descriptor import sezm as pt_sezm with mock.patch.object( pt_sezm.DescrptSeZM, From d855cf5715993a243919b26ed950d7ea6df65cdf Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 31 Jul 2026 09:20:46 +0800 Subject: [PATCH 16/26] docs: fix contradictory native-spin heading and stale ABSENCE docstring CodeRabbit review on #5939. --- doc/model/dpa4.md | 8 ++++---- source/tests/infer/gen_dpa4_spin_zbl.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 7958fd2caa..bc75985225 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -591,10 +591,10 @@ inference](#multi-gpu-mpi-inference). are `None` placeholders there, exactly as for the plain-energy dpmodel route. -The following combinations are **not yet supported** on the native-spin -graph route: none of the earlier restrictions remain. Multi-rank inference, -charge-spin FiLM conditioning (`add_chg_spin_ebd`), and ZBL zone bridging -(`bridging_method: ZBL`) all combine freely with `spin.scheme: native`. +No native-spin combination restrictions remain on the graph route: +multi-rank inference, charge-spin FiLM conditioning (`add_chg_spin_ebd`), +and ZBL zone bridging (`bridging_method: ZBL`) all combine freely with +`spin.scheme: native`. ## Embedding extraction diff --git a/source/tests/infer/gen_dpa4_spin_zbl.py b/source/tests/infer/gen_dpa4_spin_zbl.py index eea3ac37f7..5a1f2fcdb1 100644 --- a/source/tests/infer/gen_dpa4_spin_zbl.py +++ b/source/tests/infer/gen_dpa4_spin_zbl.py @@ -313,7 +313,7 @@ def _assert_zbl_term_is_active(model_dict: dict) -> float: def _check_metadata(pt2_path: str) -> None: - """Assert the frozen archive's metadata and the with-comm ABSENCE. + """Assert the frozen archive's metadata and the with-comm PRESENCE. Parameters ---------- From deabcb796a52ee2d53bb6aba4833e94d906295f5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 31 Jul 2026 09:53:59 +0800 Subject: [PATCH 17/26] style(test): silence CodeQL alerts in the pt bridging parity test Two GitHub Advanced Security findings on the new file: - 'unittest' was imported both as a module and via 'from unittest import mock'; use a plain 'import unittest.mock' instead. - the 'deepmd.pt.utils.env' import was dead: this test pins CPU explicitly, and env is pulled in transitively by get_model plus the source/tests/pt package __init__. --- .../tests/pt/model/test_sezm_parallel_bridging_parity.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/source/tests/pt/model/test_sezm_parallel_bridging_parity.py b/source/tests/pt/model/test_sezm_parallel_bridging_parity.py index 46e5dbe61f..cabde13078 100644 --- a/source/tests/pt/model/test_sezm_parallel_bridging_parity.py +++ b/source/tests/pt/model/test_sezm_parallel_bridging_parity.py @@ -17,6 +17,7 @@ """ import unittest +import unittest.mock import numpy as np import torch @@ -24,7 +25,6 @@ from deepmd.pt.model.model import ( get_model, ) -from deepmd.pt.utils import env # noqa: F401 - imports pt test env side effects from .test_sezm_parallel import ( _perturb_descriptor, @@ -187,13 +187,9 @@ def test_ablation_identity_exchange_diverges(self) -> None: """Negative contract: stubbing the exchange to identity breaks the parity -- proves the geometry actually exercises the gate. """ - from unittest import ( - mock, - ) - from deepmd.pt.model.descriptor import sezm as pt_sezm - with mock.patch.object( + with unittest.mock.patch.object( pt_sezm.DescrptSeZM, "_gate_partial_exchange", lambda self, partials, comm_dict: partials, From b61ae9a8fa23a9296fc5637e5335d51d8e75fc84 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 31 Jul 2026 18:42:18 +0800 Subject: [PATCH 18/26] fix(dpmodel): hybrid must aggregate the dense-comm capability A DescrptHybrid inherited the concrete `dense_lower_supports_comm() == True` default while aggregating `has_message_passing_across_ranks()` over its children. For hybrid([dpa4, se_e2_a]) that combination makes `_needs_with_comm_artifact(model, "nlist")` true, so an auto/nlist freeze builds a dense with-comm trace that dies in DPA4's NotImplementedError as soon as comm_dict is passed. Aggregate it (and supports_edge_parallel) with ALL, matching the veto semantics already used for the atomic-model composition. --- deepmd/dpmodel/descriptor/hybrid.py | 18 +++++++++ .../dpmodel/test_atomic_model_capabilities.py | 38 +++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 110ae1fbe0..9709ce37f0 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -212,6 +212,24 @@ def has_message_passing_across_ranks(self) -> bool: descrpt.has_message_passing_across_ranks() for descrpt in self.descrpt_list ) + def supports_edge_parallel(self) -> bool: + """Returns whether the hybrid can run under domain decomposition. + + A veto by any child vetoes the whole concatenation: the hybrid + output contains that child's block, so the composite is only as + parallel-capable as its least capable member. + """ + return all(descrpt.supports_edge_parallel() for descrpt in self.descrpt_list) + + def dense_lower_supports_comm(self) -> bool: + """Returns whether every child's DENSE lower implements comm_dict. + + ALL rather than ANY: the dense with-comm trace passes ``comm_dict`` + to every child, so one child whose dense adapter raises on it (DPA4) + makes the whole hybrid's dense comm path non-viable. + """ + return all(descrpt.dense_lower_supports_comm() for descrpt in self.descrpt_list) + def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" return True diff --git a/source/tests/common/dpmodel/test_atomic_model_capabilities.py b/source/tests/common/dpmodel/test_atomic_model_capabilities.py index d675da3d51..70b504353b 100644 --- a/source/tests/common/dpmodel/test_atomic_model_capabilities.py +++ b/source/tests/common/dpmodel/test_atomic_model_capabilities.py @@ -118,6 +118,44 @@ def test_dp_atomic_model_delegates_to_descriptor() -> None: ) +def test_hybrid_descriptor_aggregates_dense_comm() -> None: + """A hybrid is only as comm-capable as its least capable child. + + Without the ALL-aggregation a ``hybrid([dpa4, se_e2_a])`` inherits the + ``True`` default while its DPA4 child says ``False``, so the freeze + machinery emits a dense with-comm artifact whose trace then dies in + DPA4's ``NotImplementedError``. + """ + from deepmd.dpmodel.descriptor.hybrid import ( + DescrptHybrid, + ) + + dpa4 = _dpa4_descriptor(bridging=False) + sea = DescrptSeA(rcut=4.0, rcut_smth=3.5, sel=[8, 8]) + assert dpa4.dense_lower_supports_comm() is False + assert sea.dense_lower_supports_comm() is True + + mixed = DescrptHybrid(list=[dpa4, sea]) + # the hybrid itself stays on the DENSE lower ... + assert mixed.uses_graph_lower() is False + # ... and needs a ghost exchange because of the DPA4 child ... + assert mixed.has_message_passing_across_ranks() is True + # ... which the dense lower cannot provide: ALL-aggregation vetoes it + assert mixed.dense_lower_supports_comm() is False + assert mixed.supports_edge_parallel() is True + + # the True branch of the same ALL rule + all_dense = DescrptHybrid( + list=[sea, DescrptSeA(rcut=4.0, rcut_smth=3.5, sel=[8, 8])] + ) + assert all_dense.dense_lower_supports_comm() is True + assert all_dense.has_message_passing_across_ranks() is False + + # and the atomic model delegates the aggregate, not the default + assert _dp_atomic_model(mixed).dense_lower_supports_comm() is False + assert _dp_atomic_model(all_dense).dense_lower_supports_comm() is True + + class _EdgeParallelChild(InnerPotentialAtomicModel): """Stub child with settable capabilities (test_zbl_bridging.py pattern).""" From 94401ade084fbe9a89e2367314583f4b75e4ada7 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 31 Jul 2026 18:42:18 +0800 Subject: [PATCH 19/26] test(zbl): update the freeze regressions to the with-comm contract Two executable assertions still pinned the pre-#5906 single-rank contract (`has_comm_artifact is False`) for the ZBL and native-spin+ZBL graph freezes. They are skipped only when CI=true, so a local full-suite run compiled the archive and then failed on the obsolete expectation. Flip both to the multi-rank contract and also assert the nested forward_lower_with_comm.pt2 entry is present. --- source/tests/pt_expt/model/test_zbl_bridging.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/source/tests/pt_expt/model/test_zbl_bridging.py b/source/tests/pt_expt/model/test_zbl_bridging.py index 5da31ad263..c1bc55f3a4 100644 --- a/source/tests/pt_expt/model/test_zbl_bridging.py +++ b/source/tests/pt_expt/model/test_zbl_bridging.py @@ -414,8 +414,11 @@ def test_graph_freeze_and_deep_eval_parity(self, tmp_path) -> None: with zipfile.ZipFile(model_file) as z: md = json.loads(z.read("model/extra/metadata.json").decode("utf-8")) - # single-rank contract: compositions never get a with-comm artifact - assert md["has_comm_artifact"] is False + # multi-rank contract (issue #5906): the SFPG partials are completed + # across ranks, so a bridged composition DOES get a with-comm twin + assert md["has_comm_artifact"] is True + with zipfile.ZipFile(model_file) as z: + assert "model/extra/forward_lower_with_comm.pt2" in z.namelist() dp = DeepPot(str(model_file)) e, f, v = dp.eval( @@ -643,14 +646,17 @@ def test_native_spin_with_bridging_graph_freeze_and_deep_eval(tmp_path) -> None: model_file = tmp_path / "dpa4_native_spin_zbl_graph.pt2" # native spin has no dense lower at all, so the graph kind is the only - # valid one here; the composition additionally forbids a with-comm twin. + # valid one here; since issue #5906 the graph lower additionally carries + # a with-comm twin (only the dense lower stays single-rank for spin). deserialize_to_file( str(model_file), {"model": model.serialize()}, lower_kind="graph" ) with zipfile.ZipFile(model_file) as z: md = json.loads(z.read("model/extra/metadata.json").decode("utf-8")) + names = z.namelist() assert md["is_spin"] is True - assert md["has_comm_artifact"] is False + assert md["has_comm_artifact"] is True + assert "model/extra/forward_lower_with_comm.pt2" in names assert md["use_spin"] == [True, False] dp = DeepPot(str(model_file)) From 5b8608083eaac517a42cb7fa0e1f5c0bb1d4647c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 31 Jul 2026 18:42:38 +0800 Subject: [PATCH 20/26] refactor(pt_expt): the graph with-comm exporter belongs to EnergyModel make_model() produces DOS, dipole, polar, property and energy models, but forward_lower_graph_exportable_with_comm hard-codes do_grad_r("energy"), do_grad_c("energy") and _translate_energy_keys. Defining it on the common factory made every non-energy model advertise an API that would raise on the missing energy output; the serializer's energy gate only hid it. Give it the same ownership as its non-comm twin: define it on EnergyModel and alias it in LinearEnergyModel. NativeSpinEnergyModel is built on EnergyModel, so the spin route keeps it through the MRO. --- deepmd/pt_expt/model/dp_linear_model.py | 6 + deepmd/pt_expt/model/ener_model.py | 185 ++++++++++++++++++++++++ deepmd/pt_expt/model/make_model.py | 185 ------------------------ 3 files changed, 191 insertions(+), 185 deletions(-) diff --git a/deepmd/pt_expt/model/dp_linear_model.py b/deepmd/pt_expt/model/dp_linear_model.py index e4d9e0b877..1ae8255f84 100644 --- a/deepmd/pt_expt/model/dp_linear_model.py +++ b/deepmd/pt_expt/model/dp_linear_model.py @@ -41,6 +41,12 @@ class LinearEnergyModel(DPModelCommon, DPLinearModel_): # identical for any energy model -- reuse EnergyModel's verbatim so # compositions (e.g. analytical bridging) freeze like standard models. forward_lower_graph_exportable = EnergyModel.forward_lower_graph_exportable + # Same ownership for the with-comm twin: it is the SAME energy contract + # plus the border-exchange inputs, so a bridged composition gets the + # multi-rank artifact through the same alias (issue #5906). + forward_lower_graph_exportable_with_comm = ( + EnergyModel.forward_lower_graph_exportable_with_comm + ) def __init__( self, diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 77b6a84bac..06f57176cd 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -550,3 +550,188 @@ def fn( aparam, charge_spin, ) + + def forward_lower_graph_exportable_with_comm( + self, + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + send_list: torch.Tensor, + send_proc: torch.Tensor, + recv_proc: torch.Tensor, + send_num: torch.Tensor, + recv_num: torch.Tensor, + communicator: torch.Tensor, + nlocal: torch.Tensor, + nghost: torch.Tensor, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + """Trace ``forward_common_lower_graph`` with comm_dict tensors as + additional positional inputs -- the with-comm counterpart of + :meth:`forward_lower_graph_exportable` for message-passing graph + descriptors (dpa2's repformer block drives cross-rank ghost refresh + via ``deepmd_export::border_op``, see + :meth:`~deepmd.pt_expt.descriptor.repformers. + DescrptBlockRepformers._exchange_ghosts_graph`). + + Mirrors the dense ``forward_common_lower_exportable_with_comm`` + (``pt_expt/model/make_model.py``): packs the 8 trailing positional + comm tensors into a ``comm_dict`` inside the traced function. Also + derives ``n_local`` (the per-frame OWNED node count, reshaped to + ``(1,)``; single-frame -- LAMMPS always drives inference with + ``nf=1``) from the scalar ``nlocal`` tensor, so the differentiated + reduction excludes ghost (not-owned) nodes (see + :meth:`forward_common_lower_graph`'s ``n_local`` parameter). Unlike + the plain-graph export path (which traces + ``forward_common_lower_graph_exportable`` and then wraps a SECOND + make_fx trace around the key-translation closure), this method + traces ONCE: the comm-dict packing, ``n_local`` derivation, the + ``forward_common_lower_graph`` call and the key translation all live + in a single traced ``fn`` -- following the dense with-comm + precedent, which is also a single trace. + + Parameters + ---------- + atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin, do_atomic_virial + As in :meth:`forward_lower_graph_exportable`. + send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost + The 8 comm tensors (see ``_make_comm_sample_inputs`` in + ``serialization.py``), packed into ``comm_dict`` inside the + traced function. + + Runtime device contract: ALL 8 stay on CPU, symmetric with + the dense with-comm artifact -- they are consumed only by the + opaque ``border_op`` whose HOST code dereferences their + ``data_ptr`` (``send_list`` carries raw host pointers) and + reads ``nlocal``/``nghost`` via cheap host ``.item()`` calls. + Deriving the in-graph owned count from a device-placed + ``nlocal`` instead (the previous design) made every per-layer + ``border_op`` forward AND custom backward pull the scalars + back with synchronizing D2H reads (``4 * nlayers`` per MD + step). The C++ ``run_model_graph_with_comm`` implements this + placement. + n_local + (1,) int64 ON THE MODEL DEVICE: the per-frame OWNED node + count consumed IN-GRAPH by the owned-node energy mask (it + becomes a device kernel operand after + ``move_to_device_pass``, like ``n_node``; a CPU tensor fed + there is read as a device pointer -- CUDA illegal memory + access). Carries the same value as the ``nlocal`` comm + tensor; the two inputs exist precisely to separate the + device-compute role from the host-MPI-control role. + **make_fx_kwargs + Extra keyword arguments forwarded to ``make_fx`` + (e.g. ``tracing_mode="symbolic"``). + + Returns + ------- + torch.nn.Module + A traced module whose ``forward`` accepts ``(atype, n_node, + n_local, edge_index, edge_vec, edge_mask, destination_order, + destination_row_ptr, source_order, source_row_ptr, fparam, + aparam, charge_spin, send_list, send_proc, recv_proc, send_num, + recv_num, communicator, nlocal, nghost)`` and returns a dict with the + SAME public keys as :meth:`forward_lower_graph_exportable` + (``atom_energy``, ``energy``, ``force``, ``virial``, + ``atom_virial`` when ``do_atomic_virial``). + """ + model = self + do_grad_r = self.do_grad_r("energy") + do_grad_c = self.do_grad_c("energy") + + def fn( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + send_list: torch.Tensor, + send_proc: torch.Tensor, + recv_proc: torch.Tensor, + send_num: torch.Tensor, + recv_num: torch.Tensor, + communicator: torch.Tensor, + nlocal: torch.Tensor, + nghost: torch.Tensor, + ) -> dict[str, torch.Tensor]: + comm_dict = { + "send_list": send_list, + "send_proc": send_proc, + "recv_proc": recv_proc, + "send_num": send_num, + "recv_num": recv_num, + "communicator": communicator, + "nlocal": nlocal, + "nghost": nghost, + } + # ``n_local`` (slot 2, DEVICE) is the owned-count input consumed + # by the in-graph owned-node mask; the CPU ``nlocal`` comm + # tensor is host control metadata for border_op only. + model_ret = model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=True, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + comm_dict=comm_dict, + ) + return _translate_energy_keys( + model_ret, + do_grad_r=do_grad_r, + do_grad_c=do_grad_c, + do_atomic_virial=do_atomic_virial, + local=True, + ) + + return make_fx(fn, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + send_list, + send_proc, + recv_proc, + send_num, + recv_num, + communicator, + nlocal, + nghost, + ) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 32ac4e7dc5..55c05c5ec3 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -1241,191 +1241,6 @@ def fn_spin( spin, ) - def forward_lower_graph_exportable_with_comm( - self, - atype: torch.Tensor, - n_node: torch.Tensor, - n_local: torch.Tensor, - edge_index: torch.Tensor, - edge_vec: torch.Tensor, - edge_mask: torch.Tensor, - destination_order: torch.Tensor, - destination_row_ptr: torch.Tensor, - source_order: torch.Tensor, - source_row_ptr: torch.Tensor, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None, - send_list: torch.Tensor, - send_proc: torch.Tensor, - recv_proc: torch.Tensor, - send_num: torch.Tensor, - recv_num: torch.Tensor, - communicator: torch.Tensor, - nlocal: torch.Tensor, - nghost: torch.Tensor, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - """Trace ``forward_common_lower_graph`` with comm_dict tensors as - additional positional inputs -- the with-comm counterpart of - :meth:`forward_lower_graph_exportable` for message-passing graph - descriptors (dpa2's repformer block drives cross-rank ghost refresh - via ``deepmd_export::border_op``, see - :meth:`~deepmd.pt_expt.descriptor.repformers. - DescrptBlockRepformers._exchange_ghosts_graph`). - - Mirrors the dense ``forward_common_lower_exportable_with_comm`` - (``pt_expt/model/make_model.py``): packs the 8 trailing positional - comm tensors into a ``comm_dict`` inside the traced function. Also - derives ``n_local`` (the per-frame OWNED node count, reshaped to - ``(1,)``; single-frame -- LAMMPS always drives inference with - ``nf=1``) from the scalar ``nlocal`` tensor, so the differentiated - reduction excludes ghost (not-owned) nodes (see - :meth:`forward_common_lower_graph`'s ``n_local`` parameter). Unlike - the plain-graph export path (which traces - ``forward_common_lower_graph_exportable`` and then wraps a SECOND - make_fx trace around the key-translation closure), this method - traces ONCE: the comm-dict packing, ``n_local`` derivation, the - ``forward_common_lower_graph`` call and the key translation all live - in a single traced ``fn`` -- following the dense with-comm - precedent, which is also a single trace. - - Parameters - ---------- - atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin, do_atomic_virial - As in :meth:`forward_lower_graph_exportable`. - send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost - The 8 comm tensors (see ``_make_comm_sample_inputs`` in - ``serialization.py``), packed into ``comm_dict`` inside the - traced function. - - Runtime device contract: ALL 8 stay on CPU, symmetric with - the dense with-comm artifact -- they are consumed only by the - opaque ``border_op`` whose HOST code dereferences their - ``data_ptr`` (``send_list`` carries raw host pointers) and - reads ``nlocal``/``nghost`` via cheap host ``.item()`` calls. - Deriving the in-graph owned count from a device-placed - ``nlocal`` instead (the previous design) made every per-layer - ``border_op`` forward AND custom backward pull the scalars - back with synchronizing D2H reads (``4 * nlayers`` per MD - step). The C++ ``run_model_graph_with_comm`` implements this - placement. - n_local - (1,) int64 ON THE MODEL DEVICE: the per-frame OWNED node - count consumed IN-GRAPH by the owned-node energy mask (it - becomes a device kernel operand after - ``move_to_device_pass``, like ``n_node``; a CPU tensor fed - there is read as a device pointer -- CUDA illegal memory - access). Carries the same value as the ``nlocal`` comm - tensor; the two inputs exist precisely to separate the - device-compute role from the host-MPI-control role. - **make_fx_kwargs - Extra keyword arguments forwarded to ``make_fx`` - (e.g. ``tracing_mode="symbolic"``). - - Returns - ------- - torch.nn.Module - A traced module whose ``forward`` accepts ``(atype, n_node, - n_local, edge_index, edge_vec, edge_mask, destination_order, - destination_row_ptr, source_order, source_row_ptr, fparam, - aparam, charge_spin, send_list, send_proc, recv_proc, send_num, - recv_num, communicator, nlocal, nghost)`` and returns a dict with the - SAME public keys as :meth:`forward_lower_graph_exportable` - (``atom_energy``, ``energy``, ``force``, ``virial``, - ``atom_virial`` when ``do_atomic_virial``). - """ - model = self - do_grad_r = self.do_grad_r("energy") - do_grad_c = self.do_grad_c("energy") - - def fn( - atype: torch.Tensor, - n_node: torch.Tensor, - n_local: torch.Tensor, - edge_index: torch.Tensor, - edge_vec: torch.Tensor, - edge_mask: torch.Tensor, - destination_order: torch.Tensor, - destination_row_ptr: torch.Tensor, - source_order: torch.Tensor, - source_row_ptr: torch.Tensor, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None, - send_list: torch.Tensor, - send_proc: torch.Tensor, - recv_proc: torch.Tensor, - send_num: torch.Tensor, - recv_num: torch.Tensor, - communicator: torch.Tensor, - nlocal: torch.Tensor, - nghost: torch.Tensor, - ) -> dict[str, torch.Tensor]: - comm_dict = { - "send_list": send_list, - "send_proc": send_proc, - "recv_proc": recv_proc, - "send_num": send_num, - "recv_num": recv_num, - "communicator": communicator, - "nlocal": nlocal, - "nghost": nghost, - } - # ``n_local`` (slot 2, DEVICE) is the owned-count input consumed - # by the in-graph owned-node mask; the CPU ``nlocal`` comm - # tensor is host control metadata for border_op only. - model_ret = model.forward_common_lower_graph( - atype, - n_node, - n_local, - edge_index, - edge_vec, - edge_mask, - destination_order, - destination_row_ptr, - source_order, - source_row_ptr, - destination_sorted=True, - do_atomic_virial=do_atomic_virial, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - comm_dict=comm_dict, - ) - return _translate_energy_keys( - model_ret, - do_grad_r=do_grad_r, - do_grad_c=do_grad_c, - do_atomic_virial=do_atomic_virial, - local=True, - ) - - return make_fx(fn, **make_fx_kwargs)( - atype, - n_node, - n_local, - edge_index, - edge_vec, - edge_mask, - destination_order, - destination_row_ptr, - source_order, - source_row_ptr, - fparam, - aparam, - charge_spin, - send_list, - send_proc, - recv_proc, - send_num, - recv_num, - communicator, - nlocal, - nghost, - ) - def forward_common_lower_exportable_with_comm( self, extended_coord: torch.Tensor, From 4021124336c8fce310a72cb54b688ec102c4b14e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 31 Jul 2026 18:42:38 +0800 Subject: [PATCH 21/26] docs: sweep the stale single-rank and dense-freeze statements Six sites still described the pre-#5906 contract or a CLI behaviour that does not exist: - dpa4.uses_graph_lower / sezm.has_message_passing_across_ranks claimed bridging still fails multi-rank fast, contradicting the capability split this branch introduces. - NativeSpinModelKind cited "native-spin lowers are single-rank only" as its motivating example; only the DENSE lower is excluded now. - get_sezm_model advertised bridging and spin as unsupported although the factory deliberately constructs both; replace with the precise matrix (deepspin scheme, lora, use_compile, preset_out_bias remain unsupported). - doc/model/dpa4.md offered plain DPA4 a choice of dense or graph freeze via a '--lower-kind auto' default. There is no 'auto' choice (nlist and graph only, default nlist) and freeze() overrides any non-graph request to graph for graph-capable models. - gen_dpa4_spin_zbl.py still headed its contract 'Single-rank only'. --- deepmd/dpmodel/descriptor/dpa4.py | 6 ++++-- deepmd/dpmodel/model/native_spin_model.py | 6 ++++-- deepmd/pt/model/descriptor/sezm.py | 5 +++-- deepmd/pt_expt/model/get_model.py | 10 +++++++--- doc/model/dpa4.md | 12 +++++++----- source/tests/infer/gen_dpa4_spin_zbl.py | 5 +++-- 6 files changed, 28 insertions(+), 16 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 2f48a1487f..801b4c34f2 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -2367,8 +2367,10 @@ def uses_graph_lower(self) -> bool: spin and charge_spin are threaded through ``call_graph`` like any other per-node/per-frame input, and bridging is applied inside the shared ``_run_graph`` forward with no extra threading (it - reads ``self.bridging_switch`` directly). Bridging models still - fail multi-rank fast via ``has_message_passing_across_ranks``. + reads ``self.bridging_switch`` directly). Bridging models are + multi-rank capable too: their SFPG per-node partials are + completed across ranks by ``_gate_partial_exchange`` before the + gate is applied (issue #5906). """ return not self._graph_lower_disabled diff --git a/deepmd/dpmodel/model/native_spin_model.py b/deepmd/dpmodel/model/native_spin_model.py index 15943ba206..e44a9381d6 100644 --- a/deepmd/dpmodel/model/native_spin_model.py +++ b/deepmd/dpmodel/model/native_spin_model.py @@ -34,8 +34,10 @@ class NativeSpinModelKind: ``NativeSpinEnergyModel``) are parallel products with NO subclass relation between them -- an ``isinstance`` against one backend's concrete class is silently dead in the other. Backend seams that need a - cross-backend family test (e.g. the with-comm freeze gate: native-spin - lowers are single-rank only) test against this shared marker instead. + cross-backend family test test against this shared marker instead. The + motivating consumer is the with-comm freeze gate, where native spin + excludes only the DENSE lower; native-spin GRAPH lowers do participate + in the with-comm path and carry the nested artifact (issue #5906). """ diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 394c3a2eea..084f75b14d 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -2294,8 +2294,9 @@ def has_message_passing_across_ranks(self) -> bool: A domain-decomposed run must exchange them through ``border_op``, so multi-rank inference always needs the with-comm exchange. Whether - multi-rank is POSSIBLE at all is :meth:`supports_edge_parallel` - (bridging vetoes it there). + multi-rank is POSSIBLE at all is :meth:`supports_edge_parallel`, + which is ``True`` for every configuration including bridging: the + SFPG partials are completed by ``_gate_partial_exchange``. """ return True diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index b28923601d..6dedb96d9e 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -78,9 +78,13 @@ def get_sezm_model(data: dict) -> EnergyModel: training configs are interchangeable between the pt and pt_expt backends. In addition to the ``SeZM``/``sezm``/``dpa4`` aliases accepted by pt, pt_expt also accepts ``DPA4``. - The pt-only SeZM extensions (bridging, LoRA, compile, spin, - preset_out_bias) are not supported here and raise - ``NotImplementedError``. + Supported SeZM extensions: analytical bridging (e.g. ZBL), composed by + :func:`_compose_bridging`, and native-scheme spin, routed to + :func:`get_native_spin_model`; the two combine. + + Still unsupported here, each raising ``NotImplementedError``: the + virtual-atom (``deepspin``) spin scheme, ``lora``, ``use_compile``, and + ``preset_out_bias``. Notes ----- diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index bc75985225..338ccc67b2 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -570,11 +570,13 @@ inference](#multi-gpu-mpi-inference). is graph-eligible. `dp --pt_expt freeze --lower-kind graph` on a `deepspin`-scheme model raises an error at freeze time, per the dense/graph eligibility rule above. -- **Graph route only, no dense fallback.** Unlike a plain-energy DPA4/SeZM - descriptor -- which can freeze to either the dense or the graph lower -- - a native-spin descriptor has only the graph lower. `--lower-kind auto` - (the default) resolves to `graph`; `--lower-kind nlist` is not a valid - option for a native-spin model. +- **Graph route only, no dense fallback.** A native-spin descriptor has only + the graph lower. This is not a spin-specific restriction on the CLI: every + graph-capable DPA4/SeZM model, plain-energy included, is frozen to the + graph lower. `--lower-kind` accepts only `nlist` (the default) and `graph`, + and `freeze()` overrides any non-`graph` request to `graph` whenever the + model is graph-lower capable, logging a warning. A dense artifact is + therefore not selectable through this entry point for these models. - **Multi-rank capable.** The frozen archive embeds the with-comm artifact (`has_comm_artifact` is `true`), so multi-rank LAMMPS works exactly as described in [Multi-GPU (MPI) inference](#multi-gpu-mpi-inference): the diff --git a/source/tests/infer/gen_dpa4_spin_zbl.py b/source/tests/infer/gen_dpa4_spin_zbl.py index 5a1f2fcdb1..eb191f0104 100644 --- a/source/tests/infer/gen_dpa4_spin_zbl.py +++ b/source/tests/infer/gen_dpa4_spin_zbl.py @@ -21,7 +21,7 @@ freeze must carry ``is_spin`` metadata for a model whose top-level class is the linear composition -- neither single-feature fixture exercises that. -Single-rank only, and this fixture PINS that limitation +Multi-rank capable, and this fixture PINS that contract ------------------------------------------------------- Bridging enables the descriptor's Source Freeze Propagation Gate, whose per-node ``eta_j = prod_{e: src_e = j} w_e`` folds a node's FULL outgoing-edge @@ -92,7 +92,8 @@ # Bridging radii, identical to gen_dpa4_zbl.py's: they feed the descriptor's # InnerClamp AND BridgingSwitch (built together from the same radii) on the -# LEARNED child, and are what makes the composition single-rank only. +# LEARNED child, and are what makes the composition need the SFPG cross-rank +# completion (hence the with-comm artifact) under domain decomposition. _BRIDGING_R_INNER = 0.8 _BRIDGING_R_OUTER = 1.2 From 975e17e8543412d8152791205890196cb91e29ba Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 31 Jul 2026 19:45:10 +0800 Subject: [PATCH 22/26] docs: fix duplicated word in the native-spin marker docstring Left over from removing the stale parenthetical in 4021124. --- deepmd/dpmodel/model/native_spin_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/dpmodel/model/native_spin_model.py b/deepmd/dpmodel/model/native_spin_model.py index e44a9381d6..be6701447b 100644 --- a/deepmd/dpmodel/model/native_spin_model.py +++ b/deepmd/dpmodel/model/native_spin_model.py @@ -34,7 +34,7 @@ class NativeSpinModelKind: ``NativeSpinEnergyModel``) are parallel products with NO subclass relation between them -- an ``isinstance`` against one backend's concrete class is silently dead in the other. Backend seams that need a - cross-backend family test test against this shared marker instead. The + cross-backend family test check against this shared marker instead. The motivating consumer is the with-comm freeze gate, where native spin excludes only the DENSE lower; native-spin GRAPH lowers do participate in the with-comm path and carry the nested artifact (issue #5906). From eb238607e99709aaae94d3a057e5d7b113382a38 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 1 Aug 2026 09:44:23 +0800 Subject: [PATCH 23/26] fix(pt_expt): bridging is a composition, so `standard` must reject it get_standard_model honored `bridging_method` and returned a LinearEnergyModel -- the type asked for was not the type returned. It also made a second owner of the bridging build, and the two owners had already drifted: get_sezm_model promotes descriptor.exclude_types to model-level pair_exclude_types and this one never did, changing a 0.9 A Ni-O dimer by 79.97 eV (max |dF| 318.48 eV/A). Reject `bridging_method` in the standard builder instead. Rejecting rather than ignoring keeps the original fail-loud property: silently dropping the term yields a physically different model than the config requests. DPA4 (`type: "dpa4"`/`"sezm"`) is now the single bridging owner in this backend. The route was unreachable from a validated input.json anyway -- argcheck declares `bridging_method` only under the dpa4/sezm variant -- so no supported configuration changes behavior. Follow-ups: #5947 (drop the exclusion promotion), #5948 (express bridging as an explicit linear_ener composition, after which the restriction is moot). --- deepmd/pt_expt/model/get_model.py | 58 +++++++---- .../pt_expt/model/test_get_model_bridging.py | 98 ++++++++++++------- 2 files changed, 101 insertions(+), 55 deletions(-) diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 6dedb96d9e..30e27c91e9 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -180,10 +180,11 @@ def _compose_bridging(model: Any, data: dict, bridging_method: str) -> Any: Composition, not a flag (first-principles design): the analytical bridging term is its own atomic model, summed with the learned one by the existing linear composition machinery. The ONE owner of the - composition build for this backend -- both :func:`get_sezm_model` - (``type: "dpa4"``) and :func:`get_standard_model` (``type: - "standard"``) route through here, mirroring the dpmodel twin - (``deepmd/dpmodel/model/model.py``). + composition build for this backend: :func:`get_sezm_model` + (``type: "dpa4"``/``"sezm"``) is its only caller, because bridging + yields a composition and so is not expressible on a non-composite + model type -- :func:`get_standard_model` rejects it. Issue #5948 + tracks spelling the composition explicitly as ``linear_ener``. Parameters ---------- @@ -230,31 +231,46 @@ def _compose_bridging(model: Any, data: dict, bridging_method: str) -> Any: def get_standard_model(data: dict) -> Any: - """Build a pt_expt standard model, honoring ``bridging_method``. + """Build a pt_expt standard model: one descriptor plus one fitting net. - pt_expt twin of :func:`deepmd.dpmodel.model.model.get_standard_model`: - the analytical-bridging radii feed the DESCRIPTOR's - InnerClamp/BridgingSwitch and the method composes the atomic model with - its InnerPotential term. Without this wrapper a ``type: "standard"`` - config with ``bridging_method`` silently dropped the bridging term - (backend divergence from dpmodel -- issue #5906 Task 4 audit). + ``bridging_method`` is rejected here rather than honored. Analytical + bridging is a COMPOSITION -- it yields a ``LinearEnergyModel`` over + ``[learned, InnerPotential]`` -- so a builder that accepted it would + return a model of a different kind than the one requested. pt_expt + keeps exactly one bridging owner, :func:`get_sezm_model` + (``type: "dpa4"``/``"sezm"``), so the composition and its + ``exclude_types`` reconciliation cannot drift between two builders. + + Rejecting is deliberate over silently ignoring: dropping a bridging + term without a word yields a physically different model than the config + asks for. Issue #5948 tracks replacing the flag with an explicit + ``linear_ener`` composition, at which point this restriction is moot. Parameters ---------- data : dict The data to construct the model. + + Returns + ------- + Any + The constructed standard model. + + Raises + ------ + ValueError + If ``bridging_method`` is set: bridging is not expressible on a + non-composite model type. """ - data = copy.deepcopy(data) bridging_method = str(data.get("bridging_method", "none")) - bridging_enabled = bridging_method.lower() not in ("none", "") - if bridging_enabled: - data.setdefault("descriptor", {}) - data["descriptor"]["inner_clamp_r_inner"] = data.get("bridging_r_inner", 0.5) - data["descriptor"]["inner_clamp_r_outer"] = data.get("bridging_r_outer", 0.8) - model = _model_factory.get_standard_model(data) - if not bridging_enabled: - return model - return _compose_bridging(model, data, bridging_method) + if bridging_method.lower() not in ("none", ""): + raise ValueError( + "`bridging_method` is not supported for a standard model in the " + "pt_expt backend: analytical bridging builds a linear " + 'composition, not a standard model. Use model `type: "dpa4"` ' + '(or `"sezm"`) with the same descriptor and fitting net.' + ) + return _model_factory.get_standard_model(data) def get_native_spin_model(data: dict) -> NativeSpinEnergyModel: diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index f128bd25ea..d9e02492dd 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -1,12 +1,25 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""pt_expt ``get_standard_model`` must honor ``bridging_method`` like its -dpmodel twin (``deepmd/dpmodel/model/model.py``) -- issue #5906 Task 4 -variant-alignment audit, gap 2: a ``type: "standard"`` config with bridging -silently dropped the InnerPotential composition in pt_expt. +"""Analytical bridging has exactly ONE owner per backend. + +Bridging builds a COMPOSITION (``LinearEnergyModel`` over +``[learned, InnerPotential]``), so it is not expressible on a non-composite +model type: ``type: "standard"`` would have to return a model of a +different kind than the one requested. pt_expt therefore owns bridging on +the DPA4/SeZM route only and REJECTS it in the standard builder -- loudly, +because silently dropping the term yields a physically different model. + +Two builders accepting the flag is exactly how the routes drifted: +``get_sezm_model`` promotes ``descriptor.exclude_types`` to model-level +``pair_exclude_types`` and the standard route never did, which changes a +0.9 A Ni-O dimer by ~80 eV (issue #5947). Issue #5948 replaces the flag +with an explicit ``linear_ener`` composition, after which this restriction +becomes moot. """ import copy +import pytest + from deepmd.dpmodel.atomic_model.linear_atomic_model import ( LinearEnergyAtomicModel, ) @@ -42,37 +55,56 @@ def _dpa4_standard_config() -> dict: } -def test_standard_model_type_builds_bridging_composition() -> None: - """pt_expt twin of dpmodel model.py's get_standard_model: a config with - bridging_method must compose [learned, InnerPotential], not silently - drop the bridging term. - """ - data = _dpa4_standard_config() +def _bridged(data: dict) -> dict: data["bridging_method"] = "ZBL" data["bridging_r_inner"] = 0.8 data["bridging_r_outer"] = 1.2 - model = get_standard_model(copy.deepcopy(data)) - assert isinstance(model.atomic_model, LinearEnergyAtomicModel) - assert len(model.atomic_model.models) == 2 - # The descriptor radii injection must ride the same seam (a composition - # without the inner-clamp radii would be a half-applied bridging config): - desc = model.atomic_model.models[0].descriptor - assert desc.bridging_switch is not None - # And the get_model router (type omitted -> "standard") reaches the same - # composition: - routed = get_model(copy.deepcopy(data)) - assert isinstance(routed.atomic_model, LinearEnergyAtomicModel) - - -def test_standard_model_type_maps_dpa4_fitting() -> None: - """Dpmodel maps dpa4_ener/sezm_ener fitting under type:'standard' via the - model registry; pt_expt must not raise where dpmodel builds. + return data + + +def test_standard_builder_rejects_bridging() -> None: + """The standard builder must not hand back a composition.""" + with pytest.raises(ValueError, match="bridging_method"): + get_standard_model(_bridged(_dpa4_standard_config())) + + +def test_get_model_rejects_bridging_without_dpa4_model_type() -> None: + """Same contract through the dispatcher: an omitted model type defaults + to the standard route, so it must reject rather than compose. + """ + with pytest.raises(ValueError, match="bridging_method"): + get_model(_bridged(_dpa4_standard_config())) + + +def test_standard_builder_without_bridging_is_unaffected() -> None: + """The rejection keys on the flag, not on the DPA4 components: a plain + DPA4 standard model still builds and carries no bridging switch. """ model = get_standard_model(_dpa4_standard_config()) - assert model is not None + assert not isinstance(model.atomic_model, LinearEnergyAtomicModel) assert model.atomic_model.descriptor.bridging_switch is None +@pytest.mark.parametrize( + "model_type", + [ + "dpa4", # canonical spelling + "sezm", # pt-compatible alias + ], +) +def test_dpa4_model_type_owns_the_composition(model_type: str) -> None: + """The one supported spelling composes [learned, InnerPotential] and + injects the radii into the learned child's descriptor. + """ + data = _bridged(_dpa4_standard_config()) + data["type"] = model_type + model = get_model(copy.deepcopy(data)) + assert isinstance(model.atomic_model, LinearEnergyAtomicModel) + assert len(model.atomic_model.models) == 2 + # a composition without the inner-clamp radii would be half-applied + assert model.atomic_model.models[0].descriptor.bridging_switch is not None + + def test_pt_checkpoint_eval_works_for_composition(tmp_path) -> None: """``DeepEval`` on a ``.pt`` checkpoint of a bridging composition. @@ -91,11 +123,8 @@ def test_pt_checkpoint_eval_works_for_composition(tmp_path) -> None: ModelWrapper, ) - config = _dpa4_standard_config() + config = _bridged(_dpa4_standard_config()) config["type"] = "dpa4" - config["bridging_method"] = "ZBL" - config["bridging_r_inner"] = 0.8 - config["bridging_r_outer"] = 1.2 model = get_model(copy.deepcopy(config)).to(torch.float64).eval() ckpt = str(tmp_path / "dpa4_zbl.pt") wrapper = ModelWrapper(model, model_params=copy.deepcopy(config)) @@ -122,8 +151,9 @@ def test_compile_attention_probe_tolerates_composition() -> None: _warn_compiled_attention, ) - data = _dpa4_standard_config() - data["bridging_method"] = "ZBL" - model = get_standard_model(data) + data = _bridged(_dpa4_standard_config()) + data["type"] = "dpa4" + model = get_model(data) + assert isinstance(model.atomic_model, LinearEnergyAtomicModel) # must not raise _warn_compiled_attention(model, "Default") From 2e7865c719953097003d796141c834d53955245e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 4 Aug 2026 17:43:47 +0800 Subject: [PATCH 24/26] fix: truthful factory return types, and two stale multi-rank claims - get_sezm_model was annotated `-> EnergyModel`, but its bridging path returns LinearEnergyModel, which is NOT an EnergyModel. Annotate it and get_standard_model against BaseModel (the honest common contract, as pt already does) and give _compose_bridging its concrete LinearEnergyModel result, instead of erasing all three with `Any`. - pt/entrypoints/freeze_pt2.py still explained the with-comm predicate by saying bridging reports supports_edge_parallel()=False and falls back to single-rank -- inverted since this branch. - doc/model/dpa4.md called the SFPG exchange "negligible". Both kernels end in an MPI_Barrier and the force graph differentiates through them, so the transposes add a matching pair of round trips; at small system sizes or high rank counts that can dominate. It has not been benchmarked, so the claim is removed rather than restated. --- deepmd/pt/entrypoints/freeze_pt2.py | 8 +++++--- deepmd/pt_expt/model/get_model.py | 24 ++++++++++++++---------- doc/model/dpa4.md | 9 +++++++-- 3 files changed, 26 insertions(+), 15 deletions(-) diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index 53f945c446..8c88e64105 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -951,9 +951,11 @@ def freeze_sezm_to_pt2( # Second artifact: the LAMMPS multi-rank with-comm graph. It threads the # eight border_op communication tensors so cross-rank ghost features are # exchanged between interaction blocks. Gated on the edge_vec lower contract - # (energy and native spin), so virtual spin (nlist interface) is excluded; - # bridging models report supports_edge_parallel()=False (Source Freeze - # Propagation is not rank-decomposable). Both fall back to single-rank. + # (energy and native spin), so virtual spin (nlist interface) is excluded + # and falls back to single-rank. Bridging models DO take this path: they + # report supports_edge_parallel()=True since the Source Freeze Propagation + # gate's per-node partials are completed across ranks by + # DescrptSeZM._gate_partial_exchange (issue #5906). with_comm = ( model.export_lower_input_kind() == "edge_vec" and model.supports_edge_parallel() ) diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 30e27c91e9..50f60ecf49 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -9,7 +9,7 @@ import copy import logging from typing import ( - Any, + TYPE_CHECKING, ) from deepmd.dpmodel.atomic_model.dp_atomic_model import ( @@ -36,9 +36,6 @@ from deepmd.pt_expt.model.dpa4_model import ( DPA4EnergyModel, ) -from deepmd.pt_expt.model.ener_model import ( - EnergyModel, -) from deepmd.pt_expt.model.model import ( BaseModel, ) @@ -53,6 +50,11 @@ normalize_spin_use_spin, ) +if TYPE_CHECKING: + from deepmd.pt_expt.model.dp_linear_model import ( + LinearEnergyModel, + ) + log = logging.getLogger(__name__) # Warn at most once per process for backend-ignored switches (keyed by name). @@ -71,7 +73,7 @@ get_zbl_model = _model_factory.get_zbl_model -def get_sezm_model(data: dict) -> EnergyModel: +def get_sezm_model(data: dict) -> BaseModel: """Build a pt_expt energy model from a DPA4/SeZM model config. Mirrors :func:`deepmd.pt.model.model.get_sezm_model` so that dpa4/sezm @@ -174,7 +176,9 @@ def get_sezm_model(data: dict) -> EnergyModel: return model -def _compose_bridging(model: Any, data: dict, bridging_method: str) -> Any: +def _compose_bridging( + model: BaseModel, data: dict, bridging_method: str +) -> "LinearEnergyModel": """Compose the learned model with its analytical bridging term. Composition, not a flag (first-principles design): the analytical @@ -198,8 +202,8 @@ def _compose_bridging(model: Any, data: dict, bridging_method: str) -> Any: Returns ------- - Any - A :class:`LinearEnergyModel` over ``[learned, InnerPotential]``. + LinearEnergyModel + A composition over ``[learned, InnerPotential]``. """ from deepmd.dpmodel.atomic_model.inner_potential import ( InnerPotentialAtomicModel, @@ -230,7 +234,7 @@ def _compose_bridging(model: Any, data: dict, bridging_method: str) -> Any: return LinearEnergyModel(atomic_model_=composed) -def get_standard_model(data: dict) -> Any: +def get_standard_model(data: dict) -> BaseModel: """Build a pt_expt standard model: one descriptor plus one fitting net. ``bridging_method`` is rejected here rather than honored. Analytical @@ -253,7 +257,7 @@ def get_standard_model(data: dict) -> Any: Returns ------- - Any + BaseModel The constructed standard model. Raises diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 338ccc67b2..8cce7b9faf 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -462,8 +462,13 @@ spin). The Source Freeze Propagation gate folds each node's full *outgoing*-edge set, which no single rank observes for ghost owners; the with-comm artifact completes the gate's per-node `[log eta, zero count]` partials across ranks with one reverse-accumulate plus one forward-broadcast -border exchange of an `(N, 2)` tensor per forward pass — negligible next to -the per-block feature exchanges that already run. Native spin +border exchange of an `(N, 2)` tensor per forward pass. The payload is +narrow, but the cost is not bandwidth alone: both kernels end in an +`MPI_Barrier`, and the force graph differentiates through them, so their +transposes add a matching pair of round trips to the force evaluation. On +small systems or at high rank counts those synchronizations can dominate; +the exchange has not been benchmarked across system sizes and rank counts. +Native spin (`scheme: "native"`) and charge/spin conditioning likewise support multi-rank: native spin reuses the `edge_vec` interface on PT and the NeighborGraph lower on pt_expt. The remainder of this subsection describes From 0408b38863ff3f0cede417e1e5d363d3552eddbb Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 4 Aug 2026 17:50:11 +0800 Subject: [PATCH 25/26] test(lmp): one owner for the DPA4 native-spin LAMMPS harness test_lammps_dpa4_spin_zbl_pt2.py shared 421 byte-identical lines with test_lammps_dpa4_spin_graph_pt2.py: the 4-atom NiO system, the live DeepPot reference, setup/teardown, LAMMPS construction, force_mag gathering, the MPI subprocess/parser and the parity assertions. Any change to tolerances, the runner output schema or timeout handling had to be made twice. Extract them into dpa4_spin_harness.py. Each module now carries only its archive, its geometry and its variant-specific claims: spin_graph 605 -> 230 lines spin_zbl 566 -> 214 lines Kept narrowly scoped rather than folded into lammps_test_utils: these pieces are DPA4-native-spin-specific (`atom_modify map yes` for the graph route's atom map, the spin_norm/hbar `fm` convention, and a timeout-bounded runner that SIGKILLs the whole mpirun process group), and moving them into the shared module would change behaviour for unrelated tests. Same 8 tests collect with the same IDs; the graph module's two MPI tests now use the shared assertion helper the ZBL module had already factored out. Runtime validation is the C++/LAMMPS CI jobs: the .pt2 fixtures are not present locally and this box has no reliable MPI. --- source/lmp/tests/dpa4_spin_harness.py | 542 ++++++++++++++++++ .../test_lammps_dpa4_chg_spin_deepspin_pt2.py | 2 +- .../tests/test_lammps_dpa4_spin_graph_pt2.py | 453 ++------------- .../tests/test_lammps_dpa4_spin_zbl_pt2.py | 434 ++------------ source/lmp/tests/test_lammps_dpa4_zbl_pt2.py | 2 +- 5 files changed, 624 insertions(+), 809 deletions(-) create mode 100644 source/lmp/tests/dpa4_spin_harness.py diff --git a/source/lmp/tests/dpa4_spin_harness.py b/source/lmp/tests/dpa4_spin_harness.py new file mode 100644 index 0000000000..c48eb5f600 --- /dev/null +++ b/source/lmp/tests/dpa4_spin_harness.py @@ -0,0 +1,542 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared harness for the DPA4 native-spin LAMMPS ``.pt2`` tests. + +Owns the mechanics common to every DPA4 native-spin variant -- the 4-atom +NiO system, the live ``DeepPot`` reference, LAMMPS construction, magnetic +force gathering, the MPI subprocess/parser, and the single-rank and +rank-parity assertions -- so each variant module carries only its archive, +its geometry and its variant-specific claims. + +Narrowly scoped rather than folded into ``lammps_test_utils``: the pieces +here are DPA4-native-spin-specific (``atom_modify map yes`` for the graph +route's atom map, the ``spin_norm / hbar`` ``fm`` convention, and a +timeout-bounded runner that SIGKILLs the whole mpirun process group) and +would change behaviour for the unrelated tests that share that module. + +Consumers: ``test_lammps_dpa4_spin_graph_pt2.py`` (plain native spin) and +``test_lammps_dpa4_spin_zbl_pt2.py`` (native spin + ZBL bridging). +""" + +import json +import os +import signal +import subprocess as sp +import sys +import tempfile +import textwrap +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import constants +import numpy as np +import pytest +from lammps import ( + PyLammps, +) + +MPI_DEFAULT_TIMEOUT = 120.0 + +# LAMMPS's ``fm`` (what ``compute property/atom fmx fmy fmz`` reports) is +# NOT the raw DeepEval force_mag: pair_deepspin.cpp scales it by +# ``spin_norm / hbar`` per atom (metal-units ``hbar = 6.5821191e-04``, see +# ``source/lmp/pair_deepspin.cpp:531,535``). ``spin_norm`` is 0 for the two +# non-magnetic O atoms, so the scaling is a no-op there (0 stays 0). +HBAR_METAL = 6.5821191e-04 + +# The 4-atom NiO system every DPA4 native-spin variant shares (box, +# coordinates and LAMMPS type ordering reused verbatim from +# test_lammps_spin_pt2.py): 2 Ni atoms (LAMMPS type 1, deepmd atype 0, +# spin-active) + 2 O atoms (LAMMPS type 2, deepmd atype 1, non-magnetic) -- +# matching the archives' ``type_map=["Ni", "O"]`` and +# ``use_spin=[True, False]``. The Ni-Ni pair (atoms 0 and 1) sits ~0.978 A +# apart, which is inside the ZBL bridging transition zone (0.8, 1.2) used by +# the bridged variant. +BOX = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +COORD = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +SPIN = np.array( + [ + [0, 0, 1.2737], + [0, 0, 1.2737], + [0, 0, 0], + [0, 0, 0], + ] +) +TYPE_NIO = np.array([1, 1, 2, 2]) + +# LAMMPS Voigt-ish component order used by the virial/pressure checks. +_VIRIAL_ORDER = [0, 4, 8, 3, 6, 7, 1, 2, 5] + + +def cell_from_lammps_box(lmp_box: np.ndarray) -> np.ndarray: + """Convert a LAMMPS ``xlo xhi ylo yhi zlo zhi xy xz yz`` box spec to a + flat, row-major 3x3 cell matrix (deepmd's ``box`` convention). + """ + xlo, xhi, ylo, yhi, zlo, zhi, xy, xz, yz = lmp_box + return np.array( + [ + xhi - xlo, + 0.0, + 0.0, + xy, + yhi - ylo, + 0.0, + xz, + yz, + zhi - zlo, + ] + ) + + +def compute_expected( + pb_file: Path, + *, + box: np.ndarray = BOX, + coord: np.ndarray = COORD, + spin: np.ndarray = SPIN, + type_map: np.ndarray = TYPE_NIO, +) -> dict[str, Any]: + """Evaluate ``pb_file`` on the fixed system to obtain the reference. + + Runs in a subprocess to avoid importing ``deepmd`` in the LAMMPS test + process (see ``test_lammps_model_devi_pt2.py``'s ``_compute_expected`` + for the same precaution: the LAMMPS plugin already loads + ``libdeepmd_op_pt.so`` at the C++ level, and importing the Python + package on top of that can segfault). + + The archive lives in ``source/tests/infer`` next to ``gen_common.py``, + whose ``load_custom_ops()`` loads the build-tree ``libdeepmd_op_pt.so`` + (registering ``deepmd::edge_force_virial``, which the graph ``.pt2`` + inference needs). ``import deepmd.pt`` alone only loads the op library + from SHARED_LIB_DIR, which the build-test env does not populate -- so + the subprocess reuses that fallback (after importing ``deepmd.pt``, per + its docstring) before constructing ``DeepPot``. + + Parameters + ---------- + pb_file : Path + The ``.pt2`` archive to evaluate. + box : np.ndarray + LAMMPS box spec of the system. + coord : np.ndarray + Per-atom coordinates. + spin : np.ndarray + Per-atom spin vectors. + type_map : np.ndarray + LAMMPS 1-based per-atom types. + + Returns + ------- + dict + ``e`` (energy), ``ae`` (atom energies), ``f`` (forces), ``fm`` + (magnetic forces, already scaled to LAMMPS's ``fm`` convention) and + ``v`` (per-atom virial, sign-flipped to LAMMPS's convention). + """ + cell = cell_from_lammps_box(box) + atype = (type_map - 1).tolist() # LAMMPS 1-based -> deepmd 0-based + infer_dir = str(pb_file.resolve().parent) + script = textwrap.dedent(f"""\ + import json + import sys + import numpy as np + + sys.path.insert(0, {infer_dir!r}) + import deepmd.pt # noqa: F401 (triggers the base op-library load) + from gen_common import load_custom_ops + + load_custom_ops() + from deepmd.infer import DeepPot + + dp = DeepPot({str(pb_file.resolve())!r}) + e, f, v, ae, av, fm, mm = dp.eval( + np.array({coord.tolist()!r}).reshape(1, -1, 3), + np.array({cell.tolist()!r}).reshape(1, 9), + {atype!r}, + atomic=True, + spin=np.array({spin.tolist()!r}).reshape(1, -1, 3), + ) + print(json.dumps({{ + "e": float(e[0, 0]), + "ae": np.asarray(ae[0]).reshape(-1).tolist(), + "f": np.asarray(f[0]).tolist(), + "fm": np.asarray(fm[0]).tolist(), + "av": np.asarray(av[0]).tolist(), + }})) + """) + proc = sp.run([sys.executable, "-c", script], capture_output=True, text=True) + if proc.returncode != 0: + raise RuntimeError(f"Failed to compute expected values:\n{proc.stderr}") + result = json.loads(proc.stdout.strip()) + + # Raw DeepEval force_mag (dE/dspin), scaled by LAMMPS's own + # spin_norm / hbar unit convention (see HBAR_METAL) before comparison. + fm_raw = np.array(result["fm"]) + spin_norm = np.linalg.norm(spin, axis=1) + return { + "e": result["e"], + "ae": np.array(result["ae"]), + "f": np.array(result["f"]), + "fm": fm_raw * (spin_norm / HBAR_METAL)[:, None], + # Per-atom virial, sign-flipped (LAMMPS convention) relative to + # DeepPot's atomic virial output (mirrors test_lammps_spin_pt2.py). + "v": -np.array(result["av"]), + } + + +def make_lammps(data_file: Path, units: str = "metal") -> PyLammps: + """Standard DeepSpin LAMMPS system, plus ``atom_modify map yes``. + + Mirrors ``lammps_test_utils.make_spin_lammps`` (not reused directly: it + does not set ``atom_modify``), with the map turned on -- the native-spin + DPA4 GRAPH ``.pt2`` needs the LAMMPS atom-map to resolve ghost-atom + indices to local owners for single-rank inference (same requirement as + the energy graph route; see ``pair_deepspin.cpp``'s + ``DeePMD-kit Error: Single-rank LAMMPS .pt2 inference requires + `atom_modify map yes``` check). + + Parameters + ---------- + data_file : Path + LAMMPS data file to read. + units : str + Unit system; only ``"metal"`` is valid for spin. + + Returns + ------- + PyLammps + The constructed LAMMPS instance. + + Raises + ------ + ValueError + If ``units`` is not ``"metal"``. + """ + if units != "metal": + raise ValueError("units for spin should be metal") + + lammps = PyLammps() + lammps.units(units) + lammps.boundary("p p p") + lammps.atom_style("spin") + lammps.atom_modify("map yes") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + lammps.mass("1 58") + lammps.mass("2 16") + lammps.timestep(0.0005) + lammps.fix("1 all nve") + return lammps + + +def gather_force_mag(lammps: PyLammps, natoms: int) -> np.ndarray: + """Extract per-atom force_mag in atom-id order. + + LAMMPS does not expose ``fm`` through the legacy ``extract``/ + ``gather_atoms`` registry (see ``run_mpi_pair_deepmd_spin_dpa3_pt2.py``'s + module docstring), so go via ``compute property/atom fmx fmy fmz`` + + ``gather`` (id-ordered on every rank, single-rank included). + + Parameters + ---------- + lammps : PyLammps + A LAMMPS instance with an ``fmprop`` compute defined. + natoms : int + Number of atoms. + + Returns + ------- + np.ndarray + ``(natoms, 3)`` magnetic forces in atom-id order. + """ + fm_global = lammps.lmp.gather("c_fmprop", 1, 3) + return np.array(fm_global, dtype=np.float64).reshape(natoms, 3) + + +def run_mpi_spin_runner( + mpi_runner: Path, + pb_file: Path, + data_path: Path, + *, + nprocs: int = 2, + processors: str | None = None, + extra_args: list[str] | None = None, + capture: bool = False, + timeout: float | None = None, +) -> dict: + """Invoke the graph-spin MPI runner under ``mpirun -n ``. + + With ``capture=True``, return raw subprocess info (``returncode``, + ``stdout``, ``stderr``, ``timed_out``). Every invocation is bounded by + ``timeout`` (default ``MPI_DEFAULT_TIMEOUT``) so a should-fail-but- + doesn't run cannot hang the suite, and on expiry the WHOLE mpirun + process group is SIGKILLed. + + Parameters + ---------- + mpi_runner : Path + The runner script to launch. + pb_file : Path + The ``.pt2`` archive under test. + data_path : Path + LAMMPS data file for this run. + nprocs : int + Number of MPI ranks. + processors : str, optional + LAMMPS ``processors`` grid; defaults to ``1 1 1`` when ``nprocs==1``. + extra_args : list of str, optional + Extra runner arguments. + capture : bool + Return raw subprocess info instead of parsing the output. + timeout : float, optional + Wall-clock bound; defaults to ``MPI_DEFAULT_TIMEOUT``. + + Returns + ------- + dict + ``{"pe": float, "rows": np.ndarray}``, or the raw subprocess info + when ``capture`` is true. + + Raises + ------ + RuntimeError + If a non-``capture`` run exceeds ``timeout``. + subprocess.CalledProcessError + If a non-``capture`` run exits non-zero. + """ + if timeout is None: + timeout = MPI_DEFAULT_TIMEOUT + with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: + out_path = f.name + try: + argv = [ + "mpirun", + "-n", + str(nprocs), + sys.executable, + str(mpi_runner), + str(data_path.resolve()), + str(pb_file.resolve()), + out_path, + ] + if processors is not None: + argv.extend(["--processors", processors]) + elif nprocs == 1: + argv.extend(["--processors", "1 1 1"]) + if extra_args: + argv.extend(extra_args) + proc = sp.Popen( + argv, + stdout=sp.PIPE if capture else None, + stderr=sp.PIPE if capture else None, + text=True, + start_new_session=True, + ) + try: + stdout, stderr = proc.communicate(timeout=timeout) + except sp.TimeoutExpired: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + stdout, stderr = proc.communicate() + if capture: + return { + "returncode": None, + "stdout": stdout or "", + "stderr": stderr or "", + "timed_out": True, + } + raise RuntimeError( + f"mpirun timed out after {timeout}s (process group killed); " + "a should-succeed MPI regression is deadlocked." + ) from None + if capture: + return { + "returncode": proc.returncode, + "stdout": stdout, + "stderr": stderr, + "timed_out": False, + } + if proc.returncode != 0: + raise sp.CalledProcessError(proc.returncode, argv) + with open(out_path) as fh: + lines = fh.read().strip().splitlines() + pe = float(lines[0]) + rows = np.array( + [list(map(float, line.split())) for line in lines[1:]], + dtype=np.float64, + ) + return {"pe": pe, "rows": rows} + finally: + if os.path.exists(out_path): + os.remove(out_path) + + +def assert_mpi_matches_single_rank(single: dict, multi: dict) -> None: + """Compare a multi-rank MPI result against the 1-rank one: pe, per-atom + force and per-atom force_mag, all at rtol=atol=1e-10. + + Both runs report LAMMPS's own ``fm`` (already scaled by + ``spin_norm / HBAR_METAL`` inside pair_deepspin.cpp), so the scaling + cancels and the rows compare directly. + + Parameters + ---------- + single : dict + Result of the 1-rank run. + multi : dict + Result of the multi-rank run. + """ + # anti-vacuity: a degenerate fixture (all-zero forces) would make the + # comparison pass for the wrong reason. + assert np.abs(multi["rows"][:, :3]).max() > 1e-6, "forces are trivially zero" + assert np.abs(multi["rows"][:, 3:6]).max() > 1e-6, "force_mag is trivially zero" + + np.testing.assert_allclose( + multi["pe"], single["pe"], rtol=1e-10, atol=1e-10, err_msg="energy" + ) + np.testing.assert_allclose( + multi["rows"][:, :3], + single["rows"][:, :3], + rtol=1e-10, + atol=1e-10, + err_msg="force", + ) + np.testing.assert_allclose( + multi["rows"][:, 3:6], + single["rows"][:, 3:6], + rtol=1e-10, + atol=1e-10, + err_msg="force_mag", + ) + + +def check_single_rank_energy_force( + lammps: PyLammps, + pb_file: Path, + expected: dict, + *, + coord: np.ndarray = COORD, +) -> None: + """Single-rank LAMMPS energy + force + force_mag vs the DeepEval + reference, including the native-spin zero-``fm``-on-non-spin invariant. + + Parameters + ---------- + lammps : PyLammps + The LAMMPS instance from the module's fixture. + pb_file : Path + The ``.pt2`` archive under test. + expected : dict + Reference produced by :func:`compute_expected`. + coord : np.ndarray + The system coordinates (row count sets the atom count). + """ + natoms = coord.shape[0] + lammps.pair_style(f"deepspin {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("fmprop all property/atom fmx fmy fmz") + lammps.run(0) + + assert lammps.eval("pe") == pytest.approx(expected["e"]) + + forces = np.array( + [lammps.atoms[ii].force for ii in range(natoms)], dtype=np.float64 + ) + ids = np.array([lammps.atoms[ii].id for ii in range(natoms)]) + forces = forces[np.argsort(ids)] + np.testing.assert_allclose(forces, expected["f"], atol=1e-8, rtol=0) + + force_mag = gather_force_mag(lammps, natoms) + np.testing.assert_allclose(force_mag, expected["fm"], atol=1e-8, rtol=0) + # Anti-vacuity / native-spin design invariant: force_mag on the two + # non-spin (O) atoms must be exactly zero, both in the Python reference + # (baked into expected["fm"]) and as produced by LAMMPS. + np.testing.assert_array_equal(force_mag[2:], np.zeros((natoms - 2, 3))) + + lammps.run(1) + + +def check_single_rank_virial( + lammps: PyLammps, + pb_file: Path, + expected: dict, + *, + box: np.ndarray = BOX, + coord: np.ndarray = COORD, +) -> None: + """Single-rank per-atom pe/pressure/virial via ``pe/atom`` / + ``pressure`` / ``centroid/stress/atom``, atol=1e-8, rtol=1e-8. + + Parameters + ---------- + lammps : PyLammps + The LAMMPS instance from the module's fixture. + pb_file : Path + The ``.pt2`` archive under test. + expected : dict + Reference produced by :func:`compute_expected`. + box : np.ndarray + LAMMPS box spec (sets the cell volume). + coord : np.ndarray + The system coordinates (row count sets the atom count). + """ + natoms = coord.shape[0] + lammps.pair_style(f"deepspin {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("peatom all pe/atom pair") + lammps.compute("pressure all pressure NULL pair") + lammps.compute("virial all centroid/stress/atom NULL pair") + lammps.variable("eatom atom c_peatom") + for ii in range(9): + jj = _VIRIAL_ORDER[ii] + lammps.variable(f"pressure{jj} equal c_pressure[{ii + 1}]") + for ii in range(9): + jj = _VIRIAL_ORDER[ii] + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") + lammps.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps.run(0) + + assert lammps.eval("pe") == pytest.approx(expected["e"]) + + forces = np.array( + [lammps.atoms[ii].force for ii in range(natoms)], dtype=np.float64 + ) + ids = np.array([lammps.atoms[ii].id for ii in range(natoms)]) + forces = forces[np.argsort(ids)] + np.testing.assert_allclose(forces, expected["f"], atol=1e-8, rtol=0) + + idx_map = lammps.lmp.numpy.extract_atom("id")[:natoms] - 1 + np.testing.assert_allclose( + np.array(lammps.variables["eatom"].value), + expected["ae"][idx_map], + atol=1e-8, + rtol=1e-8, + ) + + vol = box[1] * box[3] * box[5] + for ii in range(6): + jj = _VIRIAL_ORDER[ii] + pressure_jj = np.array(lammps.variables[f"pressure{jj}"].value) / ( + constants.nktv2p + ) + expected_pressure_jj = -expected["v"][idx_map, jj].sum(axis=0) / vol + np.testing.assert_allclose( + pressure_jj, expected_pressure_jj, atol=1e-8, rtol=1e-8 + ) + for ii in range(9): + jj = _VIRIAL_ORDER[ii] + virial_jj = np.array(lammps.variables[f"virial{jj}"].value) / (constants.nktv2p) + np.testing.assert_allclose( + virial_jj, expected["v"][idx_map, jj], atol=1e-8, rtol=1e-8 + ) diff --git a/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py b/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py index 95b012f819..3d76b17aac 100644 --- a/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.py @@ -23,7 +23,7 @@ Reference values are computed LIVE at test-setup time via ``deepmd.infer.DeepPot.eval`` on the archive itself (mirroring -``test_lammps_dpa4_spin_graph_pt2.py``'s ``_compute_expected``, which +``dpa4_spin_harness.compute_expected``, which explains the reasoning in full) rather than read from the generator's ``.expected`` sidecar: the sidecar's evaluation uses a 6x6x6 A cell whose edge length equals DPA4's LAMMPS ghost cutoff exactly diff --git a/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py b/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py index 2603e1afb8..ea8aaf03fb 100644 --- a/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py @@ -28,7 +28,7 @@ 6x6x6 box (see its module docstring / ``_COORDS`` / ``_CELL`` / ``_SPINS``), and that box's edge length (6.0) exactly equals DPA4's ghost cutoff (rcut(4.0)+skin(2.0)=6.0) -- not a safe geometry to reuse for a LAMMPS -periodic run. Instead, ``_compute_expected`` below loads the archive and +periodic run. Instead, ``dpa4_spin_harness.compute_expected`` loads the archive and evaluates it, at test-setup time, on THIS module's own fixed geometry -- mirroring ``test_lammps_model_devi_pt2.py``'s ``_compute_expected`` pattern (subprocess-isolated, so importing ``deepmd``'s Python package does not @@ -40,23 +40,25 @@ """ import importlib.util -import json import os import shutil -import signal -import subprocess as sp -import sys -import tempfile -import textwrap from pathlib import ( Path, ) -import constants import numpy as np import pytest -from lammps import ( - PyLammps, +from dpa4_spin_harness import ( + BOX, + COORD, + SPIN, + TYPE_NIO, + assert_mpi_matches_single_rank, + check_single_rank_energy_force, + check_single_rank_virial, + compute_expected, + make_lammps, + run_mpi_spin_runner, ) from write_lmp_data import ( write_lmp_data_spin, @@ -90,36 +92,14 @@ # phantom path and none trips the genuinely-empty fail-fast. The shifted # coordinates (``coord_empty_rank``) are defined below, after ``coord``. data_file_empty_rank = Path(__file__).parent / "data_dpa4_spin_graph_pt2_empty_rank.lmp" +data_file_empty_rank = Path(__file__).parent / "data_dpa4_spin_graph_pt2_empty_rank.lmp" # The MPI runner is graph-spin-specific (no aparam / no NULL-type # extras, unlike run_mpi_pair_deepmd_spin_dpa3_pt2.py's virtual-atom-scheme # runner): the native-spin DPA4 fixture takes no fparam/aparam. mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_spin_graph_dpa4_pt2.py" -_MPI_DEFAULT_TIMEOUT = 120.0 - -# Same 4-atom NiO system as test_lammps_spin_pt2.py (box, coordinates, and -# LAMMPS type ordering all reused verbatim): 2 Ni atoms (LAMMPS type 1, -# deepmd atype 0, spin-active) + 2 O atoms (LAMMPS type 2, deepmd atype 1, -# non-magnetic) -- matches ``deeppot_dpa4_spin_graph.pt2``'s -# ``type_map=["Ni", "O"]`` and ``use_spin=[True, False]`` (gen_dpa4_spin.py). -box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) -coord = np.array( - [ - [12.83, 2.56, 2.18], - [12.09, 2.87, 2.74], - [3.51, 2.51, 2.60], - [4.27, 3.22, 1.56], - ] -) -spin = np.array( - [ - [0, 0, 1.2737], - [0, 0, 1.2737], - [0, 0, 0], - [0, 0, 0], - ] -) -type_NiO = np.array([1, 1, 2, 2]) +# The shared 4-atom NiO system (dpa4_spin_harness). +box, coord, spin, type_NiO = BOX, COORD, SPIN, TYPE_NIO # Owned-empty-rank variant of ``coord`` (see the comment above # ``data_file_empty_rank``): the two Ni atoms shift to x ~= 26.8/26.1 so @@ -134,118 +114,10 @@ coord_empty_rank = coord.copy() coord_empty_rank[:2, 0] += _EMPTY_RANK_NI_X_SHIFT -# LAMMPS's ``fm`` (what ``compute property/atom fmx fmy fmz`` reports) is -# NOT the raw DeepEval force_mag: pair_deepspin.cpp scales it by -# ``spin_norm / hbar`` per atom (metal-units ``hbar = 6.5821191e-04``, see -# ``source/lmp/pair_deepspin.cpp:531,535`` -- same convention already -# implicit, if untested against a raw LAMMPS ``fm`` read, in -# test_lammps_spin_pt2.py). ``spin_norm`` is 0 for the two non-magnetic O -# atoms, so the scaling is a no-op there (0 stays 0). -_HBAR_METAL = 6.5821191e-04 - -# Reference values (energy / atom-energy / force / force_mag / virial), -# populated by ``_compute_expected`` in ``setup_module`` -- see the module -# docstring for why these are computed live via a DeepPot subprocess call -# rather than hardcoded or read from a sidecar file. -expected_e = None -expected_ae = None -expected_f = None -expected_fm = None -expected_v = None - - -def _cell_from_lammps_box(lmp_box: np.ndarray) -> np.ndarray: - """Convert a LAMMPS ``xlo xhi ylo yhi zlo zhi xy xz yz`` box spec to a - flat, row-major 3x3 cell matrix (deepmd's ``box`` convention). - """ - xlo, xhi, ylo, yhi, zlo, zhi, xy, xz, yz = lmp_box - return np.array( - [ - xhi - xlo, - 0.0, - 0.0, - xy, - yhi - ylo, - 0.0, - xz, - yz, - zhi - zlo, - ] - ) - - -def _compute_expected() -> None: - """Load ``deeppot_dpa4_spin_graph.pt2`` via ``DeepPot`` and evaluate the - module's fixed 4-atom NiO system to obtain the Python reference. - - Runs in a subprocess to avoid importing ``deepmd`` in the LAMMPS test - process (see ``test_lammps_model_devi_pt2.py``'s ``_compute_expected`` - for the same precaution: the LAMMPS plugin already loads - ``libdeepmd_op_pt.so`` at the C++ level, and importing the Python - package on top of that can segfault). - """ - global expected_e, expected_ae, expected_f, expected_fm, expected_v - - cell = _cell_from_lammps_box(box) - atype = (type_NiO - 1).tolist() # LAMMPS 1-based -> deepmd 0-based (Ni=0, O=1) - - # ``deeppot_dpa4_spin_graph.pt2`` lives in ``source/tests/infer`` next to - # ``gen_common.py``, whose ``load_custom_ops()`` loads the build-tree - # ``libdeepmd_op_pt.so`` (registering ``deepmd::edge_force_virial``, which - # the graph ``.pt2`` inference needs). ``import deepmd.pt`` alone only loads - # the op library from SHARED_LIB_DIR, which the build-test env does not - # populate -- so the subprocess reuses that fallback (after importing - # ``deepmd.pt``, per its docstring) before constructing ``DeepPot``. - infer_dir = str(pb_file.resolve().parent) - script = textwrap.dedent(f"""\ - import json - import sys - import numpy as np - - sys.path.insert(0, {infer_dir!r}) - import deepmd.pt # noqa: F401 (triggers the base op-library load) - from gen_common import load_custom_ops - - load_custom_ops() - from deepmd.infer import DeepPot - - dp = DeepPot({str(pb_file.resolve())!r}) - e, f, v, ae, av, fm, mm = dp.eval( - np.array({coord.tolist()!r}).reshape(1, -1, 3), - np.array({cell.tolist()!r}).reshape(1, 9), - {atype!r}, - atomic=True, - spin=np.array({spin.tolist()!r}).reshape(1, -1, 3), - ) - print(json.dumps({{ - "e": float(e[0, 0]), - "ae": np.asarray(ae[0]).reshape(-1).tolist(), - "f": np.asarray(f[0]).tolist(), - "fm": np.asarray(fm[0]).tolist(), - "av": np.asarray(av[0]).tolist(), - }})) - """) - proc = sp.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - ) - if proc.returncode != 0: - raise RuntimeError(f"Failed to compute expected values:\n{proc.stderr}") - result = json.loads(proc.stdout.strip()) - - expected_e = result["e"] - expected_ae = np.array(result["ae"]) - expected_f = np.array(result["f"]) - # Raw DeepEval force_mag (dE/dspin), scaled by LAMMPS's own - # spin_norm / hbar unit convention (see the comment on ``_HBAR_METAL`` - # above) before comparison. - fm_raw = np.array(result["fm"]) - spin_norm = np.linalg.norm(spin, axis=1) - expected_fm = fm_raw * (spin_norm / _HBAR_METAL)[:, None] - # Per-atom virial, sign-flipped (LAMMPS convention) relative to DeepPot's - # atomic virial output (mirrors test_lammps_spin_pt2.py's convention). - expected_v = -np.array(result["av"]) +# Reference values, populated by ``compute_expected`` in ``setup_module`` -- +# see the module docstring for why these are computed live via a DeepPot +# subprocess call rather than hardcoded or read from a sidecar file. +expected: dict = {} def setup_module() -> None: @@ -255,7 +127,7 @@ def setup_module() -> None: ) if not pb_file.exists(): pytest.skip("deeppot_dpa4_spin_graph.pt2 not found") - _compute_expected() + expected.update(compute_expected(pb_file)) write_lmp_data_spin(box, coord, spin, type_NiO, data_file) box_empty_rank = np.array([0, 90, 0, 13, 0, 13, 0, 0, 0]) write_lmp_data_spin( @@ -264,235 +136,36 @@ def setup_module() -> None: def teardown_module() -> None: - for f in (data_file, data_file_empty_rank): - if f.exists(): - os.remove(f) - - -def _lammps(data_file, units="metal") -> PyLammps: - """Standard DeepSpin LAMMPS system, plus ``atom_modify map yes``. - - Mirrors ``lammps_test_utils.make_spin_lammps`` (not reused directly: it - does not set ``atom_modify``), with the map turned on -- the native-spin - DPA4 GRAPH ``.pt2`` needs the LAMMPS atom-map to resolve ghost-atom - indices to local owners for single-rank inference (same requirement as - the energy graph route; see ``pair_deepspin.cpp``'s - ``DeePMD-kit Error: Single-rank LAMMPS .pt2 inference requires - `atom_modify map yes``` check). - """ - if units != "metal": - raise ValueError("units for spin should be metal") - - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("spin") - lammps.atom_modify("map yes") - lammps.neighbor("2.0 bin") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - lammps.mass("1 58") - lammps.mass("2 16") - lammps.timestep(0.0005) - lammps.fix("1 all nve") - return lammps + for path in (data_file, data_file_empty_rank): + if path.exists(): + os.remove(path) @pytest.fixture def lammps(): - lmp = _lammps(data_file=data_file) + lmp = make_lammps(data_file) yield lmp lmp.close() -def _gather_force_mag(lammps: PyLammps, natoms: int) -> np.ndarray: - """Extract per-atom force_mag in atom-id order. - - LAMMPS does not expose ``fm`` through the legacy ``extract``/ - ``gather_atoms`` registry (see ``run_mpi_pair_deepmd_spin_dpa3_pt2.py``'s - module docstring), so go via ``compute property/atom fmx fmy fmz`` + - ``gather`` (id-ordered on every rank, single-rank included). - """ - fm_global = lammps.lmp.gather("c_fmprop", 1, 3) - return np.array(fm_global, dtype=np.float64).reshape(natoms, 3) - - def test_pair_deepspin(lammps) -> None: - """Single-rank LAMMPS energy + force + force_mag vs the Python DeepEval - graph-spin reference (Task 7 path), on the same 4-atom NiO system. + """Single-rank LAMMPS energy + force + force_mag vs the Python + graph-spin DeepEval reference (Task 7 path). """ - lammps.pair_style(f"deepspin {pb_file.resolve()}") - lammps.pair_coeff("* *") - lammps.compute("fmprop all property/atom fmx fmy fmz") - lammps.run(0) - - assert lammps.eval("pe") == pytest.approx(expected_e) - - forces = np.array([lammps.atoms[ii].force for ii in range(4)], dtype=np.float64) - ids = np.array([lammps.atoms[ii].id for ii in range(4)]) - order = np.argsort(ids) - forces = forces[order] - np.testing.assert_allclose(forces, expected_f, atol=1e-8, rtol=0) - - force_mag = _gather_force_mag(lammps, coord.shape[0]) - np.testing.assert_allclose(force_mag, expected_fm, atol=1e-8, rtol=0) - # Anti-vacuity / native-spin design invariant: force_mag on the two - # non-spin (O) atoms must be exactly zero, both in the Python reference - # (baked into expected_fm above) and as produced by LAMMPS. - np.testing.assert_array_equal(force_mag[2:], np.zeros((2, 3))) - - lammps.run(1) + check_single_rank_energy_force(lammps, pb_file, expected) def test_pair_deepspin_virial(lammps) -> None: - """Single-rank per-atom pe/pressure/virial via - ``pe/atom`` / ``pressure`` / ``centroid/stress/atom``, atol=1e-8, - rtol=1e-8. - """ - lammps.pair_style(f"deepspin {pb_file.resolve()}") - lammps.pair_coeff("* *") - lammps.compute("peatom all pe/atom pair") - lammps.compute("pressure all pressure NULL pair") - lammps.compute("virial all centroid/stress/atom NULL pair") - lammps.variable("eatom atom c_peatom") - for ii in range(9): - jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"pressure{jj} equal c_pressure[{ii + 1}]") - for ii in range(9): - jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") - lammps.dump( - "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) - ) - lammps.run(0) - - assert lammps.eval("pe") == pytest.approx(expected_e) + """Single-rank per-atom pe/pressure/virial on the native-spin archive.""" + check_single_rank_virial(lammps, pb_file, expected) - forces = np.array([lammps.atoms[ii].force for ii in range(4)], dtype=np.float64) - ids = np.array([lammps.atoms[ii].id for ii in range(4)]) - order = np.argsort(ids) - forces = forces[order] - np.testing.assert_allclose(forces, expected_f, atol=1e-8, rtol=0) - idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 - np.testing.assert_allclose( - np.array(lammps.variables["eatom"].value), - expected_ae[idx_map], - atol=1e-8, - rtol=1e-8, +def _run_mpi(data_path: Path, nprocs: int, processors: str) -> dict: + """This module's binding of the shared runner (archive + runner fixed).""" + return run_mpi_spin_runner( + mpi_runner, pb_file, data_path, nprocs=nprocs, processors=processors ) - vol = box[1] * box[3] * box[5] - for ii in range(6): - jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - pressure_jj = np.array(lammps.variables[f"pressure{jj}"].value) / ( - constants.nktv2p - ) - expected_pressure_jj = -expected_v[idx_map, jj].sum(axis=0) / vol - np.testing.assert_allclose( - pressure_jj, expected_pressure_jj, atol=1e-8, rtol=1e-8 - ) - for ii in range(9): - jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - virial_jj = np.array(lammps.variables[f"virial{jj}"].value) / (constants.nktv2p) - np.testing.assert_allclose( - virial_jj, expected_v[idx_map, jj], atol=1e-8, rtol=1e-8 - ) - - -# --------------------------------------------------------------------------- -# Multi-rank: the native-spin graph .pt2 carries the nested with-comm -# artifact, so a 2-rank run must REPRODUCE the 1-rank result (energy, -# force and force_mag) rather than fail fast. -# --------------------------------------------------------------------------- - - -def _run_mpi_subprocess( - extra_args: list[str] | None = None, - nprocs: int = 2, - data_path: Path | None = None, - processors: str | None = None, - capture: bool = False, - timeout: float | None = None, -) -> dict: - """Invoke the graph-spin MPI runner under ``mpirun -n `` against - the native-spin DPA4 graph ``.pt2``. - - Copied (module-global closure, not imported) from - ``test_lammps_dpa4_graph_pt2.py``'s twin. With ``capture=True``, return - raw subprocess info (``returncode``, ``stdout``, ``stderr``, - ``timed_out``) -- used by the fail-fast test below; every invocation is - bounded by ``timeout`` (default ``_MPI_DEFAULT_TIMEOUT``) so a - should-fail-but-doesn't run cannot hang the suite, and on expiry the - WHOLE mpirun process group is SIGKILLed. - """ - if data_path is None: - data_path = data_file - if timeout is None: - timeout = _MPI_DEFAULT_TIMEOUT - with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: - out_path = f.name - try: - argv = [ - "mpirun", - "-n", - str(nprocs), - sys.executable, - str(mpi_runner), - str(data_path.resolve()), - str(pb_file.resolve()), - out_path, - ] - if processors is not None: - argv.extend(["--processors", processors]) - elif nprocs == 1: - argv.extend(["--processors", "1 1 1"]) - if extra_args: - argv.extend(extra_args) - proc = sp.Popen( - argv, - stdout=sp.PIPE if capture else None, - stderr=sp.PIPE if capture else None, - text=True, - start_new_session=True, - ) - try: - stdout, stderr = proc.communicate(timeout=timeout) - except sp.TimeoutExpired: - os.killpg(os.getpgid(proc.pid), signal.SIGKILL) - stdout, stderr = proc.communicate() - if capture: - return { - "returncode": None, - "stdout": stdout or "", - "stderr": stderr or "", - "timed_out": True, - } - raise RuntimeError( - f"mpirun timed out after {timeout}s (process group killed); " - "a should-succeed MPI regression is deadlocked." - ) from None - if capture: - return { - "returncode": proc.returncode, - "stdout": stdout, - "stderr": stderr, - "timed_out": False, - } - if proc.returncode != 0: - raise sp.CalledProcessError(proc.returncode, argv) - with open(out_path) as fh: - lines = fh.read().strip().splitlines() - pe = float(lines[0]) - rows = np.array( - [list(map(float, line.split())) for line in lines[1:]], - dtype=np.float64, - ) - return {"pe": pe, "rows": rows} - finally: - if os.path.exists(out_path): - os.remove(out_path) - @pytest.mark.skipif( shutil.which("mpirun") is None, reason="MPI is not installed on this system" @@ -515,31 +188,9 @@ def test_pair_deepspin_mpi_matches_single_rank() -> None: Replaces the previous fail-fast test, which asserted the C++ throw that existed only while native spin was excluded from the with-comm export. """ - single = _run_mpi_subprocess(nprocs=1, processors="1 1 1") - multi = _run_mpi_subprocess(nprocs=2, processors="2 1 1") - - # anti-vacuity: a degenerate fixture (all-zero forces) would make the - # comparison pass for the wrong reason. - assert np.abs(multi["rows"][:, :3]).max() > 1e-6, "forces are trivially zero" - assert np.abs(multi["rows"][:, 3:6]).max() > 1e-6, "force_mag is trivially zero" - - np.testing.assert_allclose( - multi["pe"], single["pe"], rtol=1e-10, atol=1e-10, err_msg="energy" - ) - np.testing.assert_allclose( - multi["rows"][:, :3], - single["rows"][:, :3], - rtol=1e-10, - atol=1e-10, - err_msg="force", - ) - np.testing.assert_allclose( - multi["rows"][:, 3:6], - single["rows"][:, 3:6], - rtol=1e-10, - atol=1e-10, - err_msg="force_mag", - ) + single = _run_mpi(data_file, 1, "1 1 1") + multi = _run_mpi(data_file, 2, "2 1 1") + assert_mpi_matches_single_rank(single, multi) @pytest.mark.skipif( @@ -574,32 +225,6 @@ def test_pair_deepspin_mpi_empty_rank_phantom_pads_and_matches() -> None: at the file's MPI tolerances -- force_mag only exists on this route, so comparing it is what proves the spin leaf survives phantom padding. """ - single = _run_mpi_subprocess( - nprocs=1, processors="1 1 1", data_path=data_file_empty_rank - ) - multi = _run_mpi_subprocess( - nprocs=3, processors="3 1 1", data_path=data_file_empty_rank - ) - - # anti-vacuity: a degenerate fixture (all-zero forces) would make the - # comparison pass for the wrong reason. - assert np.abs(multi["rows"][:, :3]).max() > 1e-6, "forces are trivially zero" - assert np.abs(multi["rows"][:, 3:6]).max() > 1e-6, "force_mag is trivially zero" - - np.testing.assert_allclose( - multi["pe"], single["pe"], rtol=1e-10, atol=1e-10, err_msg="energy" - ) - np.testing.assert_allclose( - multi["rows"][:, :3], - single["rows"][:, :3], - rtol=1e-10, - atol=1e-10, - err_msg="force", - ) - np.testing.assert_allclose( - multi["rows"][:, 3:6], - single["rows"][:, 3:6], - rtol=1e-10, - atol=1e-10, - err_msg="force_mag", - ) + single = _run_mpi(data_file_empty_rank, 1, "1 1 1") + multi = _run_mpi(data_file_empty_rank, 3, "3 1 1") + assert_mpi_matches_single_rank(single, multi) diff --git a/source/lmp/tests/test_lammps_dpa4_spin_zbl_pt2.py b/source/lmp/tests/test_lammps_dpa4_spin_zbl_pt2.py index a287553564..234b4571cc 100644 --- a/source/lmp/tests/test_lammps_dpa4_spin_zbl_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_spin_zbl_pt2.py @@ -29,30 +29,32 @@ hardcoded (a hardcoded reference goes stale the moment DPA4 numerics shift) and NOT read from the generator's ``.expected`` sidecar (its own 6-atom system sits in a 6x6x6 box whose edge exactly equals DPA4's ghost cutoff -- -not a safe geometry for a LAMMPS periodic run). ``_compute_expected`` runs +not a safe geometry for a LAMMPS periodic run). ``dpa4_spin_harness.compute_expected`` runs in a subprocess so importing ``deepmd``'s Python package does not share a process with the LAMMPS plugin's own loaded ``libdeepmd_op_pt.so`` (see ``test_lammps_model_devi_pt2.py``). """ import importlib.util -import json import os import shutil -import signal -import subprocess as sp -import sys -import tempfile -import textwrap from pathlib import ( Path, ) -import constants import numpy as np import pytest -from lammps import ( - PyLammps, +from dpa4_spin_harness import ( + BOX, + COORD, + SPIN, + TYPE_NIO, + assert_mpi_matches_single_rank, + check_single_rank_energy_force, + check_single_rank_virial, + compute_expected, + make_lammps, + run_mpi_spin_runner, ) from write_lmp_data import ( write_lmp_data_spin, @@ -72,35 +74,11 @@ # native-spin runner is reused verbatim -- only the archive differs. mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_spin_graph_dpa4_pt2.py" -_MPI_DEFAULT_TIMEOUT = 120.0 - -# Same 4-atom NiO system as test_lammps_dpa4_spin_graph_pt2.py (box, -# coordinates, and LAMMPS type ordering all reused verbatim): 2 Ni atoms -# (LAMMPS type 1, deepmd atype 0, spin-active) + 2 O atoms (LAMMPS type 2, -# deepmd atype 1, non-magnetic) -- matches the archive's -# ``type_map=["Ni", "O"]`` and ``use_spin=[True, False]`` -# (gen_dpa4_spin_zbl.py inherits both from gen_dpa4_spin.py). The Ni-Ni -# pair (atoms 0 and 1) sits ~0.978 A apart -- inside the bridging -# transition zone (0.8, 1.2), so the ZBL channel is active even in the -# single-rank tests. -box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) -coord = np.array( - [ - [12.83, 2.56, 2.18], - [12.09, 2.87, 2.74], - [3.51, 2.51, 2.60], - [4.27, 3.22, 1.56], - ] -) -spin = np.array( - [ - [0, 0, 1.2737], - [0, 0, 1.2737], - [0, 0, 0], - [0, 0, 0], - ] -) -type_NiO = np.array([1, 1, 2, 2]) +# The shared 4-atom NiO system (dpa4_spin_harness): the Ni-Ni pair (atoms 0 +# and 1) sits ~0.978 A apart -- inside this variant's bridging transition +# zone (0.8, 1.2), so the ZBL channel is active even in the single-rank +# tests. +box, coord, spin, type_NiO = BOX, COORD, SPIN, TYPE_NIO # Close-pair geometry for the MPI-parity test: the SAME 4-atom system # shifted along x by -5.9, so the ~0.978 A Ni-Ni pair lands at x = 6.93 / @@ -113,116 +91,10 @@ coord_close_pair = coord.copy() coord_close_pair[:, 0] = np.mod(coord[:, 0] + _CLOSE_PAIR_SHIFT_X, box[1]) -# LAMMPS's ``fm`` (what ``compute property/atom fmx fmy fmz`` reports) is -# NOT the raw DeepEval force_mag: pair_deepspin.cpp scales it by -# ``spin_norm / hbar`` per atom (metal-units ``hbar = 6.5821191e-04``, see -# ``source/lmp/pair_deepspin.cpp:531,535``). ``spin_norm`` is 0 for the two -# non-magnetic O atoms, so the scaling is a no-op there (0 stays 0). -_HBAR_METAL = 6.5821191e-04 - -# Reference values (energy / atom-energy / force / force_mag / virial), -# populated by ``_compute_expected`` in ``setup_module`` -- see the module -# docstring for why these are computed live via a DeepPot subprocess call -# rather than hardcoded or read from a sidecar file. -expected_e = None -expected_ae = None -expected_f = None -expected_fm = None -expected_v = None - - -def _cell_from_lammps_box(lmp_box: np.ndarray) -> np.ndarray: - """Convert a LAMMPS ``xlo xhi ylo yhi zlo zhi xy xz yz`` box spec to a - flat, row-major 3x3 cell matrix (deepmd's ``box`` convention). - """ - xlo, xhi, ylo, yhi, zlo, zhi, xy, xz, yz = lmp_box - return np.array( - [ - xhi - xlo, - 0.0, - 0.0, - xy, - yhi - ylo, - 0.0, - xz, - yz, - zhi - zlo, - ] - ) - - -def _compute_expected() -> None: - """Load ``deeppot_dpa4_spin_zbl_graph.pt2`` via ``DeepPot`` and evaluate - the module's fixed 4-atom NiO system to obtain the Python reference. - - Runs in a subprocess to avoid importing ``deepmd`` in the LAMMPS test - process (see ``test_lammps_model_devi_pt2.py``'s ``_compute_expected`` - for the same precaution: the LAMMPS plugin already loads - ``libdeepmd_op_pt.so`` at the C++ level, and importing the Python - package on top of that can segfault). - """ - global expected_e, expected_ae, expected_f, expected_fm, expected_v - - cell = _cell_from_lammps_box(box) - atype = (type_NiO - 1).tolist() # LAMMPS 1-based -> deepmd 0-based (Ni=0, O=1) - - # The archive lives in ``source/tests/infer`` next to ``gen_common.py``, - # whose ``load_custom_ops()`` loads the build-tree ``libdeepmd_op_pt.so`` - # (registering ``deepmd::edge_force_virial``, which the graph ``.pt2`` - # inference needs). ``import deepmd.pt`` alone only loads the op library - # from SHARED_LIB_DIR, which the build-test env does not populate -- so - # the subprocess reuses that fallback (after importing ``deepmd.pt``, per - # its docstring) before constructing ``DeepPot``. - infer_dir = str(pb_file.resolve().parent) - script = textwrap.dedent(f"""\ - import json - import sys - import numpy as np - - sys.path.insert(0, {infer_dir!r}) - import deepmd.pt # noqa: F401 (triggers the base op-library load) - from gen_common import load_custom_ops - - load_custom_ops() - from deepmd.infer import DeepPot - - dp = DeepPot({str(pb_file.resolve())!r}) - e, f, v, ae, av, fm, mm = dp.eval( - np.array({coord.tolist()!r}).reshape(1, -1, 3), - np.array({cell.tolist()!r}).reshape(1, 9), - {atype!r}, - atomic=True, - spin=np.array({spin.tolist()!r}).reshape(1, -1, 3), - ) - print(json.dumps({{ - "e": float(e[0, 0]), - "ae": np.asarray(ae[0]).reshape(-1).tolist(), - "f": np.asarray(f[0]).tolist(), - "fm": np.asarray(fm[0]).tolist(), - "av": np.asarray(av[0]).tolist(), - }})) - """) - proc = sp.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - ) - if proc.returncode != 0: - raise RuntimeError(f"Failed to compute expected values:\n{proc.stderr}") - result = json.loads(proc.stdout.strip()) - - expected_e = result["e"] - expected_ae = np.array(result["ae"]) - expected_f = np.array(result["f"]) - # Raw DeepEval force_mag (dE/dspin), scaled by LAMMPS's own - # spin_norm / hbar unit convention (see the comment on ``_HBAR_METAL`` - # above) before comparison. - fm_raw = np.array(result["fm"]) - spin_norm = np.linalg.norm(spin, axis=1) - expected_fm = fm_raw * (spin_norm / _HBAR_METAL)[:, None] - # Per-atom virial, sign-flipped (LAMMPS convention) relative to DeepPot's - # atomic virial output (mirrors test_lammps_spin_pt2.py's convention). - expected_v = -np.array(result["av"]) +# Reference values, populated by ``compute_expected`` in ``setup_module`` -- +# see the module docstring for why these are computed live via a DeepPot +# subprocess call rather than hardcoded or read from a sidecar file. +expected: dict = {} def setup_module() -> None: @@ -232,7 +104,7 @@ def setup_module() -> None: ) if not pb_file.exists(): pytest.skip("deeppot_dpa4_spin_zbl_graph.pt2 not found") - _compute_expected() + expected.update(compute_expected(pb_file)) write_lmp_data_spin(box, coord, spin, type_NiO, data_file) write_lmp_data_spin(box, coord_close_pair, spin, type_NiO, data_file_close_pair) @@ -243,137 +115,27 @@ def teardown_module() -> None: os.remove(path) -def _lammps(data_file, units="metal") -> PyLammps: - """Standard DeepSpin LAMMPS system, plus ``atom_modify map yes``. - - Mirrors ``lammps_test_utils.make_spin_lammps`` (not reused directly: it - does not set ``atom_modify``), with the map turned on -- the native-spin - DPA4 GRAPH ``.pt2`` needs the LAMMPS atom-map to resolve ghost-atom - indices to local owners for single-rank inference (same requirement as - the energy graph route; see ``pair_deepspin.cpp``'s - ``DeePMD-kit Error: Single-rank LAMMPS .pt2 inference requires - `atom_modify map yes``` check). - """ - if units != "metal": - raise ValueError("units for spin should be metal") - - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("spin") - lammps.atom_modify("map yes") - lammps.neighbor("2.0 bin") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - lammps.mass("1 58") - lammps.mass("2 16") - lammps.timestep(0.0005) - lammps.fix("1 all nve") - return lammps - - @pytest.fixture def lammps(): - lmp = _lammps(data_file=data_file) + lmp = make_lammps(data_file) yield lmp lmp.close() -def _gather_force_mag(lammps: PyLammps, natoms: int) -> np.ndarray: - """Extract per-atom force_mag in atom-id order. - - LAMMPS does not expose ``fm`` through the legacy ``extract``/ - ``gather_atoms`` registry (see ``run_mpi_pair_deepmd_spin_dpa3_pt2.py``'s - module docstring), so go via ``compute property/atom fmx fmy fmz`` + - ``gather`` (id-ordered on every rank, single-rank included). - """ - fm_global = lammps.lmp.gather("c_fmprop", 1, 3) - return np.array(fm_global, dtype=np.float64).reshape(natoms, 3) - - def test_pair_deepspin(lammps) -> None: """Single-rank LAMMPS energy + force + force_mag vs the Python DeepEval - reference on the spin+ZBL archive, on the same 4-atom NiO system. - """ - lammps.pair_style(f"deepspin {pb_file.resolve()}") - lammps.pair_coeff("* *") - lammps.compute("fmprop all property/atom fmx fmy fmz") - lammps.run(0) - - assert lammps.eval("pe") == pytest.approx(expected_e) - - forces = np.array([lammps.atoms[ii].force for ii in range(4)], dtype=np.float64) - ids = np.array([lammps.atoms[ii].id for ii in range(4)]) - order = np.argsort(ids) - forces = forces[order] - np.testing.assert_allclose(forces, expected_f, atol=1e-8, rtol=0) - - force_mag = _gather_force_mag(lammps, coord.shape[0]) - np.testing.assert_allclose(force_mag, expected_fm, atol=1e-8, rtol=0) - # Anti-vacuity / native-spin design invariant: force_mag on the two - # non-spin (O) atoms must be exactly zero, both in the Python reference - # (baked into expected_fm above) and as produced by LAMMPS. The - # analytical ZBL child, which knows nothing about spin, must not leak - # into this channel either. - np.testing.assert_array_equal(force_mag[2:], np.zeros((2, 3))) - - lammps.run(1) - + reference on the spin+ZBL archive. -def test_pair_deepspin_virial(lammps) -> None: - """Single-rank per-atom pe/pressure/virial via - ``pe/atom`` / ``pressure`` / ``centroid/stress/atom``, atol=1e-8, - rtol=1e-8. + The shared check also pins force_mag == 0 on the two non-spin (O) + atoms: the analytical ZBL child, which knows nothing about spin, must + not leak into that channel. """ - lammps.pair_style(f"deepspin {pb_file.resolve()}") - lammps.pair_coeff("* *") - lammps.compute("peatom all pe/atom pair") - lammps.compute("pressure all pressure NULL pair") - lammps.compute("virial all centroid/stress/atom NULL pair") - lammps.variable("eatom atom c_peatom") - for ii in range(9): - jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"pressure{jj} equal c_pressure[{ii + 1}]") - for ii in range(9): - jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") - lammps.dump( - "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) - ) - lammps.run(0) + check_single_rank_energy_force(lammps, pb_file, expected) - assert lammps.eval("pe") == pytest.approx(expected_e) - forces = np.array([lammps.atoms[ii].force for ii in range(4)], dtype=np.float64) - ids = np.array([lammps.atoms[ii].id for ii in range(4)]) - order = np.argsort(ids) - forces = forces[order] - np.testing.assert_allclose(forces, expected_f, atol=1e-8, rtol=0) - - idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 - np.testing.assert_allclose( - np.array(lammps.variables["eatom"].value), - expected_ae[idx_map], - atol=1e-8, - rtol=1e-8, - ) - - vol = box[1] * box[3] * box[5] - for ii in range(6): - jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - pressure_jj = np.array(lammps.variables[f"pressure{jj}"].value) / ( - constants.nktv2p - ) - expected_pressure_jj = -expected_v[idx_map, jj].sum(axis=0) / vol - np.testing.assert_allclose( - pressure_jj, expected_pressure_jj, atol=1e-8, rtol=1e-8 - ) - for ii in range(9): - jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - virial_jj = np.array(lammps.variables[f"virial{jj}"].value) / (constants.nktv2p) - np.testing.assert_allclose( - virial_jj, expected_v[idx_map, jj], atol=1e-8, rtol=1e-8 - ) +def test_pair_deepspin_virial(lammps) -> None: + """Single-rank per-atom pe/pressure/virial on the spin+ZBL archive.""" + check_single_rank_virial(lammps, pb_file, expected) # --------------------------------------------------------------------------- @@ -384,120 +146,10 @@ def test_pair_deepspin_virial(lammps) -> None: # --------------------------------------------------------------------------- -def _run_mpi_subprocess( - extra_args: list[str] | None = None, - nprocs: int = 2, - data_path: Path | None = None, - processors: str | None = None, - capture: bool = False, - timeout: float | None = None, -) -> dict: - """Invoke the graph-spin MPI runner under ``mpirun -n `` against - the spin+ZBL DPA4 graph ``.pt2``. - - Copied (module-global closure, not imported) from - ``test_lammps_dpa4_spin_graph_pt2.py``'s twin. With ``capture=True``, - return raw subprocess info (``returncode``, ``stdout``, ``stderr``, - ``timed_out``); every invocation is bounded by ``timeout`` (default - ``_MPI_DEFAULT_TIMEOUT``) so a should-fail-but-doesn't run cannot hang - the suite, and on expiry the WHOLE mpirun process group is SIGKILLed. - """ - if data_path is None: - data_path = data_file - if timeout is None: - timeout = _MPI_DEFAULT_TIMEOUT - with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: - out_path = f.name - try: - argv = [ - "mpirun", - "-n", - str(nprocs), - sys.executable, - str(mpi_runner), - str(data_path.resolve()), - str(pb_file.resolve()), - out_path, - ] - if processors is not None: - argv.extend(["--processors", processors]) - elif nprocs == 1: - argv.extend(["--processors", "1 1 1"]) - if extra_args: - argv.extend(extra_args) - proc = sp.Popen( - argv, - stdout=sp.PIPE if capture else None, - stderr=sp.PIPE if capture else None, - text=True, - start_new_session=True, - ) - try: - stdout, stderr = proc.communicate(timeout=timeout) - except sp.TimeoutExpired: - os.killpg(os.getpgid(proc.pid), signal.SIGKILL) - stdout, stderr = proc.communicate() - if capture: - return { - "returncode": None, - "stdout": stdout or "", - "stderr": stderr or "", - "timed_out": True, - } - raise RuntimeError( - f"mpirun timed out after {timeout}s (process group killed); " - "a should-succeed MPI regression is deadlocked." - ) from None - if capture: - return { - "returncode": proc.returncode, - "stdout": stdout, - "stderr": stderr, - "timed_out": False, - } - if proc.returncode != 0: - raise sp.CalledProcessError(proc.returncode, argv) - with open(out_path) as fh: - lines = fh.read().strip().splitlines() - pe = float(lines[0]) - rows = np.array( - [list(map(float, line.split())) for line in lines[1:]], - dtype=np.float64, - ) - return {"pe": pe, "rows": rows} - finally: - if os.path.exists(out_path): - os.remove(out_path) - - -def _assert_mpi_matches_single_rank(single: dict, multi: dict) -> None: - """Compare a 2-rank MPI result against the 1-rank one: pe, per-atom - force and per-atom force_mag, all at rtol=atol=1e-10 (the sibling's MPI - tolerances). Both runs report LAMMPS's own ``fm`` (already scaled by - ``spin_norm / _HBAR_METAL`` inside pair_deepspin.cpp), so the scaling - cancels and the rows compare directly. - """ - # anti-vacuity: a degenerate fixture (all-zero forces) would make the - # comparison pass for the wrong reason. - assert np.abs(multi["rows"][:, :3]).max() > 1e-6, "forces are trivially zero" - assert np.abs(multi["rows"][:, 3:6]).max() > 1e-6, "force_mag is trivially zero" - - np.testing.assert_allclose( - multi["pe"], single["pe"], rtol=1e-10, atol=1e-10, err_msg="energy" - ) - np.testing.assert_allclose( - multi["rows"][:, :3], - single["rows"][:, :3], - rtol=1e-10, - atol=1e-10, - err_msg="force", - ) - np.testing.assert_allclose( - multi["rows"][:, 3:6], - single["rows"][:, 3:6], - rtol=1e-10, - atol=1e-10, - err_msg="force_mag", +def _run_mpi(data_path: Path, nprocs: int, processors: str) -> dict: + """This module's binding of the shared runner (archive + runner fixed).""" + return run_mpi_spin_runner( + mpi_runner, pb_file, data_path, nprocs=nprocs, processors=processors ) @@ -523,9 +175,9 @@ def test_pair_deepspin_mpi_matches_single_rank() -> None: the boundary-straddling bridging-zone pair is exercised by ``test_pair_deepspin_mpi_close_pair_across_ranks`` below. """ - single = _run_mpi_subprocess(nprocs=1, processors="1 1 1") - multi = _run_mpi_subprocess(nprocs=2, processors="2 1 1") - _assert_mpi_matches_single_rank(single, multi) + single = _run_mpi(data_file, 1, "1 1 1") + multi = _run_mpi(data_file, 2, "2 1 1") + assert_mpi_matches_single_rank(single, multi) @pytest.mark.skipif( @@ -557,10 +209,6 @@ def test_pair_deepspin_mpi_close_pair_across_ranks() -> None: f"boundary: Ni x = {x_lo}, {x_hi}, lx/2 = {half_lx}" ) - single = _run_mpi_subprocess( - nprocs=1, processors="1 1 1", data_path=data_file_close_pair - ) - multi = _run_mpi_subprocess( - nprocs=2, processors="2 1 1", data_path=data_file_close_pair - ) - _assert_mpi_matches_single_rank(single, multi) + single = _run_mpi(data_file_close_pair, 1, "1 1 1") + multi = _run_mpi(data_file_close_pair, 2, "2 1 1") + assert_mpi_matches_single_rank(single, multi) diff --git a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py index d8c9137635..099e236294 100644 --- a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py @@ -28,7 +28,7 @@ Reference values are computed LIVE at test-setup time via ``deepmd.infer.DeepPot.eval`` on the archive itself, mirroring -``test_lammps_dpa4_spin_graph_pt2.py``'s ``_compute_expected`` (which explains +``dpa4_spin_harness.compute_expected`` (which explains the reasoning in full). Two reasons, both load-bearing here: - A hardcoded array goes stale the moment DPA4 numerics shift, and this From 285e3c67ce1fa22a13389a1c4e3414efde41db6f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 4 Aug 2026 22:19:03 +0800 Subject: [PATCH 26/26] fix(test): drop a duplicated assignment and a stale ownership docstring - The harness extraction emitted `data_file_empty_rank` twice in a row: the preserved comment block already carried the assignment and the rewrite template added it again. A runtime no-op, but a patch artifact in the commit whose point was a clean single-owner harness. - TestNativeSpinWithBridging still described `get_standard_model` as the owner of the bridging composition, which is the opposite of the contract this branch establishes: that builder now rejects `bridging_method`, and `get_native_spin_model` routes DPA4/SeZM to `get_sezm_model` and re-classes the composition it returns. --- source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py | 1 - source/tests/pt_expt/model/test_zbl_bridging.py | 10 ++++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py b/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py index ea8aaf03fb..226f7aead5 100644 --- a/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_spin_graph_pt2.py @@ -92,7 +92,6 @@ # phantom path and none trips the genuinely-empty fail-fast. The shifted # coordinates (``coord_empty_rank``) are defined below, after ``coord``. data_file_empty_rank = Path(__file__).parent / "data_dpa4_spin_graph_pt2_empty_rank.lmp" -data_file_empty_rank = Path(__file__).parent / "data_dpa4_spin_graph_pt2_empty_rank.lmp" # The MPI runner is graph-spin-specific (no aparam / no NULL-type # extras, unlike run_mpi_pair_deepmd_spin_dpa3_pt2.py's virtual-atom-scheme # runner): the native-spin DPA4 fixture takes no fparam/aparam. diff --git a/source/tests/pt_expt/model/test_zbl_bridging.py b/source/tests/pt_expt/model/test_zbl_bridging.py index c1bc55f3a4..3fa45a1c9c 100644 --- a/source/tests/pt_expt/model/test_zbl_bridging.py +++ b/source/tests/pt_expt/model/test_zbl_bridging.py @@ -275,10 +275,12 @@ def _spin_system(): class TestNativeSpinWithBridging: """Native spin + analytical bridging compose (review 3649276109). - ``get_standard_model`` OWNS assembling the atomic model, bridging - composition included, and the native-spin wrapper re-classes whatever it - returns -- so the two features combine with no special case: the learned - child consumes ``spin``, the analytical child accepts and ignores it. + ``get_sezm_model`` OWNS the bridging composition (``get_standard_model`` + rejects ``bridging_method``: a composition is not expressible on a + non-composite model type). ``get_native_spin_model`` routes DPA4/SeZM + configs there and then RE-CLASSES the returned composition -- so the two + features combine with no special case: the learned child consumes + ``spin``, the analytical child accepts and ignores it. """ def test_construction_composes_and_keeps_spin(self) -> None: