feat: SFPG cross-rank completion — capabilities, ZBL bridging and spin+ZBL multi-rank (#5939)

Closes #5906.

The DPA4/SeZM Source Freeze Propagation Gate computes each node's `eta_j
= prod over outgoing edges of w(r_e)`; under MPI domain decomposition a
rank only holds edges with owned destinations, so the src-keyed per-node
partials are rank-incomplete and bridged models were single-rank only.
This PR completes the gate across ranks and, as a prerequisite, promotes
the export-time questions to atomic-model capabilities so compositions
answer by aggregation.

## Phase 1 — capability aggregation (issue Task 4)

- Split the conflated descriptor capability:
`has_message_passing_across_ranks` (needs the per-block exchange;
unconditionally true for SeZM) vs the new `supports_edge_parallel` (can
run under domain decomposition).
- Six capabilities on `BaseAtomicModel` with concrete defaults,
descriptor delegation on `DPAtomicModel`, and any/all aggregation on
`LinearEnergyAtomicModel`: `has_message_passing_across_ranks` (any),
`supports_edge_parallel` (all), `dense_lower_supports_comm` (all),
`uses_compact_edge_pairs` (any), `graph_edge_dtype` (float32 iff all
children), `supports_graph_export` (all).
- `forward_lower_graph_exportable_with_comm` hoisted from `EnergyModel`
into `make_model` (one owner, next to the non-comm twin) so
`LinearEnergyModel` compositions can export it.
- The four `serialization.py` helpers now consult the atomic model — no
`isinstance`-on-concrete-model checks, no `.descriptor` walks.
Regression fixed: a linear composition of two DPA2 children now gets its
with-comm artifact (previously denied by wrapper type).
- Composition-safe reach-ins outside serialization: `.pt`-checkpoint
eval no longer crashes on compositions (`ntypes` via the model API),
`enable_compile` degrades gracefully, and pt_expt `get_standard_model`
honors `bridging_method` like its dpmodel twin (`_compose_bridging` is
the single composition owner).

## Phase 2 — SFPG cross-rank completion (issue Tasks 2 and 3)

No new communication machinery: the fix is one extra invocation of the
existing `deepmd_export::border_op_backward` + `border_op` pair (they
are exact transposes, `R = B^T`) on an `(N, 2)` `[log_eta, zero_count]`
tensor before the gate is applied — reverse-accumulate ghost partials
into owners, then broadcast the completed values back. Zero C++ changes.

- `border_op_backward` gains autograd (its gradient is `border_op`'s
forward), so gate gradients cross ranks.
- dpmodel: `compute_edge_src_gate` packs the partials through an
optional `node_partial_exchange` hook; the dpmodel
`_gate_partial_exchange` raises (single-process reference), the pt_expt
subclass implements it on the border-op pair.
- pt backend wired the same way. The red run of the new pt parity test
demonstrated the issue's claim and more: pt's bridged parallel path did
not just compute a silently wrong gate — it crashed outright (the ZBL
injection indexed per-local types with extended ghost `src` indices);
fixed by reading extended types.
- Gates flipped: `supports_edge_parallel` is now `True` for bridged SeZM
in both backends; bridged (and spin+ZBL) graph freezes embed the nested
`forward_lower_with_comm.pt2`.

## Verification

Anti-vacuous discipline throughout: every parity test places a
sub-`r_outer` pair ACROSS the periodic/rank boundary (without it every
cross-rank gate contribution is `log w = 0`), covers both bridging
channels (hard-freeze `zero_count` at 0.4 Å, transition-zone `log_eta`
at ~1 Å), and carries an identity-exchange ablation that must diverge.

- Eager self-comm parity vs the folded reference at rtol/atol 1e-12
(energy, force, and force_mag for the spin variant), pt and pt_expt.
- make_fx traces both border ops symbolically (21-input with-comm ABI
unchanged); freeze embeds the nested artifact for ZBL and spin+ZBL
compositions.
- LAMMPS end-to-end on a Tesla T4: 2-rank vs 1-rank close-pair parity
for ZBL (`pair_style deepmd`) and spin+ZBL (`pair_style deepspin`, incl.
magnetic forces) — the spin+ZBL variant gets its first LAMMPS file. All
24 `*Dpa4Zbl*` C++ gtests pass (CPU + T4).
- Variant-alignment coverage: ZBL empty-rank fail-fast twin, the first
test of the DeepSpin owned-empty phantom path, charge-spin through
`pair_style deepspin`, and default-CLI `dp freeze` resolution
(nlist→graph auto-override + with-comm artifact) for both compositions.

## Known limitations

1. pt eager multi-rank bridging has no true-MPI pt test (no pt `.pth`
LAMMPS ZBL fixtures exist); its parity rung is self-comm.
2. `graph_edge_dtype` composition rule (float32 iff ALL children) is
conservative; fp64 is the universal ABI.
3. `supports_graph_export` keeps the hardcoded `"cuda"` probe inside
pt_expt DPA1 (capability promoted; probe internals unchanged).
4. The `NativeSpinModelKind` marker-base check in
`_needs_with_comm_artifact` remains (a cross-backend family test, not a
concrete-type reach-through).
5. Model-deviation coverage stays absent for all dpa4 variants
(pre-existing; Python `model_devi` has no spin support at all).
6. DeepPot vs DeepSpin empty-rank designs deliberately differ (fail-fast
vs phantom-pad, PR #5485); both are now pinned per variant, not unified.
7. Found while testing, left for a separate fix:
`source/api_c/include/deepmd.hpp` uses `&vec[0]` on possibly-empty
vectors (~33 sites) — undefined behavior that SIGABRTs under
`_GLIBCXX_ASSERTIONS` before the empty-rank guard's message can fire
(benign on non-hardened builds).

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added multi-rank inference for bridged DPA4/SeZM models, including
native-spin and ZBL configurations.
* Improved graph export detection, metadata, edge precision, and
communication-aware export.
  * Added atomic-output-only inference for statistics workflows.
  * Standard model loading now preserves bridging configurations.

* **Bug Fixes**
* Improved handling of atom types, ghost atoms, empty MPI ranks,
charge-spin inputs, and cross-rank calculations.

* **Documentation**
* Updated DPA4 and native-spin documentation for expanded multi-rank and
graph export support.

* **Tests**
* Added regression coverage for MPI parity, graph exports, bridging,
charge-spin behavior, and capability reporting.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

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

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

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

deepmd/dpmodel/descriptor/dpa1.py

@@ -542,6 +542,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()

deepmd/dpmodel/descriptor/dpa2.py

@@ -604,6 +604,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

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,
@@ -2278,20 +2322,27 @@ 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`.
 
         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 self.bridging_switch is None
+        return True
+
+    def supports_edge_parallel(self) -> bool:
+        """Bridging included: multi-rank is supported for every SeZM config.
+
+        The SFPG per-node partials are completed across ranks by
+        ``_gate_partial_exchange`` (reverse-accumulate + broadcast) before
+        the gate is applied (issue #5906).
+        """
+        return True
 
     def dense_lower_supports_comm(self) -> bool:
         """The DPA4 dense (nlist) lower has no comm_dict implementation.
@@ -2316,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
 

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(

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