fix(dpa4): align pt_expt training and native-spin fine-tuning

Author: OutisLi

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -863,6 +863,10 @@ def _get_forward_wrapper_func(self) -> Callable[..., dict[str, np.ndarray]]:
         :meth:`get_sel` otherwise. Sizing a dense list from ``get_sel`` is not
         merely wasteful for a graph-native model -- such a model reports no
         finite capacity, so the allocation is unbounded.
+
+        A native-spin model conditions on a per-atom magnetic moment, which the
+        wrapper forwards on the graph route alone: that scheme implements only
+        the graph lower, so the dense route never carries a moment.
         """
         import array_api_compat
 
@@ -880,6 +884,7 @@ def model_forward(
             fparam: np.ndarray | None = None,
             aparam: np.ndarray | None = None,
             charge_spin: np.ndarray | None = None,
+            spin: np.ndarray | None = None,
         ) -> dict[str, np.ndarray]:
             # Get reference array to determine the target array type and device
             # Use out_bias as reference since it's always present
@@ -901,6 +906,8 @@ def model_forward(
                 aparam = xp.asarray(aparam, device=device)
             if charge_spin is not None:
                 charge_spin = xp.asarray(charge_spin, device=device)
+            if spin is not None:
+                spin = xp.asarray(spin, device=device)
 
             if self.uses_graph_lower():
                 nframes, nloc = atype.shape
@@ -927,6 +934,7 @@ def model_forward(
                         else None
                     ),
                     charge_spin=charge_spin,
+                    spin=None if spin is None else xp.reshape(spin, (-1, 3)),
                 )
                 # The graph route works on a flat node axis; restore the
                 # per-frame layout the dense route returns.

deepmd/dpmodel/descriptor/dpa4.py

@@ -595,7 +595,7 @@ class DescrptDPA4(NativeOP, BaseDescriptor):
     """
 
     _ENV_DIM: int = 1  # Use se_r style (radial only) for EnvMatStatSe compatibility
