refactor(pt_expt): NativeSpinEnergyModel via make_native_spin_model; registry dispatch replaces native-spin seams

The pt_expt twin instantiates the shared factory on THIS backend's
EnergyModel and registers under 'native_spin' in the PT_EXPT BaseModel
registry, so the public BaseModel.deserialize returns the torch class
(backend-aware, no hard-coded branches). Consequences:

- The bespoke deserialize override and the __getattr__ delegation shim
  are deleted: the factory's deserialize closure rebuilds through the
  pt_expt EnergyModel, and inheritance replaces delegation. forward and
  forward_lower_graph_exportable (spin@index-10 ABI unchanged) now call
  the inherited call_common/forward_common_lower_graph_exportable on
  self instead of a backbone attribute.
- get_native_spin_model (public, matching its sibling builders) routes
  the backbone through get_sezm_model for the DPA4/SeZM family (keeping
  its guards) or get_standard_model otherwise, gates on
  descriptor.supports_native_spin(), and re-classes via atomic_model_.
  The standard-typed get_model path now dispatches scheme='native' too,
  so a future native-spin descriptor works from a plain standard config.
- deep_eval and deserialize_to_file drop their native-spin special
  cases (registry + has_spin() cover them); the eval_typeebd unwrap is
  isinstance-precise on the virtual-atom SpinModel only.
- NEW NativeSpinModelKind marker base in the factory module: each
  backend's concrete class is a parallel factory product with NO
  subclass relation, so the old isinstance-against-dpmodel-class
  with-comm gate went silently dead -- the freeze then compiled a
  with-comm artifact for the single-rank-only spin ABI (caught by
  test_native_spin_graph_freeze cold-cache A/B). Seams test the shared
  marker instead.

Validated locally: native-spin+spin batteries 25, export file 4/5
(test_dpa4_freeze_to_pt2[auto-graph] fails identically at the
pre-refactor commit with a cold inductor cache -- pre-existing local
inductor codegen bug, not a regression), deep_eval 93, dpmodel 19.

Author: Han Wang

deepmd/dpmodel/model/native_spin_model.py

@@ -26,6 +26,19 @@
 )
 
 
+class NativeSpinModelKind:
+    """Marker base identifying classes produced by ``make_native_spin_model``.
+
+    Each backend instantiates the factory on its OWN standard model class,
+    so the concrete classes (e.g. dpmodel's and pt_expt's
+    ``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.
+    """
+
+
 def make_native_spin_model(T_Model: type) -> type:
     """Make a native-spin model class from a standard model class.
 
@@ -56,7 +69,7 @@ def make_native_spin_model(T_Model: type) -> type:
         The derived native-spin model class.
     """
 
-    class NSM(T_Model):
+    class NSM(T_Model, NativeSpinModelKind):
         """Native-spin variant of ``T_Model`` (see ``make_native_spin_model``)."""
 
         def __init__(self, *args: Any, spin: Spin, **kwargs: Any) -> None:

deepmd/pt_expt/infer/deep_eval.py

