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

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

@@ -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()

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

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