feat(pt_expt): native spin supports multi-rank on the graph lower

pt_expt refused a with-comm artifact for EVERY ``NativeSpinModelKind``,
making native spin single-rank only -- a divergence from pt, whose
``SeZMNativeSpinModel`` does NOT override ``supports_edge_parallel`` and so
participates in the edge-parallel artifact like any energy model.

Nothing about spin needs its own cross-rank exchange: the spin input is
per-node and its ghost rows arrive via the LAMMPS ``sp`` forward-comm, so
only the per-block ghost FEATURE refresh has to travel -- the same
``border_op`` the energy model already drives.

- ``NativeSpinEnergyModel.forward_lower_graph_exportable_with_comm``: the
  22-input ABI -- the non-comm spin prefix verbatim (``spin`` keeps
  positional slot 10) with the 8 comm tensors appended after the
  conditional fparam/aparam/charge_spin tail, exactly how the energy model
  appends them.  Single make_fx trace, following the energy with-comm
  precedent.
- ``_needs_with_comm_artifact``: native spin is admitted on the GRAPH lower
  and still refused on nlist, which has no spin with-comm wrapper at all.
- The dyn-shape and sample-input builders learn the spin slot for the
  with-comm ABI.
- C++ ``DeepSpinPTExpt::run_model_graph_with_comm`` (twin of the DeepPot
  one) plus the compute_inner branch: extended node set, device
  ``n_local`` for the owned-energy mask, all 8 comm tensors on CPU, and the
  same build-time pair-exclusion seam as the single-rank graph branch.  The
  blanket multi-rank fail-fast becomes a narrow guard for archives frozen
  before native spin participated in the with-comm export.

Tests: the gate (graph True / nlist False) and the 22-input ABI with
``force_mag`` and ``mask_mag`` surviving the with-comm trace.  Two existing
assertions encoded the old behaviour and now expect the artifact.

Author: Han Wang

deepmd/pt_expt/model/native_spin_model.py