@@ -95,7 +95,7 @@ def _graph_spin_output_key(odef: "OutputVariableDef") -> str | None:
     """Map a native-spin request def to its graph-spin public model key.
 
     Twin of ``_GRAPH_CATEGORY_TO_KEY`` for
-    ``DPA4NativeSpinModel.forward_lower_graph_exportable``'s output dict
+    ``NativeSpinEnergyModel.forward_lower_graph_exportable``'s output dict
     (``_translate_spin_energy_keys``: ``atom_energy``, ``energy``,
     ``force``, ``force_mag``, ``virial``, ``atom_virial``).  Category alone
     is NOT enough here: ``do_derivative`` (``deepmd.dpmodel.output_def``)
@@ -338,19 +338,13 @@ def _init_from_model_json(self, model_json_str: str) -> None:
 
             self._dpmodel = SpinModel.deserialize(model_data)
             self._is_spin = True
-        elif model_type in ("dpa4_native_spin", "sezm_native_spin"):
-            # Native-spin (NeighborGraph route) DPA4/SeZM: no virtual atoms,
-            # spin rides the graph lower directly (Task 6/7 of the DPA4
-            # native-spin plan). Mirrors the ``spin_ener`` branch above.
-            from deepmd.pt_expt.model.dpa4_native_spin_model import (
-                DPA4NativeSpinModel,
-            )
-
-            self._dpmodel = DPA4NativeSpinModel.deserialize(model_data)
-            self._is_spin = True
         else:
+            # Registry-dispatched: wrapper classes registered in the pt_expt
+            # BaseModel registry (e.g. the native-spin models, type
+            # "native_spin") come back as their pt_expt torch classes and
+            # declare spin via the base-model capability method.
             self._dpmodel = BaseModel.deserialize(model_data)
-            self._is_spin = False
+            self._is_spin = self._dpmodel.has_spin()
 
         self._rcut = self._dpmodel.get_rcut()
         self._type_map = self._dpmodel.get_type_map()
@@ -1662,7 +1656,7 @@ def _eval_model_spin(
             # extended/nlist ABI at all -- dispatch to the graph-native fast
             # path (mirrors _eval_model's dispatch to _eval_model_graph for
             # the non-spin case). charge_spin has no slot in this ABI (see
-            # DPA4NativeSpinModel.forward_lower_graph_exportable).
+            # NativeSpinEnergyModel.forward_lower_graph_exportable).
             return self._eval_model_graph_spin(
                 coords, cells, atom_types, spins, fparam, aparam, request_defs
             )
@@ -1851,7 +1845,7 @@ def _eval_model_graph_spin(
         construction (SAME builder, SAME positional ABI up through
         ``source_row_ptr``), then inserts the owned-atom ``spin`` tensor
         ``(N, 3)`` at positional index 10 of
-        ``DPA4NativeSpinModel.forward_lower_graph_exportable`` -- the node
+        ``NativeSpinEnergyModel.forward_lower_graph_exportable`` -- the node
         axis IS the owned-local-atom axis for single-rank eval (no ghost
         nodes), so ``spin`` needs no extension/mapping, unlike the dense
         spin path's ``ext_spin_t``. There is no ``charge_spin`` slot in this
@@ -2353,8 +2347,14 @@ def eval_typeebd(self) -> np.ndarray:
 
         from deepmd.dpmodel.utils.type_embed import TypeEmbedNet as TypeEmbedNetDP
 
+        from deepmd.pt_expt.model.spin_model import (
+            SpinModel,
+        )
+
         model = self._dpmodel
-        if self._is_spin_model():
+        if isinstance(model, SpinModel):
+            # Virtual-atom wrapper: type-embed nets live on the backbone.
+            # Native-spin models ARE the model (is-a); no unwrap.
             model = model.backbone_model
         out = []
         for mm in model.modules():

deepmd/pt_expt/model/__init__.py

@@ -15,8 +15,8 @@
 from .dp_zbl_model import (
     DPZBLModel,
 )
-from .dpa4_native_spin_model import (
-    DPA4NativeSpinModel,
+from .native_spin_model import (
+    NativeSpinEnergyModel,
 )
 from .ener_model import (
     EnergyModel,
@@ -43,12 +43,12 @@
 __all__ = [
     "BaseModel",
     "DOSModel",
-    "DPA4NativeSpinModel",
     "DPZBLModel",
     "DipoleModel",
     "EnergyModel",
     "FrozenModel",
     "LinearEnergyModel",
+    "NativeSpinEnergyModel",
     "PolarModel",
     "PropertyModel",
     "SpinEnergyModel",

deepmd/pt_expt/model/get_model.py

@@ -26,8 +26,8 @@
 from deepmd.pt_expt.model.dos_model import (
     DOSModel,
 )
-from deepmd.pt_expt.model.dpa4_native_spin_model import (
-    DPA4NativeSpinModel,
+from deepmd.pt_expt.model.native_spin_model import (
+    NativeSpinEnergyModel,
 )
 from deepmd.pt_expt.model.ener_model import (
     EnergyModel,
@@ -151,7 +151,7 @@ def get_sezm_model(data: dict) -> EnergyModel:
                 "scheme are not supported in the pt_expt backend; use spin "
                 "scheme 'native' instead."
             )
-        return _get_dpa4_native_spin_model(data)
+        return get_native_spin_model(data)
     if str(data.get("bridging_method", "none")).lower() != "none":
         raise NotImplementedError(
             "`bridging_method` is not supported for DPA4/SeZM in the pt_expt backend."
@@ -213,17 +213,20 @@ def get_sezm_model(data: dict) -> EnergyModel:
     )
 
 
-def _get_dpa4_native_spin_model(data: dict) -> DPA4NativeSpinModel:
-    """Build a pt_expt DPA4/SeZM native (virtual-atom-free) spin model.
+def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
+    """Build a pt_expt native (virtual-atom-free) spin model.
 
-    Mirrors :func:`deepmd.dpmodel.model.model.get_dpa4_native_spin_model`:
-    no virtual atoms or doubled type map are introduced, ``use_spin`` is
+    Mirrors :func:`deepmd.dpmodel.model.model.get_native_spin_model`: no
+    virtual atoms or doubled type map are introduced, and ``use_spin`` is
     injected into the descriptor config (consumed by the descriptor's
-    equivariant spin embedding), and the non-spin backbone is built by
-    recursing into :func:`get_sezm_model` (now with ``"spin"`` popped),
-    which supplies the bridging/lora/compile/preset_out_bias rejections,
-    descriptor/fitting defaulting and ``exclude_types`` consistency check
-    shared with the non-spin DPA4/SeZM path.
+    equivariant spin embedding). The non-spin backbone is built by the
+    standard builder for the config's model type -- :func:`get_sezm_model`
+    for the DPA4/SeZM family (keeping its bridging/lora/compile/
+    preset_out_bias rejections and ``exclude_types`` consistency check),
+    else :func:`get_standard_model` -- then re-classed through the
+    registered :class:`NativeSpinEnergyModel`. Eligibility is the
+    ``descriptor.supports_native_spin()`` capability, not a descriptor-type
+    list.
 
     Parameters
     ----------
@@ -234,10 +237,6 @@ def _get_dpa4_native_spin_model(data: dict) -> DPA4NativeSpinModel:
     data = copy.deepcopy(data)
     spin_cfg = data.pop("spin")
     data.setdefault("descriptor", {})
-    if data["descriptor"].get("type", "dpa4") not in ("dpa4", "DPA4", "sezm", "SeZM"):
-        raise NotImplementedError(
-            "spin scheme 'native' requires the DPA4/SeZM descriptor"
-        )
     if data["descriptor"].get("add_chg_spin_ebd", False):
         raise NotImplementedError(
             "charge-spin FiLM combined with native spin on the graph route "
@@ -252,8 +251,28 @@ def _get_dpa4_native_spin_model(data: dict) -> DPA4NativeSpinModel:
         allow_missing_label=spin_cfg.get("allow_missing_label", False),
     )
     data["descriptor"]["use_spin"] = use_spin
-    backbone_model = get_sezm_model(data)
-    return DPA4NativeSpinModel(backbone_model=backbone_model, spin=spin)
+    model_type = str(data.get("type", "standard")).lower()
+    backbone_builder = (
+        get_sezm_model if model_type in ("dpa4", "sezm") else get_standard_model
+    )
+    try:
+        backbone_model = backbone_builder(data)
+    except TypeError as err:
+        # A descriptor without native spin support rejects the injected
+        # ``use_spin`` keyword at construction; translate to the
+        # capability-gate error.
+        raise NotImplementedError(
+            "spin scheme 'native' requires a descriptor with native spin "
+            "support (supports_native_spin()); descriptor type "
+            f"{data['descriptor'].get('type')!r} does not accept `use_spin`"
+        ) from err
+    descriptor = backbone_model.atomic_model.descriptor
+    if not descriptor.supports_native_spin():
+        raise NotImplementedError(
+            "spin scheme 'native' requires a descriptor declaring "
+            "supports_native_spin()"
+        )
+    return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin)
 
 
 def get_linear_model(model_params: dict) -> BaseModel:
@@ -355,6 +374,11 @@ def get_model(data: dict) -> BaseModel:
     model_type = data.get("type", "standard")
     if model_type == "standard":
         if "spin" in data:
+            if str(data["spin"].get("scheme", "deepspin")) == "native":
+                # Descriptor-agnostic entry: any standard-typed config whose
+                # descriptor declares supports_native_spin() rides the
+                # native scheme with zero model/dispatch changes.
+                return get_native_spin_model(data)
             return get_spin_model(data)
         return get_standard_model(data)
     elif model_type == "linear_ener":

…/pt_expt/model/dpa4_native_spin_model.py deepmd/pt_expt/model/native_spin_model.pydeepmd/pt_expt/model/dpa4_native_spin_model.py renamed to deepmd/pt_expt/model/native_spin_model.py deepmd/pt_expt/model/dpa4_native_spin_model.py renamed to deepmd/pt_expt/model/native_spin_model.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""pt_expt DPA4/SeZM native-spin model (NeighborGraph route, autograd force_mag)."""
+"""pt_expt native-spin energy model (NeighborGraph route, autograd force_mag)."""
 
 from typing import (
     Any,
@@ -10,11 +10,14 @@
     make_fx,
 )
 