-    LATEST_VERSION: float = 1.1
+    LATEST_VERSION: float = 1.2
 
     def __init__(
         self,
@@ -2687,6 +2687,68 @@ def load(module: Any, prefix: str) -> Any:
         # === Output FFN ===
         self.output_ffn._load_variables(take_prefix("output_ffn."))
 
+    def _migrate_variables(
+        self,
+        variables: dict[str, Any],
+        version: float,
+        prefix: str = "",
+    ) -> float:
+        """Rewrite stored variables whose meaning changed since ``version``.
+
+        Operates on the flat mapping keyed by ``state_dict`` names, BEFORE
+        anything is assigned to a module: ``load_state_dict`` restores a
+        module's own buffers before descending into its children, so a
+        migration applied to live attributes would rewrite values the child
+        load is about to overwrite. Only representations are upgraded here;
+        a difference no rewrite can absorb stays a forward-time branch on
+        :attr:`version`, so a migrated descriptor never changes its own math.
+
+        Version 1.2 moved the env-seed spin gate from the spin coordinate to
+        the resulting environment quadratic form. For an active-spin model,
+        squaring the stored amplitude preserves the represented function.
+        Legacy native-spin models with no magnetic types instead carry
+        dormant, unconstrained spin-route values; those output-controlling
+        values are canonicalized to the zero function before the routes can
+        be activated by fine-tuning. Versions below 1.1 predate the
+        native-spin route and retain their original forward semantics.
+
+        Parameters
+        ----------
+        variables
+            Stored variables keyed by ``state_dict`` name, mutated in place.
+        version
+            Version the variables were written at.
+        prefix
+            Key prefix of this descriptor within ``variables``.
+
+        Returns
+        -------
+        float
+            Version the variables express after migration.
+        """
+        if not 1.1 <= version < 1.2:
+            return version
+
+        gate_key = prefix + "env_seed_embedding.spin_scale"
+        if self.use_spin is not None and not any(self.use_spin):
+            # dpmodel serialization names NativeLayer weights ``matrix``;
+            # pt_expt state dictionaries expose the wrapped attribute as ``w``.
+            dormant_keys = (
+                "spin_embedding.mag_layer2.matrix",
+                "spin_embedding.mag_layer2.w",
+                "spin_embedding.adam_spin_vec_weight",
+                "spin_embedding.adam_spin_nbr_weight",
+                "env_seed_embedding.spin_scale",
+            )
+            for name in dormant_keys:
+                key = prefix + name
+                if key in variables:
+                    xp = array_api_compat.array_namespace(variables[key])
+                    variables[key] = xp.zeros_like(variables[key])
+        elif gate_key in variables:
+            variables[gate_key] = variables[gate_key] ** 2
+        return 1.2
+
     def serialize(self) -> dict[str, Any]:
         return {
             "@class": "Descriptor",
@@ -2776,7 +2838,7 @@ def deserialize(cls, data: dict[str, Any]) -> DescrptDPA4:
         data.pop("env_mat", None)
         config.pop("s2_grid_resolution", None)
         obj = cls(**config)
-        obj.version = version
+        obj.version = obj._migrate_variables(variables, version)
         obj._load_variables(variables)
         return obj
 

deepmd/dpmodel/descriptor/dpa4_nn/embedding.py

@@ -445,6 +445,8 @@ class EnvironmentInitialEmbedding(NativeOP):
         Random seed for reproducibility.
     """
 
+    CONFIG_DERIVED_ARRAYS = ("spin_mask",)
+
     def __init__(
         self,
         *,
@@ -488,7 +490,10 @@ def __init__(
         # plus, for the native spin scheme, the 3 envelope-gated neighbor-spin
         # components, so the inner product ``D = M^T M`` yields the neighbor
         # spin-spin invariants alongside the geometric ones.
-        self.coord_dim = 4 + (3 if self.spin_flags is not None else 0)
+        self.geometry_coord_dim = 4
+        self.coord_dim = self.geometry_coord_dim + (
+            3 if self.spin_flags is not None else 0
+        )
 
         # === RBF projection: n_radial -> rbf_out_dim (two-layer MLP) ===
         # rbf_out_dim = max(32, embed_dim - 2*type_dim) to align G-network width to embed_dim
@@ -566,19 +571,18 @@ def __init__(
             dtype=PRECISION_DICT[self.precision.lower()],
         )
 
-        # === Native spin: per-type mask and isotropic channel scale ===
+        # === Native spin: per-type mask and post-quadratic activation gate ===
         # The mask gates the neighbor-spin channel by source type, so a
         # non-magnetic neighbor contributes zero and (critically) carries zero
-        # magnetic force ``-dE/ds``. The single scalar scale (shared across
-        # x/y/z) keeps the spin coordinates transforming with the geometry, so
-        # the env-matrix invariant stays SO(3)-invariant; ``output_proj`` is
-        # zero-initialized, so the spin contribution starts neutral regardless.
+        # magnetic force ``-dE/ds``. ``spin_scale`` multiplies the spin-only
+        # contribution after the environment quadratic form, providing a
+        # linear gate that can start from exactly zero.
         if self.spin_flags is not None:
             self.spin_mask = np.array(
                 [1.0 if flag else 0.0 for flag in self.spin_flags],
                 dtype=PRECISION_DICT[self.precision.lower()],
             )
-            self.spin_scale = np.ones(
+            self.spin_scale = np.zeros(
                 (1,), dtype=PRECISION_DICT[self.precision.lower()]
             )
 
@@ -648,11 +652,7 @@ def call(
                     xp.take(xp.astype(atype_flat, xp.int64), src_i, axis=0),
                     axis=0,
                 )[:, None]  # (E, 1)
-                spin_scale = xp.astype(
-                    xp_asarray_nodetach(xp, self.spin_scale[...], device=device),
-                    r_tilde.dtype,
-                )
-                spin_chan = edge_env * spin_scale * spin_src * mask  # (E, 3)
+                spin_chan = edge_env * spin_src * mask  # (E, 3)
             else:
                 spin_chan = xp.zeros(
                     (r_tilde.shape[0], 3), dtype=r_tilde.dtype, device=device
@@ -720,9 +720,26 @@ def call(
         # Summing over the coordinate axis makes D invariant to a joint rotation
         # of the geometry and the spin channels; with the spin channels present,
         # D additionally carries the neighbor spin-spin invariants.
-        env_agg_t = xp.permute_dims(env_agg, (0, 2, 1))  # (N, embed_dim, coord_dim)
-        env_agg_axis = env_agg[:, :, : self.axis_dim]  # (N, coord_dim, axis_dim)
-        D = xp.matmul(env_agg_t, env_agg_axis)  # (N, embed_dim, axis_dim)
+        if self.spin_flags is None:
+            env_agg_t = xp.permute_dims(env_agg, (0, 2, 1))
+            env_agg_axis = env_agg[:, :, : self.axis_dim]
+            D = xp.matmul(env_agg_t, env_agg_axis)
+        else:
+            geometry_agg = env_agg[:, : self.geometry_coord_dim, :]
+            spin_agg = env_agg[:, self.geometry_coord_dim :, :]
+            D_geometry = xp.matmul(
+                xp.permute_dims(geometry_agg, (0, 2, 1)),
+                geometry_agg[:, :, : self.axis_dim],
+            )
+            D_spin = xp.matmul(
+                xp.permute_dims(spin_agg, (0, 2, 1)),
+                spin_agg[:, :, : self.axis_dim],
+            )
+            spin_scale = xp.astype(
+                xp_asarray_nodetach(xp, self.spin_scale[...], device=device),
+                D_spin.dtype,
+            )
+            D = D_geometry + spin_scale * D_spin
 
         # === Step 6. Output projection for FiLM logits ===
         D_flat = xp.reshape(
@@ -994,6 +1011,8 @@ class SpinEmbedding(NativeOP):
         Whether parameters are trainable.
     """
 
+    CONFIG_DERIVED_ARRAYS = ("spin_mask",)
+
     def __init__(
         self,
         *,
@@ -1020,8 +1039,9 @@ def __init__(
         self.spin_flags = [bool(flag) for flag in use_spin]
 
         # === Per-type spin gate ===
-        # Non-persistent: rebuilt from config on construction and moved with the
-        # module, so the deterministic mask never enters the serialized state.
+        # Configuration-derived (hence ``CONFIG_DERIVED_ARRAYS``): rebuilt on
+        # construction and moved with the module, so the deterministic mask
+        # never enters the serialized state.
         self.spin_mask = np.array(
             [1.0 if bool(flag) else 0.0 for flag in use_spin], dtype=prec
         )
@@ -1053,23 +1073,26 @@ def __init__(
             seed=child_seed(seed_scalar, 1),
             trainable=self.trainable,
         )
+        self.mag_layer2.w = np.zeros(
+            (self.channels, self.channels),
+            dtype=prec,
+        )
 
         # === l=1 per-type per-channel weight ===
         # ``adam_`` prefix routes the table to Adam in HybridMuon, matching the
         # type-embedding treatment for per-type lookup parameters.
-        init_std = 1.0 / math.sqrt(float(self.ntypes + self.channels))
-        rng_vec = np.random.default_rng(child_seed(seed, 1))
-        self.adam_spin_vec_weight = rng_vec.normal(
-            0.0, init_std, size=(self.ntypes, self.channels)
-        ).astype(prec)
+        self.adam_spin_vec_weight = np.zeros(
+            (self.ntypes, self.channels),
+            dtype=prec,
+        )
 
         # === l=1 per-source-type per-channel weight for neighbor aggregation ===
         # Separate from the on-site weight: this scales the neighbor's spin
         # direction before it is aggregated into the center node's l=1 seed.
-        rng_nbr = np.random.default_rng(child_seed(seed, 2))
-        self.adam_spin_nbr_weight = rng_nbr.normal(
-            0.0, init_std, size=(self.ntypes, self.channels)
-        ).astype(prec)
+        self.adam_spin_nbr_weight = np.zeros(
+            (self.ntypes, self.channels),
+            dtype=prec,
+        )
 
     def call(self, spin: Any, atype: Any) -> tuple[Any, Any]:
         """

deepmd/dpmodel/descriptor/dpa4_nn/norm.py

@@ -360,18 +360,25 @@ def __init__(
         self.trainable = bool(trainable)
         prec = PRECISION_DICT[self.precision.lower()]
 
-        self.degree_index_m = np.asarray(degree_index_m, dtype=np.int64)
+        # A backend wrapping this module holds array attributes as framework
+        # tensors, possibly on an accelerator, and the caller's index table is
+        # one of them. Normalize it to NumPy once, then drive the setup below
+        # from that local binding rather than from the stored attribute: the
+        # numpy-only surface it relies on (``.size``, boolean-mask assignment)
+        # does not survive the backend's conversion either.
+        degree_index_m = to_numpy_array(degree_index_m).astype(np.int64, copy=False)
+        self.degree_index_m = degree_index_m
 
         # Pre-fuse degree balancing and channel averaging into a single weight:
         #   w_d = 1 / (n_coeff_l * (lmax+1) * C)
         # where n_coeff_l is the number of retained coefficients for degree l in
         # the reduced layout.
-        weights = np.zeros(self.degree_index_m.size, dtype=prec)
+        weights = np.zeros(degree_index_m.size, dtype=prec)
         scale = 1.0 / ((self.lmax + 1) * self.channels)
         for l in range(self.lmax + 1):
             n_coeff_l = 2 * min(l, self.mmax) + 1
             w_l = scale / float(n_coeff_l)
-            weights[self.degree_index_m == l] = w_l
+            weights[degree_index_m == l] = w_l
         if np.any(weights == 0):
             raise ValueError(
                 "ReducedEquivariantRMSNorm: balance_weight has zeros; "
@@ -411,15 +418,15 @@ def call(self, x: Any) -> Any:
         # === Step 2. Compute a shared degree-balanced RMS ===
         balance_weight = xp_asarray_nodetach(xp, self.balance_weight, device=device)
         mean_variance = xp.sum(x0 * x0, axis=(2, 3)) * balance_weight[0]
-        if self.degree_index_m.size > 1:
+        if xt.shape[2] > 0:
             mean_variance = mean_variance + xp.sum(
                 (xt * xt) * balance_weight[1:][None, None, :, None], axis=(2, 3)
             )
         inv_rms = 1.0 / xp.sqrt(mean_variance + self.eps)
         inv_rms = inv_rms[:, :, None, None]  # (F, E, 1, 1)
 
         x0 = x0 * inv_rms
-        if self.degree_index_m.size > 1:
+        if xt.shape[2] > 0:
             xt = xt * inv_rms
 
         # === Step 3. Apply per-degree affine parameters ===
@@ -428,7 +435,7 @@ def call(self, x: Any) -> Any:
         expanded_scale = xp.take(adam_scale, degree_index_m, axis=1)
         expanded_scale = expanded_scale[:, None, ...]  # (F, 1, D_m_trunc, C)
         x0 = x0 * expanded_scale[:, :, :1, :]
-        if self.degree_index_m.size > 1:
+        if xt.shape[2] > 0:
             xt = xt * expanded_scale[:, :, 1:, :]
 
         # === Step 4. Add scalar bias and restore layout ===
@@ -438,7 +445,7 @@ def call(self, x: Any) -> Any:
         )  # (F, 1, 1, C)
         x0 = x0 + bias0
 
-        out = x0 if self.degree_index_m.size == 1 else xp.concat([x0, xt], axis=2)
+        out = x0 if xt.shape[2] == 0 else xp.concat([x0, xt], axis=2)
         out = xp.astype(out, in_dtype)
         return out
 

deepmd/dpmodel/model/model.py

@@ -156,7 +156,6 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
     use_spin = normalize_spin_use_spin(spin_cfg["use_spin"], data["type_map"])
     spin = Spin(
         use_spin=use_spin,
-        virtual_scale=spin_cfg.get("virtual_scale", 1.0),
         allow_missing_label=spin_cfg.get("allow_missing_label", False),
     )
     data.setdefault("descriptor", {})

deepmd/dpmodel/model/native_spin_model.py

@@ -74,11 +74,14 @@ def make_native_spin_model(T_Model: type) -> type:
     class NSM(T_Model, NativeSpinModelKind):
         """Native-spin variant of ``T_Model`` (see ``make_native_spin_model``)."""
 
+        CONFIG_DERIVED_ARRAYS = ("spin_mask",)
+
         def __init__(self, *args: Any, spin: Spin, **kwargs: Any) -> None:
             super().__init__(*args, **kwargs)
             self.spin = spin
             self.ntypes_real = self.spin.ntypes_real
-            # Per-real-type 0/1 spin gate.
+            # Per-real-type 0/1 spin gate, derived from ``use_spin`` and hence
+            # rebuilt here rather than adopted from a checkpoint.
             self.spin_mask = self.spin.get_spin_mask()
 
         @staticmethod

deepmd/dpmodel/model/spin_model.py

@@ -55,6 +55,8 @@ class SpinModel(NativeOP):
         \boldsymbol{\tau}_i = \mathbf{F}_i^{\mathrm{virtual}} \times \boldsymbol{\sigma}_i.
     """
 
+    CONFIG_DERIVED_ARRAYS = ("spin_mask", "virtual_scale_mask")
+
     def __init__(
         self,
         backbone_model: DPAtomicModel,
@@ -76,6 +78,8 @@ def __init__(
                 # concrete default).
                 descriptor.disable_graph_lower()
         self.ntypes_real = self.spin.ntypes_real
+        # Both per-type tables follow from ``use_spin`` and ``virtual_scale``,
+        # so they are rebuilt here rather than adopted from a checkpoint.
         self.virtual_scale_mask = self.spin.get_virtual_scale_mask()
         self.spin_mask = self.spin.get_spin_mask()
 

deepmd/dpmodel/utils/stat.py

@@ -208,9 +208,20 @@ def _compute_model_predict(
         fparam = to_numpy_array(system.get("fparam", None))
         aparam = to_numpy_array(system.get("aparam", None))
         charge_spin = to_numpy_array(system.get("charge_spin", None))
+        # A native-spin model conditions on the per-atom moment, so the bias it
+        # predicts here is only the bias it will predict during training if the
+        # moment travels with the sample. The virtual-atom scheme never reaches
+        # this key: it expands the moment into virtual atoms before sampling.
+        spin = to_numpy_array(system.get("spin", None))
 
         sample_predict = model_forward(
-            coord, atype, box, fparam=fparam, aparam=aparam, charge_spin=charge_spin
+            coord,
+            atype,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            charge_spin=charge_spin,
+            spin=spin,
         )
         for kk in keys:
             model_predict[kk].append(