@@ -133,8 +133,8 @@ def forward_lower_graph_exportable(
         before the conditional ``fparam``/``aparam`` tail; the conditional
         ``charge_spin`` tail follows at index 13 (combined native-spin +
         charge-spin FiLM models, review 3638047227; ``None`` otherwise).
-        There is NO with-comm variant (single-rank only; multi-rank
-        graph-spin is a follow-up).
+        :meth:`forward_lower_graph_exportable_with_comm` extends this SAME
+        prefix with the 8 comm tensors for multi-rank.
 
         Two-layer make_fx trace, mirroring
         :meth:`~deepmd.pt_expt.model.ener_model.EnergyModel.forward_lower_graph_exportable`:
@@ -278,3 +278,163 @@ 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,
+        spin: 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 the multi-rank graph-spin lower into an exportable module.
+
+        The with-comm counterpart of
+        :meth:`forward_lower_graph_exportable`: same positional prefix,
+        ``spin`` still at index 10, then the 8 comm tensors appended after
+        the conditional ``fparam``/``aparam``/``charge_spin`` tail (indices
+        14-21) -- exactly how
+        :meth:`~deepmd.pt_expt.model.ener_model.EnergyModel.forward_lower_graph_exportable_with_comm`
+        appends them for the energy model.
+
+        ``spin`` is the EXTENDED per-node spin ``(N, 3)``: ghost rows carry
+        their owner's spin, delivered by the LAMMPS ``sp`` forward-comm
+        before the call, so the descriptor's spin embedding sees the same
+        value on every rank that holds the node.  The per-block ghost
+        FEATURE refresh rides ``deepmd_export::border_op`` exactly as in the
+        energy model; spin needs no border exchange of its own because it is
+        an input, not a derived feature.
+
+        Single make_fx trace (the energy with-comm precedent), unlike the
+        two-layer trace of the non-comm spin path: the comm-dict packing,
+        the ``forward_common_lower_graph`` call with ``spin`` as a second
+        autograd leaf, and the public-key translation all live in one traced
+        ``fn``.
+
+        Parameters
+        ----------
+        atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, spin, fparam, aparam, charge_spin
+            As in :meth:`forward_lower_graph_exportable`.
+        send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost
+            The 8 comm tensors, packed into ``comm_dict`` inside the traced
+            function.  Same runtime device contract as the energy model's:
+            ALL 8 stay on CPU (host control metadata for ``border_op``),
+            while the device-side owned count is the separate ``n_local``
+            input at slot 2.
+        do_atomic_virial
+            Whether to also return ``atom_virial``.
+        **make_fx_kwargs
+            Extra keyword arguments forwarded to ``make_fx``.
+
+        Returns
+        -------
+        torch.nn.Module
+            A traced module accepting the 22-input ABI above and returning
+            the same public keys as :meth:`forward_lower_graph_exportable`
+            (``atom_energy``, ``energy``, ``force``, ``force_mag``,
+            ``virial``, ``mask_mag``, plus ``atom_virial`` when requested).
+        """
+        model = self
+
+        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,
+            spin: 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,
+            }
+            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,
+                spin=spin,
+                comm_dict=comm_dict,
+            )
+            # Same single-owner translation as the eager forward and the
+            # non-comm lower, so mask_mag is emitted here too.
+            return self._translate_eager_call(
+                model_ret, atype, do_atomic_virial=do_atomic_virial
+            )
+
+        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,
+            spin,
+            fparam,
+            aparam,
+            charge_spin,
+            send_list,
+            send_proc,
+            recv_proc,
+            send_num,
+            recv_num,
+            communicator,
+            nlocal,
+            nghost,
+        )

deepmd/pt_expt/utils/serialization.py

@@ -148,17 +148,13 @@ def _needs_with_comm_artifact(
     (absent on descriptors, such as dpa2/dpa3, whose dense lower always
     supports comm — treated as ``True``).
 
-    FIRST rule (checked before any descriptor-based logic): the native-spin
-    native-spin model (``NativeSpinEnergyModel``, type ``native_spin`` /
-    ``sezm_native_spin``) always returns ``False``, regardless of
-    ``lower_kind`` or the wrapped backbone descriptor's own
-    ``has_message_passing_across_ranks()``. Ghost-atom SPIN exchange across
-    MPI ranks is not implemented -- the wrapper's graph-spin ``.pt2`` ABI
-    (``forward_lower_graph_exportable``, spin at positional index 10) is
-    single-rank only. The backbone's own (energy-only) descriptor may well
-    report ``True`` for cross-rank message passing -- that capability
-    belongs to the energy-only backbone, not to this wrapper's (unported)
-    spin threading, so it must not leak through.
+    Native spin participates on the GRAPH lower, matching pt's
+    ``SeZMModel.supports_edge_parallel`` (which ``SeZMNativeSpinModel`` does
+    not override): the spin input is per-node and its ghost rows arrive via
+    the LAMMPS ``sp`` forward-comm, so nothing about spin needs its own
+    cross-rank exchange -- the per-block ghost FEATURE refresh is the same
+    ``border_op`` the energy model uses. It is excluded only on the dense
+    (nlist) lower, which has no spin with-comm wrapper.
 
     Parameters
     ----------
@@ -180,7 +176,9 @@ def _needs_with_comm_artifact(
     # Cross-backend family test: the dpmodel and pt_expt concrete classes
     # are parallel factory products with no subclass relation, so the shared
     # marker base -- not a concrete class -- is the membership check.
-    if isinstance(model, NativeSpinModelKind):
+    # Native spin rides the GRAPH lower only; its dense lower has no
+    # with-comm wrapper at all.
+    if isinstance(model, NativeSpinModelKind) and lower_kind != "graph":
         return False
 
     # Analytical bridging models are single-rank only (pt's
@@ -835,6 +833,7 @@ def _build_graph_dynamic_shapes(
 
 def _build_graph_dynamic_shapes_with_comm(
     *sample_inputs: torch.Tensor | None,
+    is_native_spin: bool = False,
 ) -> tuple:
     """Build dynamic-shape specs for the with-comm graph-form export.
 
@@ -855,16 +854,23 @@ def _build_graph_dynamic_shapes_with_comm(
         source_row_ptr, fparam, aparam, charge_spin, send_list, send_proc,
         recv_proc, send_num, recv_num, communicator, nlocal, nghost)`` --
         21 entries matching ``forward_lower_graph_exportable_with_comm``.
+        Native-spin ABI (``is_native_spin=True``): 22 entries, with ``spin``
+        inserted at slot 10 and the conditional tail shifted to 11-13, so
+        the comm block starts at 14.
+    is_native_spin : bool
+        Whether ``sample_inputs`` follows the native-spin positional ABI.
 
     Returns
     -------
     tuple
         Per-input dynamic-shape specs (dicts of ``torch.export.Dim`` or
         ``None``) in the same order as ``sample_inputs``.
     """
-    fparam = sample_inputs[10]
-    aparam = sample_inputs[11]
-    charge_spin = sample_inputs[12]
+    tail_start = 11 if is_native_spin else 10
+    spin = sample_inputs[10] if is_native_spin else None
+    fparam = sample_inputs[tail_start]
+    aparam = sample_inputs[tail_start + 1]
+    charge_spin = sample_inputs[tail_start + 2]
     nframes_val = 1
     n_node_total_dim = torch.export.Dim("n_node_total", min=1)
     nedge_dim = torch.export.Dim("nedge", min=2)
@@ -879,6 +885,13 @@ def _build_graph_dynamic_shapes_with_comm(
         {0: n_node_total_dim + 1},  # destination_row_ptr: (N + 1,)
         {0: nedge_dim},  # source_order: (E,)
         {0: n_node_total_dim + 1},  # source_row_ptr: (N + 1,)
+        # spin: (N, 3) — EXTENDED node axis, shares atype's symbol; present
+        # only in the native-spin ABI, where it occupies slot 10.
+        *(
+            ({0: n_node_total_dim} if spin is not None else None,)
+            if is_native_spin
+            else ()
+        ),
         {0: nframes_val} if fparam is not None else None,  # fparam
         # aparam: (N, nda) — flat on the SAME extended node axis as atype
         # (owned prefix + ghost rows).
@@ -1672,6 +1685,7 @@ def _trace_and_export(
                 dtype=torch.float64,
                 edge_dtype=edge_dtype,
                 device=torch.device("cpu"),
+                want_spin=is_native_spin,
             )
             comm_inputs = _make_comm_sample_inputs(
                 nloc=nlocal_sample,
@@ -1687,7 +1701,9 @@ def _trace_and_export(
                 tracing_mode="symbolic",
                 _allow_non_fake_inputs=True,
             )
-            dynamic_shapes = _build_graph_dynamic_shapes_with_comm(*sample_inputs)
+            dynamic_shapes = _build_graph_dynamic_shapes_with_comm(
+                *sample_inputs, is_native_spin=is_native_spin
+            )
         else:
             edge_dtype = (
                 torch.float32

source/api_cc/include/DeepSpinPTExpt.h

@@ -267,6 +267,32 @@ class DeepSpinPTExpt : public DeepSpinBackend {
       const torch::Tensor& fparam,
       const torch::Tensor& aparam);
 
+  /**
+   * @brief Run the native-spin parallel GRAPH artifact: run_model_graph's
+   * ABI (spin at positional index 10) with the 8 border_op comm tensors
+   * appended after the conditional fparam/aparam/charge_spin tail.
+   *
+   * ``spin`` is the EXTENDED per-node spin -- ghost rows carry their
+   * owner's value from the LAMMPS ``sp`` forward-comm, so spin itself needs
+   * no cross-rank exchange; only the per-block ghost FEATURE refresh rides
+   * ``border_op``.  Twin of ``DeepPotPTExpt::run_model_graph_with_comm``.
+   */
+  std::vector<torch::Tensor> run_model_graph_with_comm(
+      const torch::Tensor& atype,
+      const torch::Tensor& n_node,
+      const torch::Tensor& n_local,
+      const torch::Tensor& edge_index,
+      const torch::Tensor& edge_vec,
+      const torch::Tensor& edge_mask,
+      const torch::Tensor& destination_order,
+      const torch::Tensor& destination_row_ptr,
+      const torch::Tensor& source_order,
+      const torch::Tensor& source_row_ptr,
+      const torch::Tensor& spin,
+      const torch::Tensor& fparam,
+      const torch::Tensor& aparam,
+      const std::vector<at::Tensor>& comm_tensors);
+
   /**
    * @brief Run the native-spin parallel edge artifact: the energy edge
    * with-comm schema (coord and extended types span the extended node set)