-from deepmd.dpmodel.model.dpa4_native_spin_model import (
-    DPA4NativeSpinModel as DPA4NativeSpinModelDP,
+from deepmd.dpmodel.model.native_spin_model import (
+    make_native_spin_model,
 )
-from deepmd.pt_expt.common import (
-    torch_module,
+from deepmd.pt_expt.model.ener_model import (
+    EnergyModel,
+)
+from deepmd.pt_expt.model.model import (
+    BaseModel,
 )
 
 
@@ -59,86 +62,27 @@ def _translate_spin_energy_keys(
     return out
 
 
-@torch_module
-class DPA4NativeSpinModel(DPA4NativeSpinModelDP):
-    """pt_expt native-spin DPA4/SeZM model.
+@BaseModel.register("native_spin")
+class NativeSpinEnergyModel(make_native_spin_model(EnergyModel)):
+    """pt_expt native-spin energy model.
 
-    Mirrors :class:`deepmd.dpmodel.model.dpa4_native_spin_model.DPA4NativeSpinModel`
-    (construction, delegation, output defs are all inherited unchanged), but
-    overrides two methods:
+    ``make_native_spin_model`` applied to THIS backend's
+    :class:`~deepmd.pt_expt.model.ener_model.EnergyModel` (construction,
+    output defs, serialization all come from the factory; the deserialize
+    closure rebuilds through the pt_expt model class, so a registry round
+    trip yields a real ``torch.nn.Module``), plus two torch-specific
+    overrides:
 
-    - :meth:`forward`: the pt_expt backbone's ``call_common`` (Task 3 of the
-      "DPA4 native spin on the NeighborGraph route" plan) produces REAL
-      autograd ``energy_derv_r``/``energy_derv_r_mag``/``energy_derv_c_redu``
-      tensors, unlike the dpmodel parent's energy-only ``call`` (which is
+    - :meth:`forward`: the pt_expt ``call_common`` produces REAL autograd
+      ``energy_derv_r``/``energy_derv_r_mag``/``energy_derv_c_redu``
+      tensors, unlike the dpmodel factory's energy-only ``call`` (which is
       restricted to ``force``/``force_mag``/``virial`` as ``None``
       placeholders because dpmodel has no autograd).
-    - :meth:`deserialize`: the dpmodel parent's version hardcodes the
-      DPMODEL (numpy) ``BaseModel`` to rebuild ``backbone_model``, which
-      would produce a backbone missing the pt_expt/torch export machinery
-      (see the override's docstring); this class rebuilds it through the
-      pt_expt registry instead.
+    - :meth:`forward_lower_graph_exportable`: the graph-spin ``.pt2``
+      positional ABI (``spin`` at index 10) over the inherited
+      ``forward_common_lower_graph_exportable``.
     """
 
-    def __getattr__(self, name: str) -> Any:
-        """Get attribute from the wrapped model.
-
-        In torch.nn.Module, submodules (e.g. ``backbone_model``) are stored
-        in ``_modules``, not ``__dict__``. The dpmodel parent's
-        ``__getattr__`` guards its ``backbone_model`` delegation with
-        ``"backbone_model" not in self.__dict__``, which is always true here
-        and would incorrectly raise ``AttributeError`` for every submodule
-        access (including ``self.backbone_model`` itself). Mirrors
-        :class:`deepmd.pt_expt.model.spin_model.SpinModel`'s override: try
-        ``torch.nn.Module``'s own ``__getattr__`` (checks ``_parameters``,
-        ``_buffers``, ``_modules``) first, then fall back to
-        ``backbone_model`` delegation for arbitrary attributes.
-        """
-        try:
-            return torch.nn.Module.__getattr__(self, name)
-        except AttributeError:
-            pass
-        # backbone_model is in _modules, access via _modules directly to
-        # avoid re-entering __getattr__.
-        modules = self.__dict__.get("_modules", {})
-        backbone = modules.get("backbone_model")
-        if backbone is not None:
-            return getattr(backbone, name)
-        raise AttributeError(name)
-
-    @classmethod
-    def deserialize(cls, data: dict) -> "DPA4NativeSpinModel":
-        """Rebuild ``backbone_model`` through the pt_expt registry.
-
-        The dpmodel parent's ``deserialize`` (inherited otherwise) imports
-        ``deepmd.dpmodel.model.base_model.BaseModel`` at ITS OWN module
-        level, hardcoded regardless of which subclass's ``deserialize`` is
-        actually invoked -- so calling it on this pt_expt subclass would
-        still rebuild a plain numpy dpmodel ``backbone_model`` (missing
-        every pt_expt/torch export method, e.g.
-        ``forward_common_lower_graph_exportable``) and then let the
-        ``@torch_module`` auto-wrap machinery paper over it with a
-        dynamically-generated wrapper that only covers the dpmodel object's
-        OWN methods -- which never included the pt_expt-only export
-        machinery to begin with. Mirrors
-        :meth:`~deepmd.pt_expt.model.spin_model.SpinModel.deserialize`'s
-        override for exactly the same reason.
-        """
-        from deepmd.pt_expt.model.model import (
-            BaseModel,
-        )
-        from deepmd.utils.spin import (
-            Spin,
-        )
-
-        data = data.copy()
-        data.pop("@class", None)
-        data.pop("@version", None)
-        data.pop("type", None)
-        spin = Spin.deserialize(data.pop("spin"))
-        backbone_model = BaseModel.deserialize(data.pop("backbone_model"))
-        return cls(backbone_model=backbone_model, spin=spin)
-
     def forward(
         self,
         coord: torch.Tensor,
@@ -188,7 +132,7 @@ def forward(
         # ``spin=`` rides the NeighborGraph lower only; ``neighbor_graph_method``
         # is left at its default (None) so pt_expt's own default-flip resolves
         # it to the carry-all graph builder for this (DPA4) descriptor.
-        model_ret = self.backbone_model.call_common(
+        model_ret = self.call_common(
             coord,
             atype,
             box=box,
@@ -253,8 +197,9 @@ def forward_lower_graph_exportable(
         Two-layer make_fx trace, mirroring
         :meth:`~deepmd.pt_expt.model.ener_model.EnergyModel.forward_lower_graph_exportable`:
         the inner layer
-        (:meth:`~deepmd.pt_expt.model.make_model.make_model.forward_common_lower_graph_exportable`
-        on ``self.backbone_model``) traces ``forward_common_lower_graph``
+        (the inherited
+        :meth:`~deepmd.pt_expt.model.make_model.make_model.forward_common_lower_graph_exportable`)
+        traces ``forward_common_lower_graph``
         with ``spin`` as a SECOND autograd leaf next to ``edge_vec`` (fixing
         ``charge_spin=None`` -- this wrapper's ABI never exposes that slot);
         this outer layer re-traces with the PUBLIC positional ABI above and
@@ -312,7 +257,7 @@ def forward_lower_graph_exportable(
             ``atom_energy``, ``energy``, ``force``, ``force_mag``,
             ``virial``, and (when ``do_atomic_virial``) ``atom_virial``.
         """
-        traced = self.backbone_model.forward_common_lower_graph_exportable(
+        traced = self.forward_common_lower_graph_exportable(
             atype,
             n_node,
             n_local,