diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 98e6e0cd3a..d5a3a3eabd 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/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. @@ -935,6 +943,11 @@ def model_forward( for kk, vv in atomic_ret.items() } else: + if spin is not None: + raise NotImplementedError( + "native-spin output-bias calibration requires the " + "NeighborGraph lower" + ) ( extended_coord, extended_atype, diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 801b4c34f2..fd21e7798a 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/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 diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py index c93e982473..82c55dd8c4 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py +++ b/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]: """ diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/norm.py b/deepmd/dpmodel/descriptor/dpa4_nn/norm.py index c6492bd561..984482ee8f 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/norm.py +++ b/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,7 +418,7 @@ 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) ) @@ -419,7 +426,7 @@ def call(self, x: Any) -> Any: 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 diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index a66f40e7f3..2d378525dc 100644 --- a/deepmd/dpmodel/model/model.py +++ b/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", {}) diff --git a/deepmd/dpmodel/model/native_spin_model.py b/deepmd/dpmodel/model/native_spin_model.py index be6701447b..cd90b763b6 100644 --- a/deepmd/dpmodel/model/native_spin_model.py +++ b/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 diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 8744cc7b9f..faf0a8dacf 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/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() diff --git a/deepmd/dpmodel/utils/stat.py b/deepmd/dpmodel/utils/stat.py index de01a10792..900e101b7d 100644 --- a/deepmd/dpmodel/utils/stat.py +++ b/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( diff --git a/deepmd/kernels/utils.py b/deepmd/kernels/utils.py index df2997b7d3..1a0d02bc04 100644 --- a/deepmd/kernels/utils.py +++ b/deepmd/kernels/utils.py @@ -165,9 +165,9 @@ def use_amp_infer() -> bool: """Return whether bf16 autocast is enabled for inference. The flag is controlled by the ``DP_AMP_INFER`` environment variable and is - read at module construction time. It only affects inference when the - descriptor's ``use_amp`` option is also enabled; training follows - ``use_amp`` regardless of this environment variable. + read at module construction time. It controls inference independently of + the descriptor's ``use_amp`` option; training follows ``use_amp`` regardless + of this environment variable. Returns ------- diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 084f75b14d..6e05ae884f 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -439,7 +439,7 @@ class DescrptSeZM(BaseDescriptor, nn.Module): """ _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, @@ -2203,28 +2203,23 @@ def _compute_mode_ctx(self, device: torch.device) -> Generator[None, None, None] Notes ----- - - When `use_amp=True`, enables torch.autocast with bfloat16 on CUDA - during training. Eval/inference enables the same autocast region only - when ``DP_AMP_INFER`` was truthy at construction time. - This can improve speed and reduce memory usage on GPUs with native - bfloat16 support. - Disable AMP on GPUs without native bfloat16 support to avoid runtime - errors or additional conversion overhead. - - Only affects autocast-eligible operations. - - Does nothing during inference (`self.training=False`) unless - ``DP_AMP_INFER`` is enabled, on non-CUDA devices, or when - `use_amp=False`. + Training follows ``use_amp`` and evaluation follows ``DP_AMP_INFER`` + (captured at construction as ``use_amp_infer``). The two are + independent: mixed precision at inference is a throughput choice that + must not require a model to have been trained with it, and a + checkpoint therefore never carries the training switch into a + deployment. + + Autocast reaches only eligible operations, and only on CUDA. Leave it + off on GPUs without native bfloat16 to avoid conversion overhead. Yields ------ None Runs the wrapped region under the configured AMP setting. """ - if ( - not self.use_amp - or device.type != "cuda" - or (not self.training and not self.use_amp_infer) - ): + enabled = self.use_amp if self.training else self.use_amp_infer + if not enabled or device.type != "cuda": yield return @@ -2589,8 +2584,6 @@ def deserialize(cls, data: dict[str, Any]) -> DescrptSeZM: data.pop("env_mat", None) config.pop("s2_grid_resolution", None) obj = cls(**config) - obj.version = version - obj.version_tensor.fill_(version) template = obj.state_dict() state = { key: safe_numpy_to_tensor( @@ -2598,6 +2591,7 @@ def deserialize(cls, data: dict[str, Any]) -> DescrptSeZM: ) for key, value in variables.items() } + state["version_tensor"] = obj.version_tensor.new_tensor(version) obj.load_state_dict(state) return obj @@ -2638,6 +2632,64 @@ def update_sel( local_jdata_cpy["sel"] = sel[0] return local_jdata_cpy, min_nbor_dist + def _migrate_variables( + self, + variables: dict[str, torch.Tensor], + version: float, + prefix: str = "", + ) -> float: + """Rewrite stored state whose meaning changed since ``version``. + + Operates on the incoming ``state_dict``, 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 state keyed by ``state_dict`` name, mutated in place. + version + Version the state was written at. + prefix + Key prefix of this descriptor within ``variables``. + + Returns + ------- + float + Version the state expresses 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): + dormant_keys = ( + "spin_embedding.mag_layer2.matrix", + "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: + variables[key] = torch.zeros_like(variables[key]) + elif gate_key in variables: + variables[gate_key] = variables[gate_key] ** 2 + return 1.2 + def _load_from_state_dict( self, state_dict: dict[str, torch.Tensor], @@ -2683,7 +2735,14 @@ def _load_from_state_dict( if version_key not in state_dict: state_dict[version_key] = self.version_tensor.new_tensor(1.0) - # === Step 3. Drop transient descriptor state rebuilt at construction === + # === Step 3. Bring the incoming state up to the current semantics === + state_dict[version_key] = self.version_tensor.new_tensor( + self._migrate_variables( + state_dict, float(state_dict[version_key].item()), prefix + ) + ) + + # === Step 4. Drop transient descriptor state rebuilt at construction === expected_keys = {prefix + key for key in self.state_dict().keys()} for full_key in list(state_dict.keys()): if full_key.startswith(prefix) and full_key not in expected_keys: diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py index c01c3060b5..b70357a943 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py +++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py @@ -422,7 +422,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 + ) self.register_buffer( "eps_sq_tensor", torch.tensor(self.eps * self.eps, dtype=self.dtype, device=self.device), @@ -496,13 +499,13 @@ def __init__( seed=seed_out, ) - # === 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. This preserves + # SO(3) invariance and provides a linear, learnable gate that can start + # from exactly zero during spin-free fine-tuning. if self.spin_flags is not None: spin_mask = torch.tensor( [1.0 if flag else 0.0 for flag in self.spin_flags], @@ -511,7 +514,7 @@ def __init__( ) self.register_buffer("spin_mask", spin_mask, persistent=False) self.spin_scale = nn.Parameter( - torch.ones(1, dtype=self.dtype, device=self.device), + torch.zeros(1, dtype=self.dtype, device=self.device), requires_grad=trainable, ) @@ -575,7 +578,7 @@ def forward( mask = self.spin_mask.index_select( 0, atype_flat.index_select(0, src) ).unsqueeze(-1) # (E, 1) - spin_chan = edge_env * self.spin_scale * spin_src * mask # (E, 3) + spin_chan = edge_env * spin_src * mask # (E, 3) else: spin_chan = r_tilde.new_zeros(r_tilde.shape[0], 3) r_tilde = torch.cat([r_tilde, spin_chan], dim=-1) # (E, coord_dim) @@ -618,9 +621,22 @@ def forward( # 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 = env_agg.permute(0, 2, 1) # (N, embed_dim, coord_dim) - env_agg_axis = env_agg[:, :, : self.axis_dim] # (N, coord_dim, axis_dim) - D = torch.bmm(env_agg_t, env_agg_axis) # (N, embed_dim, axis_dim) + if self.spin_flags is None: + env_agg_t = env_agg.permute(0, 2, 1) # (N, embed_dim, coord_dim) + env_agg_axis = env_agg[:, :, : self.axis_dim] + D = torch.bmm(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 = torch.bmm( + geometry_agg.permute(0, 2, 1), + geometry_agg[:, :, : self.axis_dim], + ) + D_spin = torch.bmm( + spin_agg.permute(0, 2, 1), + spin_agg[:, :, : self.axis_dim], + ) + D = D_geometry + self.spin_scale * D_spin # === Step 6. Output projection for FiLM logits === D_flat = D.reshape( @@ -865,6 +881,7 @@ def __init__( self.channels, bias=False, activation_function=None, + init="final", precision=self.precision, seed=child_seed(seed_scalar, 1), trainable=trainable, @@ -874,32 +891,19 @@ def __init__( # ``adam_`` prefix routes the table to Adam in HybridMuon, matching the # type-embedding treatment for per-type lookup parameters. self.adam_spin_vec_weight = nn.Parameter( - torch.empty( + torch.zeros( self.ntypes, self.channels, device=self.device, dtype=self.dtype ) ) - init_std = 1.0 / math.sqrt(float(self.ntypes + self.channels)) - nn.init.normal_( - self.adam_spin_vec_weight, - mean=0.0, - std=init_std, - generator=get_generator(child_seed(seed, 1)), - ) # === 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. self.adam_spin_nbr_weight = nn.Parameter( - torch.empty( + torch.zeros( self.ntypes, self.channels, device=self.device, dtype=self.dtype ) ) - nn.init.normal_( - self.adam_spin_nbr_weight, - mean=0.0, - std=init_std, - generator=get_generator(child_seed(seed, 2)), - ) for p in self.parameters(): p.requires_grad = trainable diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 8671a1e94e..887ce1754c 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -505,10 +505,9 @@ def _get_sezm_native_spin_model(model_params: dict) -> BaseModel: use_spin = [bool(flag) for flag in model_params["spin"]["use_spin"]] # ``virtual_scale`` is a virtual-atom geometric device; the native scheme - # only needs ``use_spin`` for masking, so default it when absent. + # only needs ``use_spin`` for masking, so it stays unset here. spin = Spin( use_spin=use_spin, - virtual_scale=model_params["spin"].get("virtual_scale", 1.0), allow_missing_label=model_params["spin"].get("allow_missing_label", False), ) diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py index 5676652ba3..5e0cc1f50e 100644 --- a/deepmd/pt_expt/common.py +++ b/deepmd/pt_expt/common.py @@ -5,7 +5,8 @@ classes (array_api_compat-based) as PyTorch modules. The key insight is to detect attributes by their **value type** rather than by hard-coded names: -- numpy arrays → torch buffers (persistent state like statistics, masks) +- numpy arrays → torch buffers (persistent, unless the owning class lists the + array in ``CONFIG_DERIVED_ARRAYS``) - dpmodel objects → pt_expt torch.nn.Module wrappers (via registry lookup) - None values → clear existing buffers @@ -222,8 +223,12 @@ def dpmodel_setattr(obj: torch.nn.Module, name: str, value: Any) -> tuple[bool, the need to hard-code attribute names in each wrapper's __setattr__ method. It handles three cases: - 1. **numpy arrays → torch buffers**: Persistent state like statistics (davg, dstd) - or masks that should be saved in state_dict and moved with .to(device). + 1. **numpy arrays → torch buffers**: State such as statistics (davg, dstd) that + is saved in state_dict and moved with .to(device). An array the owning + class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer + instead: being a pure function of the configuration it is rebuilt by + ``__init__``, so adopting a stored copy would let a checkpoint whose + configuration differs override the built value. 2. **None values → clear buffers**: Setting an existing buffer to None. 3. **dpmodel objects → pt_expt modules**: Nested dpmodel objects like AtomExcludeMaskDP or NetworkCollectionDP are converted to their pt_expt @@ -305,7 +310,11 @@ def dpmodel_setattr(obj: torch.nn.Module, name: str, value: Any) -> tuple[bool, # deserialize), remove it first so register_buffer doesn't conflict. if hasattr(obj, name) and name not in obj._buffers: delattr(obj, name) - obj.register_buffer(name, tensor) + obj.register_buffer( + name, + tensor, + persistent=name not in getattr(type(obj), "CONFIG_DERIVED_ARRAYS", ()), + ) return True, tensor # clear an existing buffer to None @@ -448,6 +457,19 @@ def __setattr__(self, name: str, value: Any) -> None: if not handled: super().__setattr__(name, value) + def _load_from_state_dict( + self, state_dict: dict, prefix: str, *args: Any, **kwargs: Any + ) -> None: + # A non-persistent buffer is configuration-derived: ``__init__`` + # already rebuilt it, so any archived copy is discarded rather than + # reported as an unexpected key. This keeps checkpoints written + # while the buffer was still persistent loadable, and keeps a + # checkpoint from a differently configured model (e.g. a spin-free + # pretraining, whose spin gate is all zero) from overriding it. + for name in self._non_persistent_buffers_set: + state_dict.pop(prefix + name, None) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + # Auto-generate forward -> call redirect if not explicitly defined if hasattr(module, "call") and "forward" not in module.__dict__: diff --git a/deepmd/pt_expt/descriptor/dpa4.py b/deepmd/pt_expt/descriptor/dpa4.py index de9e4e3824..c497c5e6a4 100644 --- a/deepmd/pt_expt/descriptor/dpa4.py +++ b/deepmd/pt_expt/descriptor/dpa4.py @@ -143,6 +143,11 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: # skips the missing buffer, so listing both concrete subclasses is safe) "S2GridNet": ("residual_scale",), "SO3GridNet": ("residual_scale",), + # dpa4_nn.grid_net frame mixing, built only by ``mode="cross"`` grid nets. + # Unlike the surrounding projections these are plain numpy arrays rather + # than NativeLayer objects, so they need an explicit entry here. + "FrameExpand": ("weight",), + "FrameContract": ("weight",), # descriptor-level FiLM strengths "DescrptDPA4": ("film_scale_strength_log", "film_shift_strength_log"), } @@ -205,6 +210,15 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: "graph_lower_disabled", torch.zeros((), dtype=torch.bool, device="cpu"), ) + # Persisted descriptor version, for the same reason: pt_expt rebuilds + # the module from config before loading, so without a buffer every + # checkpoint would come back claiming the semantics of the running + # code and silently skip ``_migrate_variables``. + torch.nn.Module.register_buffer( + self, + "version_tensor", + torch.tensor(self.version, dtype=torch.float64, device="cpu"), + ) self.use_amp_infer = use_amp_infer() _promote_trainable_tree(self) @@ -213,6 +227,9 @@ def deserialize(cls, data: dict) -> "DescrptDPA4": # deserialize assigns numpy arrays after __init__, which demotes # promoted Parameters back to buffers; re-promote at the end. obj = super().deserialize(data) + # The buffer carries the version of the restored variables, not the + # version the fresh construction started from. + obj.version_tensor.fill_(obj.version) return _promote_trainable_tree(obj) def _in_training_mode(self) -> bool: @@ -303,7 +320,20 @@ def _load_from_state_dict( # data-dependent ``bool(FakeTensor)`` guard that breaks # torch.export (GuardOnDataDependentSymNode Eq(u0, 1)). self._graph_lower_disabled = bool(state_dict[key]) + + # Back-compat: checkpoints predating the version buffer were written + # under version 1.1, the last one released before it existed. + version_key = prefix + "version_tensor" + if version_key not in state_dict: + state_dict[version_key] = self.version_tensor.new_tensor(1.1) + state_dict[version_key] = self.version_tensor.new_tensor( + self._migrate_variables( + state_dict, float(state_dict[version_key].item()), prefix + ) + ) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + self.version = float(self.version_tensor.item()) def forward(self, *args: Any, **kwargs: Any) -> Any: return self.call(*args, **kwargs) @@ -317,18 +347,17 @@ def _forward_blocks(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Any: geometry, edge cache, radial, env-seed, GIE and output FFN stages stay in fp32 (or higher). The dpmodel base stores ``use_amp`` only as a config flag and never autocasts (array-API has no autocast), so the - real automatic mixed precision lives here. ``x`` is the node-feature - tensor entering the blocks; its device equals the working device, so - autocast engages when ``self.use_amp`` is set, the inputs live on a - CUDA device, and either the module is training or eval-time AMP was - opted in through ``DP_AMP_INFER`` (captured once at construction as - ``self.use_amp_infer``). + real automatic mixed precision lives here. + + Training follows ``use_amp`` and evaluation follows ``DP_AMP_INFER`` + (captured once at construction as ``use_amp_infer``). The two are + independent: mixed precision at inference is a throughput choice that + must not require a model to have been trained with it. ``x`` is the + node-feature tensor entering the blocks, and its device is the working + device. """ - if ( - self.use_amp - and x.device.type == "cuda" - and (self.training or self.use_amp_infer) - ): + enabled = self.use_amp if self.training else self.use_amp_infer + if enabled and x.device.type == "cuda": with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): return super()._forward_blocks(x, *args, **kwargs) return super()._forward_blocks(x, *args, **kwargs) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 9d6396d990..fa3421528e 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -615,12 +615,13 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: model = get_model(deepcopy(model_params)).to(DEVICE) # Strip the `_CompiledModel` wrapper that pt_expt training applies - # after compilation (training.py:996). The saved state_dict has - # `model.Default.original_model.X` keys (the real weights) plus + # after compilation. The saved state_dict has + # `model.Default.original_model.X` keys (the real weights). Some + # checkpoints additionally carry # `model.Default.compiled_forward_lower._orig_mod._param_constant*` - # / `_tensor_constant*` keys (graph constants baked into the - # compiled forward — duplicates of the real weights, useless for - # eager inference). Drop the latter and unwrap the former. + # / `_tensor_constant*` keys (graph constants baked into a compiled + # forward — duplicates of the real weights, useless for eager + # inference). Drop the latter and unwrap the former. cleaned: dict[str, Any] = {} compiled_marker = ".compiled_forward_lower." # Per-task buffer copies registered on _CompiledModel (bias_atom_e, diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 50f60ecf49..f0d87a9118 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -57,7 +57,7 @@ log = logging.getLogger(__name__) -# Warn at most once per process for backend-ignored switches (keyed by name). +# Warn at most once per process for backend-relocated switches (keyed by name). _WARNED_ONCE: set[str] = set() @@ -85,22 +85,20 @@ def get_sezm_model(data: dict) -> BaseModel: :func:`get_native_spin_model`; the two combine. Still unsupported here, each raising ``NotImplementedError``: the - virtual-atom (``deepspin``) spin scheme, ``lora``, ``use_compile``, and + virtual-atom (``deepspin``) spin scheme, ``lora``, and ``preset_out_bias``. Notes ----- - ``enable_tf32`` is accepted but ignored: the pt backend uses it to toggle - TF32 matmul precision, while the pt_expt backend always runs at full - ("highest") matmul precision, which is numerically conservative. + ``model.use_compile`` and ``model.enable_tf32`` are accepted but ignored. + The pt backend hangs both switches off the model because only SeZM + implements the compile and precision paths there, whereas pt_expt applies + them to every model and so reads them from ``training.enable_compile`` and + ``training.enable_tf32``. Requesting the relocated compile switch warns + once; ``enable_tf32`` cannot, because normalization fills its model-level + default in on every run and an explicit request is indistinguishable. """ data = copy.deepcopy(data) - if bool(data.get("enable_tf32", True)) and "enable_tf32" not in _WARNED_ONCE: - log.warning( - "`enable_tf32` has no effect on the pt_expt backend, which " - "always runs at full ('highest') matmul precision; ignoring it." - ) - _WARNED_ONCE.add("enable_tf32") if "spin" in data: if str(data["spin"].get("scheme", "deepspin")) != "native": raise NotImplementedError( @@ -118,10 +116,12 @@ def get_sezm_model(data: dict) -> BaseModel: raise NotImplementedError( "`lora` is not supported for DPA4/SeZM in the pt_expt backend." ) - if data.get("use_compile"): - raise NotImplementedError( - "`use_compile` is not supported for DPA4/SeZM in the pt_expt backend." + if data.get("use_compile") and "use_compile" not in _WARNED_ONCE: + log.warning( + "`model.use_compile` has no effect on the pt_expt backend; " + "set `training.enable_compile` instead." ) + _WARNED_ONCE.add("use_compile") if data.get("preset_out_bias"): raise NotImplementedError( "`preset_out_bias` is not supported for DPA4/SeZM in the pt_expt backend." @@ -306,7 +306,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["descriptor"]["use_spin"] = use_spin diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f277452fcb..a489ec2a75 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -13,6 +13,9 @@ Callable, Mapping, ) +from contextlib import ( + nullcontext, +) from copy import ( deepcopy, ) @@ -107,7 +110,9 @@ clip_grad_norm_, ) from deepmd.pt_expt.train.utils import ( + MatmulPrecisionPolicy, count_parameters, + infer_compile_enabled, infer_env_defaults, resolve_best_checkpoint_dir, scoped_env_defaults, @@ -266,6 +271,11 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: After ``share_params``, the fitting net's child sub-modules are the same Python objects across tasks, so ``id(first_child)`` is equal for all shared tasks and unique across unrelated models. + + ``has_spin()`` leads the key because native-spin models trace an extra + per-node moment input and return additional magnetic outputs. A spin task + must therefore never reuse a spin-free task's compiled graph even when the + descriptor and fitting parameters are shared. """ descriptor_id: int = 0 try: @@ -278,13 +288,15 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: except AttributeError: pass + fitting_id: int = id(model) try: fitting = model.get_fitting_net() for _, child in fitting.named_children(): - return (descriptor_id, id(child)) + fitting_id = id(child) + break except AttributeError: pass - return (descriptor_id, id(model)) + return (int(model.has_spin()), descriptor_id, fitting_id) # --------------------------------------------------------------------------- @@ -484,14 +496,12 @@ def _trace_and_compile( make_fx, ) - was_training = model.training - # Trace in train mode so that create_graph=True is captured inside - # task_deriv_one. Without this, the autograd.grad that computes - # forces is traced with create_graph=False (eval mode), producing - # force tensors that are detached from model parameters — force loss - # backprop cannot reach the weights and force RMSE never decreases. - model.train() - + # The model's current mode decides what gets traced, and the caller keys + # the resulting graph by that same mode. It matters inside + # task_deriv_one: the autograd.grad that computes forces is decomposed + # with create_graph=self.training, so a train-mode trace keeps the + # double-backward the force loss needs to reach the weights, while an + # eval-mode trace omits it. task_buf_order: tuple[str, ...] = tuple(task_buffers.keys()) if task_buffers else () task_buf_vals_trace: tuple[torch.Tensor, ...] = ( tuple(task_buffers[k] for k in task_buf_order) if task_buffers else () @@ -626,15 +636,13 @@ def fn( ) return ( - _finalize_compiled_lower(traced_lower, model, was_training, compile_opts), + _finalize_compiled_lower(traced_lower, compile_opts), task_buf_order, ) def _finalize_compiled_lower( traced_lower: "torch.fx.GraphModule", - model: torch.nn.Module, - was_training: bool, compile_opts: dict[str, Any] | None, extra_options: dict[str, Any] | None = None, ) -> torch.nn.Module: @@ -654,9 +662,6 @@ def _finalize_compiled_lower( # left by erase_node(), which can cause segfaults during dynamo re-trace. traced_lower = _rebuild_graph_module(traced_lower) - if not was_training: - model.eval() - # This is the common boundary immediately before every pt_expt # ``torch.compile`` call. Applying the idempotent process-global patches # here leaves eager-only imports untouched while still preceding all @@ -683,6 +688,7 @@ def _trace_and_compile_graph( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, + spin: torch.Tensor | None, compile_opts: dict[str, Any] | None = None, task_buffers: dict[str, torch.Tensor] | None = None, ) -> tuple[torch.nn.Module, tuple[str, ...]]: @@ -698,13 +704,18 @@ def _trace_and_compile_graph( and returns those public keys on the FLAT node axis (``N == sum(n_node)``); the caller (:meth:`_CompiledModel.forward`) unravels them to ``(nf, nloc, *)``. + Native-spin models additionally carry the per-node moment as a second + autograd leaf. Their compiled lower returns ``force_mag`` and ``mask_mag`` + through the same model-owned translation used by eager execution. + Parameters ---------- model The (uncompiled) graph-eligible energy model. - fparam, aparam, charge_spin + fparam, aparam, charge_spin, spin Representative optional inputs (or ``None``) so the traced branch matches what :meth:`_CompiledModel.forward` passes at run time. + ``spin`` is the flat ``(N, 3)`` per-node moment. compile_opts User-supplied inductor options (merged over the built-in defaults). task_buffers @@ -722,11 +733,9 @@ def _trace_and_compile_graph( _translate_energy_keys, ) - was_training = model.training - # Trace in train mode so create_graph=True is captured inside the graph - # force backward (forward_common_lower_graph passes create_graph=self.training). - model.train() - + # Traced in the model's current mode, which the caller keys the graph by: + # forward_common_lower_graph passes create_graph=self.training to the force + # backward, so the mode decides whether the graph carries it. task_buf_order: tuple[str, ...] = tuple(task_buffers.keys()) if task_buffers else () task_buf_vals_trace: tuple[torch.Tensor, ...] = ( tuple(task_buffers[k] for k in task_buf_order) if task_buffers else () @@ -813,6 +822,7 @@ def _trace_and_compile_graph( want_fparam=fparam is not None, want_aparam=aparam is not None, want_charge_spin=charge_spin is not None, + want_spin=spin is not None, ) ( s_atype, @@ -825,10 +835,12 @@ def _trace_and_compile_graph( s_destination_row_ptr, s_source_order, s_source_row_ptr, - s_fparam, - s_aparam, - s_charge_spin, + *s_conditioning, ) = sample + # The synthetic native-spin ABI puts spin before the optional conditioning + # tensors, while forward_common_lower_graph takes it after them. + s_spin = s_conditioning.pop(0) if spin is not None else None + s_fparam, s_aparam, s_charge_spin = s_conditioning def fn( atype: torch.Tensor, @@ -844,6 +856,7 @@ def fn( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, + spin: torch.Tensor | None, *task_buf_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: # Patch task-specific buffers with the proxy tensors so make_fx records @@ -861,8 +874,8 @@ def fn( originals[name] = _fitting._buffers.get(name) _fitting._buffers[name] = val try: - # forward_common_lower_graph makes edge_vec the autograd leaf - # internally, so no outer detach/requires_grad_ here. + # forward_common_lower_graph makes edge_vec and, when present, + # spin the autograd leaves internally. model_ret = model.forward_common_lower_graph( atype, n_node, @@ -878,13 +891,20 @@ def fn( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + spin=spin, ) - return _translate_energy_keys( + if spin is None: + return _translate_energy_keys( + model_ret, + do_grad_r=do_grad_r, + do_grad_c=do_grad_c, + do_atomic_virial=False, + local=True, + ) + return model._translate_eager_call( model_ret, - do_grad_r=do_grad_r, - do_grad_c=do_grad_c, + atype, do_atomic_virial=False, - local=True, ) finally: for name, orig in originals.items(): @@ -917,6 +937,7 @@ def fn( s_fparam, s_aparam, s_charge_spin, + s_spin, *task_buf_vals_trace, ) @@ -928,8 +949,6 @@ def fn( return ( _finalize_compiled_lower( traced_lower, - model, - was_training, compile_opts, extra_options={"cpp.simdlen": 0}, ), @@ -944,6 +963,17 @@ class _CompiledModel(torch.nn.Module): ``forward()`` invocation using that batch's tensors, so no extra ``get_data()`` call is needed during ``__init__``. Tasks that share the same model structure reuse the compiled graph via ``compiled_by_structure``. + + Training always gets a compiled graph. Evaluation delegates to the + original eager model unless compiled inference is enabled, in which case it + gets a separate graph keyed by ``self.training``. The two compiled modes + differ in more than speed. The force ``autograd.grad`` is decomposed with + ``create_graph=self.training``, so a training graph carries the + double-backward needed by the force loss while an evaluation graph does + not, and train-only behaviour such as the random local-Z roll is baked in + at trace time. Serving evaluation from the training graph would therefore + both cost a second-order graph per validation batch and randomize its + frames. """ def __init__( @@ -955,14 +985,24 @@ def __init__( compile_opts: dict[str, Any] | None = None, compiled_by_structure: dict | None = None, task_key: str = DEFAULT_TASK_KEY, + compile_eval: bool = False, ) -> None: super().__init__() self.original_model = original_model - self.compiled_forward_lower: torch.nn.Module | None = None + # Compiled graphs are execution artifacts, not model state, and are + # deliberately held in a plain dict rather than as submodules: a + # registered graph copies its baked-in constants (``_param_constant*``, + # duplicates of the real weights) into every checkpoint, once per + # compiled mode. Keyed by the ``self.training`` flag they serve. + self._compiled_lower_by_mode: dict[bool, torch.nn.Module] = {} + # ``(attributes, start time)`` of a graph whose Inductor compile the + # next call will trigger; see :meth:`_report_pending_compile`. + self._pending_compile: tuple[str, float] | None = None self._task_buf_order = task_buf_order self._structure_key = structure_key self._task_key = task_key self._compile_opts = compile_opts + self._compile_eval = compile_eval # Stored only for the first-forward compile call; freed afterwards. self._task_buffers = task_buffers # Shared dict across all _CompiledModel instances in the same Trainer. @@ -982,8 +1022,8 @@ def _compiled_lower_for( ) -> tuple[torch.nn.Module, tuple[str, ...]]: """Return the compiled graph of this model, tracing it at most once. - Tasks that share a model structure share one compiled graph, so a task - reaching this point second only reports the reuse. + Tasks that share a model structure share one compiled graph per mode, + so a task reaching this point second only reports the reuse. Parameters ---------- @@ -998,21 +1038,43 @@ def _compiled_lower_for( tuple[torch.nn.Module, tuple[str, ...]] The compiled graph and its buffer order. """ - attributes = f"task={self._task_key}, path={path}" - cached = self._compiled_by_structure.get(self._structure_key) + mode = "train" if self.training else "eval" + attributes = f"task={self._task_key}, path={path}, mode={mode}" + cache_key = (*self._structure_key, self.training) + cached = self._compiled_by_structure.get(cache_key) if cached is not None: log.info("Reusing the graph compiled for an earlier task (%s).", attributes) return cached log.info("Tracing and compiling the model (%s).", attributes) started = time.perf_counter() compiled = trace() + # ``torch.compile`` only schedules the Inductor compile; it runs on the + # first call. Timing the trace alone would report a fraction of the + # real cost, so the elapsed time is reported by the caller once that + # first call has completed. + self._pending_compile = (attributes, started) + self._compiled_by_structure[cache_key] = compiled + return compiled + + def _report_pending_compile(self) -> None: + """Log the compile time of a graph whose first call just returned. + + CUDA work is synchronized first: the launch queue would otherwise + return before the device finished, and the measurement would omit the + very work it is meant to cover. + """ + pending = self._pending_compile + if pending is None: + return + self._pending_compile = None + attributes, started = pending + if torch.cuda.is_available(): + torch.cuda.synchronize() log.info( "Finished compiling (%s) in %.1f s.", attributes, time.perf_counter() - started, ) - self._compiled_by_structure[self._structure_key] = compiled - return compiled def __getattr__(self, name: str) -> Any: # Delegate unknown lookups to original_model so that callers such as @@ -1035,7 +1097,20 @@ def forward( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: + if not self.training and not self._compile_eval: + kwargs = { + "box": box, + "fparam": fparam, + "aparam": aparam, + "do_atomic_virial": do_atomic_virial, + "charge_spin": charge_spin, + } + if spin is not None: + kwargs["spin"] = spin + return self.original_model(coord, atype, **kwargs) + from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -1055,7 +1130,22 @@ def forward( self._graph_eligible = model_uses_graph_lower(self.original_model) if self._graph_eligible: return self._forward_graph( - coord, atype, box, fparam, aparam, charge_spin, nframes, nloc, rcut + coord, + atype, + box, + fparam, + aparam, + charge_spin, + spin, + nframes, + nloc, + rcut, + ) + + if spin is not None: + raise NotImplementedError( + "model-level spin requires the NeighborGraph lower; this model " + "compiles the dense neighbor-list lower" ) sel = self.original_model.get_sel() @@ -1143,7 +1233,8 @@ def forward( # batch's tensors (prime-padded inside _trace_and_compile). # Mirrors DPA4's on-cache-miss compile so no separate get_data() # is needed during __init__. - if self.compiled_forward_lower is None: + compiled_lower = self._compiled_lower_by_mode.get(self.training) + if compiled_lower is None: # Optional inputs (fparam / charge_spin) are normalized to their # defaults above, so their presence is now config-driven (a # function of the model's ``dim_*``) rather than data-driven. @@ -1165,9 +1256,8 @@ def forward( compile_opts=self._compile_opts, ), ) - self.compiled_forward_lower = compiled_lower + self._compiled_lower_by_mode[self.training] = compiled_lower self._task_buf_order = buf_order - self._task_buffers = None # free; no longer needed after compile ext_coord = ext_coord.detach().requires_grad_(True) @@ -1191,16 +1281,21 @@ def forward( ) from exc else: task_buf_vals = () - result = self.compiled_forward_lower( - ext_coord, - ext_atype, - nlist, - mapping, - fparam, - aparam, - charge_spin, - *task_buf_vals, - ) + # The evaluation graph carries no double-backward, so nothing outside + # it needs an autograd tape; forces come from the decomposed backward + # ops baked into the graph itself. + with nullcontext() if self.training else torch.no_grad(): + result = compiled_lower( + ext_coord, + ext_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + *task_buf_vals, + ) + self._report_pending_compile() # Translate forward_lower keys -> forward keys. OUTPUT-AGNOSTIC: # every key passes through unchanged (energy models emit @@ -1249,6 +1344,7 @@ def _forward_graph( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, + spin: torch.Tensor | None, nframes: int, nloc: int, rcut: float, @@ -1270,12 +1366,12 @@ def _forward_graph( coord_3d = coord.detach().reshape(nframes, nloc, 3) box_flat = box.detach().reshape(nframes, 9) if box is not None else None - # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) -- like - # every per-node tensor of the graph schema (the trace sample from - # build_synthetic_graph_inputs is flat too, so the compiled lower's - # input spec expects it). + # Graph-lower per-node inputs are flat: aparam is (N, nda) and spin is + # (N, 3), matching the synthetic trace schema. if aparam is not None: aparam = aparam.reshape(nframes * nloc, -1) + if spin is not None: + spin = spin.reshape(nframes * nloc, 3) # Mirror the optional-input defaulting of the dense path / eager # call_common: a model configured with fparam / charge_spin substitutes @@ -1324,8 +1420,9 @@ def _forward_graph( ) atype_flat = atype.reshape(nframes * nloc) - # Lazy compile of the GRAPH lower (cached per structure key). - if self.compiled_forward_lower is None: + # Lazy compile of the GRAPH lower (cached per structure key and mode). + compiled_lower = self._compiled_lower_by_mode.get(self.training) + if compiled_lower is None: compiled_lower, buf_order = self._compiled_lower_for( "neighbor-graph", lambda: _trace_and_compile_graph( @@ -1333,18 +1430,20 @@ def _forward_graph( fparam, aparam, charge_spin, + spin, task_buffers=self._task_buffers, compile_opts=self._compile_opts, ), ) - self.compiled_forward_lower = compiled_lower + self._compiled_lower_by_mode[self.training] = compiled_lower self._task_buf_order = buf_order - self._task_buffers = None # Feed a detached, grad-enabled edge_vec leaf: the traced graph's internal # ``edge_vec.detach()`` is stripped by ``_strip_saved_tensor_detach`` (as # for the dense ext_coord leaf), so the force backward roots at this input. edge_vec = ng.edge_vec.detach().requires_grad_(True) + if spin is not None: + spin = spin.detach().requires_grad_(True) if self._task_buf_order: try: @@ -1367,22 +1466,26 @@ def _forward_graph( else: task_buf_vals = () - result = self.compiled_forward_lower( - atype_flat, - ng.n_node, - ng.n_node, - ng.edge_index, - edge_vec, - ng.edge_mask, - ng.destination_order, - ng.destination_row_ptr, - ng.source_order, - ng.source_row_ptr, - fparam, - aparam, - charge_spin, - *task_buf_vals, - ) + # See the dense path: an evaluation graph needs no surrounding tape. + with nullcontext() if self.training else torch.no_grad(): + result = compiled_lower( + atype_flat, + ng.n_node, + ng.n_node, + ng.edge_index, + edge_vec, + ng.edge_mask, + ng.destination_order, + ng.destination_row_ptr, + ng.source_order, + ng.source_row_ptr, + fparam, + aparam, + charge_spin, + spin, + *task_buf_vals, + ) + self._report_pending_compile() # The compiled graph lower emits PUBLIC keys on the FLAT node axis # (``atom_energy`` / ``force`` are (N, *); ``energy`` / ``virial`` are @@ -1395,7 +1498,14 @@ def _forward_graph( # shape heuristic keeps the single-atom case (nloc == 1, where # N == nframes) correct -- node-level outputs still reshape to # (nf, 1, *) instead of staying (nf, *). - node_level_keys = {"atom_energy", "force", "atom_virial", "mask"} + node_level_keys = { + "atom_energy", + "force", + "force_mag", + "atom_virial", + "mask", + "mask_mag", + } out: dict[str, torch.Tensor] = {} for key, val in result.items(): if ( @@ -1517,6 +1627,7 @@ def __init__( training_params.get("change_bias_after_training", False) ) self.enable_compile = bool(training_params.get("enable_compile", False)) + self.enable_tf32 = bool(training_params.get("enable_tf32", False)) self._raise_if_sharding_unsupported() # Model --------------------------------------------------------------- @@ -1526,8 +1637,12 @@ def __init__( ) # Descriptors sample the eval-time policy variables exactly once, while # they are being constructed; keep the config-derived defaults scoped to - # construction so they do not leak into the rest of the process. + # construction so they do not leak into the rest of the process. The + # matmul precision policy reads the same variables and so is resolved + # in the same scope. with scoped_env_defaults(infer_env_defaults(validating_params)): + self.compile_infer = infer_compile_enabled() + self.matmul_precision = MatmulPrecisionPolicy(self.enable_tf32) for model_key in self.model_keys: self.models[model_key] = get_model( deepcopy(self.model_params_by_task[model_key]) @@ -1668,7 +1783,12 @@ def initialize_statistics( self.nonfinite_grad_guard = NonFiniteGradGuard() # Model wrapper ------------------------------------------------------- - self.wrapper = ModelWrapper(self.model, self.loss, model_params=model_params) + self.wrapper = ModelWrapper( + self.model, + self.loss, + model_params=model_params, + matmul_precision=self.matmul_precision, + ) self.start_step = 0 # Shared params (multi-task) ------------------------------------------ @@ -2023,10 +2143,13 @@ def _create_full_validators( validation_data: Any | None, ) -> tuple[FullValidator | None, FullValidator | None]: """Create the live-weight and EMA-weight full validators.""" + validation_model = self.model + if self.enable_compile and not self.multi_task: + validation_model = self._unwrapped.model[DEFAULT_TASK_KEY] return build_full_validators( validating_params=validating_params, validation_data=validation_data, - model=self.model, + model=validation_model, state_store=self._unwrapped.train_infos, num_steps=self.num_steps, rank=self.rank, @@ -2038,7 +2161,9 @@ def _create_full_validators( validation_data ), model_ema=self.model_ema, + ema_weight_model=self.model, sharding=self.sharding, + matmul_precision=self.matmul_precision, ) def _raise_if_full_validation_unsupported( @@ -2192,6 +2317,7 @@ def _compile_model(self, compile_opts: dict[str, Any]) -> None: compile_opts=compile_opts, compiled_by_structure=_compiled_by_structure, task_key=task_key, + compile_eval=self.compile_infer, ) log.info( "Compilation enabled (task=%s); the graph is traced and compiled " diff --git a/deepmd/pt_expt/train/utils.py b/deepmd/pt_expt/train/utils.py index 5786ff10d0..63746bae54 100644 --- a/deepmd/pt_expt/train/utils.py +++ b/deepmd/pt_expt/train/utils.py @@ -17,12 +17,25 @@ Any, ) +import torch + if TYPE_CHECKING: from collections.abc import ( Generator, ) - import torch +#: Accepted ``DP_TF32_INFER`` values and the matmul precision each selects. +_TF32_INFER_PRECISIONS = {"0": "highest", "1": "high", "2": "medium"} +_BOOL_ENV_VALUES = { + "1": True, + "true": True, + "yes": True, + "on": True, + "0": False, + "false": False, + "no": False, + "off": False, +} def count_parameters(module: torch.nn.Module) -> tuple[int, int]: @@ -89,6 +102,111 @@ def scoped_env_defaults(defaults: dict[str, str]) -> Generator[None, None, None] os.environ[key] = value +def infer_matmul_precision() -> str: + """ + Resolve the eval-time matmul precision from ``DP_TF32_INFER``. + + Returns + ------- + str + A ``torch.set_float32_matmul_precision`` level: ``"highest"`` for + ``0`` (the default), ``"high"`` for ``1``, ``"medium"`` for ``2``. + + Raises + ------ + ValueError + If ``DP_TF32_INFER`` is set to anything other than 0, 1 or 2. + """ + value = os.environ.get("DP_TF32_INFER", "0").strip().lower() + if value not in _TF32_INFER_PRECISIONS: + raise ValueError(f"DP_TF32_INFER must be one of 0/1/2, got {value!r}") + return _TF32_INFER_PRECISIONS[value] + + +def infer_compile_enabled() -> bool: + """ + Resolve whether evaluation forwards use the compiled model. + + Returns + ------- + bool + Whether ``DP_COMPILE_INFER`` enables compiled evaluation. + + Raises + ------ + ValueError + If ``DP_COMPILE_INFER`` is not a supported boolean value. + """ + value = os.environ.get("DP_COMPILE_INFER", "0").strip().lower() + if value not in _BOOL_ENV_VALUES: + choices = "/".join(_BOOL_ENV_VALUES) + raise ValueError(f"DP_COMPILE_INFER must be one of {choices}, got {value!r}") + return _BOOL_ENV_VALUES[value] + + +class MatmulPrecisionPolicy: + """ + The fp32 matmul precision that each forward mode runs at. + + Training and evaluation are controlled separately, mirroring the split the + pt backend applies inside ``SeZMModel``: training follows + ``training.enable_tf32``, while evaluation follows ``DP_TF32_INFER``, so a + run may train on TF32 yet validate at full precision. + + The eval level is sampled at construction, which must therefore happen + while the ``validating`` section's environment defaults are in scope; see + :func:`scoped_env_defaults`. + + Parameters + ---------- + enable_tf32 : bool + Whether training forwards may use TF32 tensor cores. + + Attributes + ---------- + train_precision : str + Matmul precision of training forwards. + eval_precision : str + Matmul precision of evaluation forwards. + """ + + def __init__(self, enable_tf32: bool) -> None: + self.train_precision = "high" if enable_tf32 else "highest" + self.eval_precision = infer_matmul_precision() + + @contextmanager + def applied(self, *, training: bool) -> Generator[None, None, None]: + """ + Run a block at the matmul precision of the requested mode. + + The previous process-wide level is restored on exit, so a caller that + evaluates in the middle of training does not disturb the surrounding + training precision. The setting only reaches CUDA matmuls, so on a + CPU-only build the block runs unchanged. + + Parameters + ---------- + training : bool + Whether the block is a training forward. + + Yields + ------ + None + Control returns to the caller with the precision in effect. + """ + if not torch.cuda.is_available(): + yield + return + previous = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision( + self.train_precision if training else self.eval_precision + ) + try: + yield + finally: + torch.set_float32_matmul_precision(previous) + + def resolve_best_checkpoint_dir( validating_params: dict[str, Any], save_ckpt: str ) -> Path: diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py index edd89e6d04..28b285d434 100644 --- a/deepmd/pt_expt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -44,6 +44,9 @@ from deepmd.pt_expt.train.ema import ( get_ema_validation_log_path, ) +from deepmd.pt_expt.train.utils import ( + MatmulPrecisionPolicy, +) from deepmd.pt_expt.utils.env import ( DEVICE, GLOBAL_PT_FLOAT_PRECISION, @@ -216,9 +219,15 @@ def __init__( stale_state_keys: tuple[str, ...] = STALE_FULL_VALIDATION_INFO_KEYS, emit_best_save_log: bool = True, model_eval_context: Callable[[], Any] | None = None, + matmul_precision: MatmulPrecisionPolicy | None = None, ) -> None: self.validation_data = validation_data self.model = model + self.matmul_precision = ( + matmul_precision + if matmul_precision is not None + else MatmulPrecisionPolicy(enable_tf32=False) + ) self.profile = select_metric_profile(model) self.state_store = state_store self.rank = rank @@ -380,7 +389,10 @@ def _evaluate(self, display_step: int) -> FullValidationResult: was_training = bool(getattr(self.model, "training", True)) self.model.eval() try: - with self.model_eval_context(): + with ( + self.matmul_precision.applied(training=False), + self.model_eval_context(), + ): # === Step 2. Evaluate All Systems === metrics = self.evaluate_all_systems() finally: @@ -925,7 +937,9 @@ def build_full_validators( checkpoint_dir: Path, ensure_supported: Callable[[], None], model_ema: Any | None = None, + ema_weight_model: torch.nn.Module | dict[str, torch.nn.Module] | None = None, sharding: ShardingPolicy | None = None, + matmul_precision: MatmulPrecisionPolicy | None = None, ) -> tuple[FullValidator | None, FullValidator | None]: """Build the full validators of a training run. @@ -964,9 +978,15 @@ def build_full_validators( The EMA state of the run. Without it the EMA flow stays inactive, so that ``ema_full_validation`` is ignored rather than rejected when EMA itself is disabled. + ema_weight_model : torch.nn.Module or dict, optional + The parameter owner tracked by ``model_ema`` when it differs from the + module used for evaluation. Defaults to ``model``. sharding : ShardingPolicy, optional The distribution strategy of the run, which decides whether checkpoint collection is a collective operation. Defaults to no sharding. + matmul_precision : MatmulPrecisionPolicy, optional + The precision policy of the run, of which both flows use the eval + level. Defaults to full precision. Returns ------- @@ -1000,6 +1020,7 @@ def make(**overrides: Any) -> FullValidator: sharding=ShardingPolicy() if sharding is None else sharding, restart_training=restart_training, checkpoint_dir=checkpoint_dir, + matmul_precision=matmul_precision, **overrides, ) @@ -1022,6 +1043,8 @@ def make(**overrides: Any) -> FullValidator: ), best_checkpoint_prefix=EMA_BEST_CKPT_PREFIX, emit_best_save_log=False, - model_eval_context=lambda: model_ema.apply_shadow(model), + model_eval_context=lambda: model_ema.apply_shadow( + model if ema_weight_model is None else ema_weight_model + ), ) return live_validator, ema_validator diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py index f59b707217..7e3bc76ff1 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/deepmd/pt_expt/train/wrapper.py @@ -15,6 +15,9 @@ from deepmd.dpmodel.utils.multi_task import ( apply_shared_links, ) +from deepmd.pt_expt.train.utils import ( + MatmulPrecisionPolicy, +) log = logging.getLogger(__name__) @@ -69,6 +72,10 @@ class ModelWrapper(torch.nn.Module): Single loss or dict of losses keyed by task name. model_params : dict, optional Model parameters to store as extra state. + matmul_precision : MatmulPrecisionPolicy, optional + The fp32 matmul precision each forward mode runs at. Defaults to full + precision for training, which leaves a wrapper built outside a trainer + (to hold pre-trained weights, say) at the conservative setting. """ def __init__( @@ -76,9 +83,15 @@ def __init__( model: torch.nn.Module | dict, loss: torch.nn.Module | dict | None = None, model_params: dict[str, Any] | None = None, + matmul_precision: MatmulPrecisionPolicy | None = None, ) -> None: super().__init__() self.model_params = model_params if model_params is not None else {} + self.matmul_precision = ( + matmul_precision + if matmul_precision is not None + else MatmulPrecisionPolicy(enable_tf32=False) + ) self.train_infos: dict[str, Any] = { "lr": 0, "step": 0, @@ -173,23 +186,29 @@ def forward( if self.model[task_key].has_spin(): input_dict["spin"] = spin - if self.inference_only: - with self._frozen_parameter_context(): - model_pred = self._forward_without_loss(task_key, input_dict) - return model_pred, None, None - - model_pred = self._forward_without_loss(task_key, input_dict) - if label is None: - return model_pred, None, None - - natoms = atype.shape[-1] - loss, more_loss = self.loss[task_key]( - cur_lr, - natoms, - model_pred, - label, - ) - return model_pred, loss, more_loss + # The module flag distinguishes a training step from the evaluation the + # trainer interleaves with it, which is exactly the split the precision + # policy is defined over. A compiled model is traced inside this block + # on its first call, so its kernels are selected at the same precision + # they later run at. + with self.matmul_precision.applied(training=self.training): + if self.inference_only: + with self._frozen_parameter_context(): + model_pred = self._forward_without_loss(task_key, input_dict) + return model_pred, None, None + + model_pred = self._forward_without_loss(task_key, input_dict) + if label is None: + return model_pred, None, None + + natoms = atype.shape[-1] + loss, more_loss = self.loss[task_key]( + cur_lr, + natoms, + model_pred, + label, + ) + return model_pred, loss, more_loss @contextmanager def _frozen_parameter_context(self) -> Generator[None, None, None]: diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index eae723bc3d..bbbf15acb4 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5742,6 +5742,19 @@ def training_args( "formatted lower-forward path. " "The first training step will be slower due to one-time compilation.", ), + Argument( + "enable_tf32", + bool, + optional=True, + default=False, + doc=supported_backends("pt_expt") + + "Enable TF32 matmul precision for CUDA training forwards. " + "Independent of `enable_compile`; eval-time TF32 is controlled " + "separately by `validating.tf32_infer` or `DP_TF32_INFER`. The " + "PyTorch backend takes the same switch as `model.enable_tf32`, " + "because there only the SeZM model implements the compile and " + "precision path, while here it applies to every model.", + ), ] def _validate_stat_file_mode(data: dict[str, Any], scope: str) -> None: @@ -5928,13 +5941,12 @@ def validating_args() -> Argument: ) doc_compiled_infer = ( "Whether to route eval-time forwards (including full validation) " - "through the DPA4/SeZM `torch.compile` path instead of eager. When `true`, " - "this flag is translated into `DP_COMPILE_INFER=1` at trainer " - "startup before any model is constructed, which is the env var SeZM " - "samples inside `SeZMModel.__init__`. A manually exported " - "`DP_COMPILE_INFER` takes precedence over this option. Only " - "meaningful when `model.use_compile=true`; has no effect on models " - "that do not implement the SeZM-style eval compile path." + "through `torch.compile` instead of eager. When `true`, this flag is " + "translated into `DP_COMPILE_INFER=1` at trainer startup before any " + "model is constructed. A manually exported `DP_COMPILE_INFER` takes " + "precedence over this option. In the PyTorch backend it applies when " + "`model.use_compile=true`; in the PyTorch Exportable backend it applies " + "when `training.enable_compile=true`." ) doc_tf32_infer = ( "Whether to enable TF32 `high` matmul precision for eval-time forwards " @@ -5942,18 +5954,17 @@ def validating_args() -> Argument: "flag is translated into `DP_TF32_INFER=1` at trainer startup before any " "model is constructed. A manually exported `DP_TF32_INFER` takes " "precedence over this option. This does not affect training forwards, " - "which are controlled by `model.enable_tf32`. The PyTorch Exportable " - "backend always runs at full ('highest') matmul precision, so the " - "option has no effect there." + "which are controlled by `model.enable_tf32` (PyTorch) or " + "`training.enable_tf32` (PyTorch Exportable)." ) doc_amp_infer = ( "Whether to enable bf16 automatic mixed precision for eval-time forwards " "(including regular validation and full validation). When `true`, this " "flag is translated into `DP_AMP_INFER=1` at trainer startup before any " "model is constructed. A manually exported `DP_AMP_INFER` takes " - "precedence over this option. This only affects SeZM/DPA4 descriptors " - "with `descriptor.use_amp=true`; training AMP remains controlled by " - "`descriptor.use_amp`." + "precedence over this option. This controls SeZM/DPA4 inference " + "independently of `descriptor.use_amp`; training AMP remains controlled " + "by `descriptor.use_amp`." ) args = [ Argument( @@ -6032,14 +6043,14 @@ def validating_args() -> Argument: bool, optional=True, default=False, - doc=supported_backends("pt") + doc_compiled_infer, + doc=supported_backends("pt", "pt_expt") + doc_compiled_infer, ), Argument( "tf32_infer", bool, optional=True, default=False, - doc=supported_backends("pt") + doc_tf32_infer, + doc=supported_backends("pt", "pt_expt") + doc_tf32_infer, ), Argument( "amp_infer", diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py index 7263566641..5dd3623cba 100644 --- a/deepmd/utils/finetune.py +++ b/deepmd/utils/finetune.py @@ -13,6 +13,9 @@ from deepmd.utils.model_branch_dict import ( get_model_dict, ) +from deepmd.utils.spin import ( + normalize_spin_use_spin, +) log = logging.getLogger(__name__) @@ -197,6 +200,47 @@ def warn_configuration_mismatch_during_finetune( ) +def _native_spin_types( + model_config: Mapping[str, Any], +) -> frozenset[str] | None: + """Return the magnetic element set of a native-spin model config.""" + spin_config = model_config.get("spin") + if ( + not isinstance(spin_config, Mapping) + or str(spin_config.get("scheme", "deepspin")).lower() != "native" + ): + return None + + type_map = model_config["type_map"] + use_spin = normalize_spin_use_spin(spin_config["use_spin"], type_map) + return frozenset( + element for element, enabled in zip(type_map, use_spin, strict=True) if enabled + ) + + +def _validate_native_spin_finetune( + target_config: Mapping[str, Any], + pretrained_config: Mapping[str, Any], +) -> None: + """Validate the native-spin transfer contract before model construction.""" + target_types = _native_spin_types(target_config) + if target_types is None: + return + + pretrained_types = _native_spin_types(pretrained_config) + if pretrained_types is None: + raise ValueError( + "Native-spin fine-tuning requires a native-spin pretrained model; " + "automatic conversion from a spin-free model is not supported." + ) + if pretrained_types and pretrained_types != target_types: + raise ValueError( + "Changing the active magnetic element set during native-spin " + f"fine-tuning is not supported: pretrained={sorted(pretrained_types)}, " + f"target={sorted(target_types)}." + ) + + class FinetuneRuleItem: def __init__( self, @@ -400,6 +444,8 @@ def build_single_rule( model_branch_chosen = model_alias_dict[model_branch_chosen] single_config_chosen = deepcopy(model_dict_params[model_branch_chosen]) + _validate_native_spin_finetune(single_config, single_config_chosen) + old_type_map = single_config_chosen["type_map"] new_type_map = single_config["type_map"] finetune_rule = FinetuneRuleItem( diff --git a/deepmd/utils/spin.py b/deepmd/utils/spin.py index 112ad91b3d..02c75b4b5f 100644 --- a/deepmd/utils/spin.py +++ b/deepmd/utils/spin.py @@ -67,7 +67,7 @@ class Spin: use_spin: list[bool] A list of boolean values indicating whether to use atomic spin for each atom type. True for spin and False for not. List of bool values with shape of [ntypes]. - virtual_scale: list[float], float + virtual_scale: list[float], float, optional The scaling factor to determine the virtual distance between a virtual atom representing spin and its corresponding real atom for each atom type with spin. This factor is defined as the virtual distance @@ -75,6 +75,9 @@ class Spin: The virtual coordinate is defined as the real coordinate plus spin * virtual_scale. List of float values with shape of [ntypes] or [ntypes_spin] or one single float value for all types, only used when use_spin is True for each atom type. + This is a device of the virtual-atom (deepspin) scheme alone. The native + scheme creates no virtual atom, so it leaves this unset and the virtual + scale accessors then raise rather than report a fabricated distance. allow_missing_label: bool Whether a training system that lacks a ``spin`` data file is admitted by filling its per-atom spin with zeros instead of raising. Supported only by @@ -85,7 +88,7 @@ class Spin: def __init__( self, use_spin: list[bool], - virtual_scale: list[float] | float, + virtual_scale: list[float] | float | None = None, allow_missing_label: bool = False, ) -> None: type_dtype = np.int32 @@ -108,31 +111,84 @@ def __init__( np.arange(self.ntypes_real, dtype=type_dtype) + self.ntypes_real ) self.input_type = np.arange(self.ntypes_real * 2, dtype=type_dtype) + if virtual_scale is None: + self.virtual_scale = None + self.virtual_scale_mask = None + else: + self.virtual_scale = self._expand_virtual_scale(virtual_scale) + self.virtual_scale_mask = (self.virtual_scale * self.use_spin).reshape([-1]) + self.pair_exclude_types = [] + self.init_pair_exclude_types_placeholder() + self.atom_exclude_types_ps = [] + self.init_atom_exclude_types_placeholder_spin() + self.atom_exclude_types_p = [] + self.init_atom_exclude_types_placeholder() + + def _expand_virtual_scale(self, virtual_scale: list[float] | float) -> np.ndarray: + """Expand ``virtual_scale`` into one value per real atom type. + + Parameters + ---------- + virtual_scale : list[float] or float + One value per real type, one value per magnetic type (scattered + onto the magnetic types), or a single value shared by all types. + + Returns + ------- + np.ndarray + The per-real-type virtual scale, with shape ``(ntypes_real,)``. + + Raises + ------ + ValueError + If the list length matches neither the real nor the magnetic type + count, or if the value is neither a list nor a float. + """ if isinstance(virtual_scale, list): if len(virtual_scale) == self.ntypes_real: - self.virtual_scale = virtual_scale + expanded = virtual_scale elif len(virtual_scale) == self.ntypes_spin: - self.virtual_scale = np.zeros( - self.ntypes_real, dtype=GLOBAL_NP_FLOAT_PRECISION - ) - self.virtual_scale[self.use_spin] = virtual_scale + expanded = np.zeros(self.ntypes_real, dtype=GLOBAL_NP_FLOAT_PRECISION) + expanded[self.use_spin] = virtual_scale else: raise ValueError( f"Invalid length of virtual_scale for spin atoms" f": Expected {self.ntypes_real} or {self.ntypes_spin} but got {len(virtual_scale)}!" ) elif isinstance(virtual_scale, float): - self.virtual_scale = [virtual_scale for _ in range(self.ntypes_real)] + expanded = [virtual_scale for _ in range(self.ntypes_real)] else: raise ValueError(f"Invalid virtual scale type: {type(virtual_scale)}") - self.virtual_scale = np.array(self.virtual_scale) - self.virtual_scale_mask = (self.virtual_scale * self.use_spin).reshape([-1]) - self.pair_exclude_types = [] - self.init_pair_exclude_types_placeholder() - self.atom_exclude_types_ps = [] - self.init_atom_exclude_types_placeholder_spin() - self.atom_exclude_types_p = [] - self.init_atom_exclude_types_placeholder() + return np.array(expanded) + + @staticmethod + def _require_virtual_scale(value: np.ndarray | None) -> np.ndarray: + """Return a virtual-scale table, rejecting an unconfigured one. + + Parameters + ---------- + value : np.ndarray or None + The table to return, or None when ``virtual_scale`` was not + configured. + + Returns + ------- + np.ndarray + The table unchanged. + + Raises + ------ + ValueError + If the table is absent. + """ + if value is None: + raise ValueError( + "`spin.virtual_scale` is not configured. It is required by the " + "virtual-atom (deepspin) scheme, which places a virtual atom at " + "this distance from its real atom; the native scheme has no " + "virtual atoms and never reads it." + ) + return value def get_ntypes_real(self) -> int: """Returns the number of real atom types.""" @@ -156,7 +212,7 @@ def get_use_spin(self) -> list[bool]: def get_virtual_scale(self) -> np.ndarray: """Returns the list of magnitude of atomic spin for each atom type.""" - return self.virtual_scale + return self._require_virtual_scale(self.virtual_scale) def init_pair_exclude_types_placeholder(self) -> None: """ @@ -244,14 +300,16 @@ def get_virtual_scale_mask(self) -> np.ndarray: Return the virtual scale mask of shape [ntypes], with spin types being its virtual scale, and non-spin types being 0. """ - return self.virtual_scale_mask + return self._require_virtual_scale(self.virtual_scale_mask) def serialize( self, ) -> dict: return { "use_spin": self.use_spin.tolist(), - "virtual_scale": self.virtual_scale.tolist(), + "virtual_scale": ( + None if self.virtual_scale is None else self.virtual_scale.tolist() + ), } @classmethod diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 8cce7b9faf..461b0b163d 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -369,7 +369,7 @@ equivalent input-file option used during training validation: | -------------------- | --------------------------- | ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `DP_COMPILE_INFER` | `validating.compiled_infer` | off | Use the compile path for evaluation/inference. Same `torch==2.11` / CUDA ≥ 12.6 requirements as `model.use_compile`. | | `DP_TF32_INFER` | `validating.tf32_infer` | `0` (highest) | float32 matmul precision for inference: `0` highest, `1` high, `2` medium. Higher values improve throughput but make the potential energy surface less smooth. | -| `DP_AMP_INFER` | `validating.amp_infer` | off | bf16 autocast inside the descriptor interaction blocks for inference when `descriptor.use_amp=true`. Usually keeps aggregate MAE similar but can make the potential energy surface less smooth. | +| `DP_AMP_INFER` | `validating.amp_infer` | off | bf16 autocast inside the descriptor interaction blocks for inference, independently of `descriptor.use_amp`. Training AMP remains controlled by `descriptor.use_amp`. Usually keeps aggregate MAE similar but can make the potential energy surface less smooth. | | `DP_TRITON_INFER` | — | `0` | Triton inference kernel level `0`-`3` (CUDA eval only, compatible with `DP_COMPILE_INFER`). `1`: universal fused kernels, numerically equivalent to the dense path with full float32 accumulation. `2`: adds the table-configured fused SO(2) value path and edge-block backward kernels (still exact float32). `3`: additionally runs the SO(2) mixing stack on fp16 tensor cores with split compensation — roughly float32-level accuracy (maximum force deviation about 4e-6 eV/Å on a 4-thousand-atom system) at a substantial speedup; only shapes validated by the tuning sweep are affected. Levels 2 and 3 read launch tables tuned per GPU model (H20 ships built in); on other GPUs the kernels fall back to conservative configurations, and `dp --pt freeze` tunes the missing entries automatically on the local GPU before exporting (a one-off sweep of a few minutes, baked into the `.pt2`). | Accepted boolean values for the other switches are `1`/`true`/`yes`/`on` and diff --git a/source/tests/common/dpmodel/test_descrpt_dpa4.py b/source/tests/common/dpmodel/test_descrpt_dpa4.py index 93b94f76e7..b0678b3711 100644 --- a/source/tests/common/dpmodel/test_descrpt_dpa4.py +++ b/source/tests/common/dpmodel/test_descrpt_dpa4.py @@ -265,6 +265,67 @@ def test_supported_feature_roundtrip(self, overrides) -> None: out2 = np.asarray(dd2.call(coord.reshape(nf, -1), atype, nlist)[0]) np.testing.assert_array_equal(out1, out2) + def test_legacy_spin_gate_is_squared_on_deserialize(self) -> None: + """Version 1.2 stores the env-seed spin gate after the quadratic form. + + The stored scalar used to multiply the spin coordinate channel + BEFORE that form, so an amplitude ``a`` contributed ``a**2 * D_spin`` + and squaring reproduces the stored function exactly. ``@version`` is + the source of truth, and a migrated payload is retagged so a second + load leaves the gate alone. + """ + dd = make_descriptor(use_spin=[True, False, False]) + data = dd.serialize() + assert data["@version"] == DescrptDPA4.LATEST_VERSION == 1.2 + data["@version"] = 1.1 + data["@variables"]["env_seed_embedding.spin_scale"] = np.full( + (1,), 3.0, dtype=np.float64 + ) + + migrated = DescrptDPA4.deserialize(data) + np.testing.assert_allclose(migrated.env_seed_embedding.spin_scale, 9.0) + assert migrated.version == 1.2 + + reloaded = DescrptDPA4.deserialize(migrated.serialize()) + np.testing.assert_allclose(reloaded.env_seed_embedding.spin_scale, 9.0) + assert reloaded.version == 1.2 + + def test_legacy_spin_free_routes_are_zeroed_on_deserialize(self) -> None: + data = make_descriptor(use_spin=[False, False, False]).serialize() + data["@version"] = 1.1 + dormant_keys = ( + "spin_embedding.mag_layer2.matrix", + "spin_embedding.adam_spin_vec_weight", + "spin_embedding.adam_spin_nbr_weight", + "env_seed_embedding.spin_scale", + ) + for key in dormant_keys: + data["@variables"][key] = np.full_like(data["@variables"][key], 3.0) + mag_layer1_key = "spin_embedding.mag_layer1.matrix" + data["@variables"][mag_layer1_key] = np.full_like( + data["@variables"][mag_layer1_key], 5.0 + ) + + migrated = DescrptDPA4.deserialize(data) + + variables = migrated.serialize()["@variables"] + for key in dormant_keys: + np.testing.assert_array_equal(variables[key], np.zeros_like(variables[key])) + np.testing.assert_array_equal( + variables[mag_layer1_key], np.full_like(variables[mag_layer1_key], 5.0) + ) + assert migrated.version == 1.2 + + def test_pre_spin_versions_keep_their_own_tag(self) -> None: + """Version 1.0 predates the spin route and the 1.1 forward math. + + Promoting it would silently switch ``deg_norm_floor`` and the radial + envelope, so a payload below 1.1 is left exactly as it was written. + """ + data = make_descriptor().serialize() + data["@version"] = 1.0 + assert DescrptDPA4.deserialize(data).version == 1.0 + def test_value_errors(self) -> None: with pytest.raises(ValueError): # kmax must be <= lmax make_descriptor(kmax=4, lmax=3) diff --git a/source/tests/common/dpmodel/test_dpa4_native_spin_model.py b/source/tests/common/dpmodel/test_dpa4_native_spin_model.py index 86ee99309f..f44af3f8a6 100644 --- a/source/tests/common/dpmodel/test_dpa4_native_spin_model.py +++ b/source/tests/common/dpmodel/test_dpa4_native_spin_model.py @@ -113,6 +113,18 @@ def test_dense_route_spin_raises(self): neighbor_graph_method="legacy", ) + def test_output_bias_predictor_rejects_dense_spin_route(self): + self.model.atomic_model.descriptor.disable_graph_lower() + model_forward = self.model.atomic_model._get_forward_wrapper_func() + + with pytest.raises(NotImplementedError, match="NeighborGraph"): + model_forward( + self.coord, + self.atype, + self.box, + spin=self.spin, + ) + def test_deepspin_scheme_with_dpa4_raises(self): cfg = {**NATIVE_SPIN_CONFIG, "spin": {"use_spin": [True, False]}} with pytest.raises(NotImplementedError): diff --git a/source/tests/common/test_finetune_utils.py b/source/tests/common/test_finetune_utils.py index cce2ca1850..33547b34d0 100644 --- a/source/tests/common/test_finetune_utils.py +++ b/source/tests/common/test_finetune_utils.py @@ -30,6 +30,14 @@ def _model_config( } +def _native_spin_model_config( + type_map: list[str], use_spin: list[bool | int | str] +) -> dict: + config = _model_config(type_map) + config["spin"] = {"scheme": "native", "use_spin": use_spin} + return config + + def test_descriptor_normalization_uses_descriptor_type_count(): assert finetune._infer_synthetic_type_count({"sel": [16, 24, 32]}) == 3 assert finetune._infer_synthetic_type_count({"exclude_types": [[0, 3]]}) == 4 @@ -272,3 +280,52 @@ def test_finetune_rule_builder_rejects_multitask_cli_branch(): assert "Multi-task fine-tuning" in str(exc) else: raise AssertionError("expected ValueError") + + +def test_finetune_rule_builder_accepts_compatible_native_spin_transfers(): + compatible_pairs = ( + ( + _native_spin_model_config(["Fe", "C"], [False, False]), + _native_spin_model_config(["Fe", "C"], [True, False]), + ), + ( + _native_spin_model_config(["Fe", "C"], ["Fe"]), + _native_spin_model_config(["C", "Fe"], ["Fe"]), + ), + ( + _native_spin_model_config(["Fe", "C"], ["Fe"]), + _native_spin_model_config(["Fe", "C", "O"], ["Fe"]), + ), + ) + for pretrained, target in compatible_pairs: + finetune.FinetuneRuleBuilder( + pretrained, + target, + change_model_params=False, + ).build() + + +def test_finetune_rule_builder_rejects_non_native_spin_pretraining(): + try: + finetune.FinetuneRuleBuilder( + _model_config(["Fe", "C"]), + _native_spin_model_config(["Fe", "C"], [True, False]), + change_model_params=False, + ).build() + except ValueError as exc: + assert "requires a native-spin pretrained model" in str(exc) + else: + raise AssertionError("expected ValueError") + + +def test_finetune_rule_builder_rejects_changed_magnetic_element_set(): + try: + finetune.FinetuneRuleBuilder( + _native_spin_model_config(["Fe", "C"], ["Fe"]), + _native_spin_model_config(["Fe", "C"], ["C"]), + change_model_params=False, + ).build() + except ValueError as exc: + assert "active magnetic element set" in str(exc) + else: + raise AssertionError("expected ValueError") diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py index a3b70be2f8..5283088b92 100644 --- a/source/tests/pt/model/test_descriptor_sezm.py +++ b/source/tests/pt/model/test_descriptor_sezm.py @@ -85,6 +85,23 @@ def _tiny_two_atom_system( return coord, atype, nlist +def _spin_edge_system( + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """A two-atom edge system with one magnetic and one non-magnetic atom.""" + coord = torch.tensor( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=dtype, device=device + ).view(1, -1, 3) + atype = torch.tensor([[0, 1]], dtype=torch.int64, device=device) + edge_index = torch.tensor([[1, 0], [0, 1]], dtype=torch.long, device=device) + edge_vec = torch.tensor( + [[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]], dtype=dtype, device=device + ) + edge_mask = torch.ones(2, dtype=torch.bool, device=device) + return coord, atype, edge_index, edge_vec, edge_mask + + def _descriptor_kwargs(**overrides) -> dict: """Build a compact SeZM descriptor config for tests.""" kwargs = { @@ -175,36 +192,44 @@ def _assert_forward_backward_smoke(self, **model_kwargs) -> DescrptSeZM: self.assertTrue(torch.all(torch.isfinite(extended_coord.grad))) return model - def test_amp_infer_env_controls_eval_autocast(self) -> None: - """Inference AMP is sampled from env and still gated by ``use_amp``.""" - with mock.patch.dict(os.environ, {"DP_AMP_INFER": "1"}, clear=False): - enabled_model = DescrptSeZM(**_descriptor_kwargs(use_amp=True)) - disabled_model = DescrptSeZM(**_descriptor_kwargs(use_amp=False)) - - enabled_model.eval() - disabled_model.eval() - - with mock.patch("torch.autocast", return_value=nullcontext()) as autocast_mock: - with enabled_model._compute_mode_ctx(torch.device("cuda")): - pass - autocast_mock.assert_called_once_with( - device_type="cuda", - dtype=torch.bfloat16, - enabled=True, - ) - - with mock.patch("torch.autocast", return_value=nullcontext()) as autocast_mock: - with disabled_model._compute_mode_ctx(torch.device("cuda")): - pass - autocast_mock.assert_not_called() - - with mock.patch.dict(os.environ, {"DP_AMP_INFER": "0"}, clear=False): - default_model = DescrptSeZM(**_descriptor_kwargs(use_amp=True)) - default_model.eval() - with mock.patch("torch.autocast", return_value=nullcontext()) as autocast_mock: - with default_model._compute_mode_ctx(torch.device("cuda")): - pass - autocast_mock.assert_not_called() + def test_train_and_eval_amp_switches_are_independent(self) -> None: + """Training follows ``use_amp``, evaluation follows ``DP_AMP_INFER``. + + Neither switch may leak into the other's mode: mixed precision at + inference is a throughput choice that must not require the model to + have been trained with it, and a model trained under AMP must still + deploy at full precision. + """ + models = {} + for amp_infer in (False, True): + with mock.patch.dict( + os.environ, {"DP_AMP_INFER": "1" if amp_infer else "0"}, clear=False + ): + for use_amp in (False, True): + models[amp_infer, use_amp] = DescrptSeZM( + **_descriptor_kwargs(use_amp=use_amp) + ) + + for (amp_infer, use_amp), model in models.items(): + for training in (False, True): + expected = use_amp if training else amp_infer + with self.subTest( + amp_infer=amp_infer, use_amp=use_amp, training=training + ): + model.train(training) + with mock.patch( + "torch.autocast", return_value=nullcontext() + ) as autocast_mock: + with model._compute_mode_ctx(torch.device("cuda")): + pass + if expected: + autocast_mock.assert_called_once_with( + device_type="cuda", + dtype=torch.bfloat16, + enabled=True, + ) + else: + autocast_mock.assert_not_called() def test_cartesian_config_wiring(self) -> None: """Each Cartesian/mixing config builds the intended submodules. @@ -915,23 +940,6 @@ def test_seed_reproducibility(self) -> None: class TestSeZMSpinEmbedding(_SeZMTestCase): """Test the native per-atom spin embedding and its descriptor injection.""" - def _spin_edges( - self, dtype: torch.dtype - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """A two-atom edge system with one magnetic and one non-magnetic atom.""" - coord = torch.tensor( - [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=dtype, device=self.device - ).view(1, -1, 3) - atype = torch.tensor([[0, 1]], dtype=torch.int64, device=self.device) - edge_index = torch.tensor( - [[1, 0], [0, 1]], dtype=torch.long, device=self.device - ) - edge_vec = torch.tensor( - [[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]], dtype=dtype, device=self.device - ) - edge_mask = torch.ones(2, dtype=torch.bool, device=self.device) - return coord, atype, edge_index, edge_vec, edge_mask - def test_cart_to_l1_intertwines_wigner_rotation(self) -> None: """The l=1 spin map must rotate with the descriptor's Wigner-D block.""" dtype = torch.float64 @@ -1041,7 +1049,9 @@ def test_descriptor_spin_joint_rotation_invariance(self) -> None: the backbone branch carries the on-site and neighbor-aggregated l=1. """ dtype = torch.float64 - coord, atype, edge_index, edge_vec, edge_mask = self._spin_edges(dtype) + coord, atype, edge_index, edge_vec, edge_mask = _spin_edge_system( + self.device, dtype + ) spin = torch.zeros(1, 2, 3, dtype=dtype, device=self.device) spin[0, 0] = torch.tensor([0.3, -0.7, 0.5], dtype=dtype, device=self.device) quat = _random_quaternion(1, device=self.device, dtype=dtype) @@ -1096,6 +1106,175 @@ def test_descriptor_spin_joint_rotation_invariance(self) -> None: self.assertFalse(torch.allclose(desc, desc_zero, atol=1e-6)) +class TestSeZMEnvSeedSpinGate(_SeZMTestCase): + """The env-seed spin gate: placement, zero point and version migration.""" + + def _kwargs(self) -> dict: + """Config of a spin descriptor with the env-seed route enabled.""" + return _descriptor_kwargs( + precision="float64", + use_spin=[True, False], + use_env_seed=True, + seed=7, + ) + + def _descriptor(self) -> DescrptSeZM: + """A spin descriptor with non-route weights moved off initialization. + + The env-seed output projection is zero-initialized, so an unperturbed + descriptor is insensitive to the environment matrix. The four + output-controlling spin routes stay at zero so the tests exercise the + same starting point used by a fresh or migrated spin-free model. + """ + model = DescrptSeZM(**self._kwargs()) + spin_routes = { + "spin_embedding.mag_layer2.matrix", + "spin_embedding.adam_spin_vec_weight", + "spin_embedding.adam_spin_nbr_weight", + "env_seed_embedding.spin_scale", + } + torch.manual_seed(0) + with torch.no_grad(): + for name, parameter in model.named_parameters(): + if name not in spin_routes: + parameter.copy_(torch.randn_like(parameter) * 0.1) + model.eval() + return model + + def _inputs(self) -> tuple[dict, torch.Tensor]: + """Edge-route keyword arguments and a spin on the magnetic atom.""" + dtype = torch.float64 + coord, atype, edge_index, edge_vec, edge_mask = _spin_edge_system( + self.device, dtype + ) + spin = torch.zeros(1, 2, 3, dtype=dtype, device=self.device) + spin[0, 0] = torch.tensor([0.3, -0.7, 0.5], dtype=dtype, device=self.device) + return { + "extended_coord": coord, + "extended_atype": atype, + "edge_index": edge_index, + "edge_vec": edge_vec, + "edge_mask": edge_mask, + }, spin + + def test_zero_gate_receives_a_gradient(self) -> None: + """A zero gate is a starting point, not a fixed point. + + The gate multiplies the spin block of the environment matrix after + the quadratic form, so the descriptor is LINEAR in it and its + gradient at zero is that block. A pre-quadratic amplitude would carry + a gradient proportional to itself and could never leave zero, which + is why the gate sits where it does. + """ + model = self._descriptor() + kwargs, spin = self._inputs() + desc, _ = model.forward_with_edges(**kwargs, spin=spin) + desc.sum().backward() + gate_grad = model.env_seed_embedding.spin_scale.grad + self.assertIsNotNone(gate_grad) + self.assertGreater(float(gate_grad.abs().max()), 1e-8) + + def test_gate_placement_differs_from_the_legacy_one_by_a_square(self) -> None: + """The version-1.2 gate ``a**2`` is the version-1.1 amplitude ``a``. + + The amplitude scaled the neighbor-spin channel, which enters the + environment matrix linearly, so the legacy forward is reproduced by a + unit gate on a spin scaled by ``a``. With the remaining spin routes + reset, the env-seed gate is the only difference between the two. + """ + model = self._descriptor() + kwargs, spin = self._inputs() + amplitude = 2.0 + with torch.no_grad(): + model.env_seed_embedding.spin_scale.fill_(amplitude**2) + migrated, _ = model.forward_with_edges(**kwargs, spin=spin) + with torch.no_grad(): + model.env_seed_embedding.spin_scale.fill_(1.0) + legacy, _ = model.forward_with_edges(**kwargs, spin=amplitude * spin) + torch.testing.assert_close(migrated, legacy, atol=1e-12, rtol=1e-12) + + def test_loading_a_legacy_state_squares_the_gate(self) -> None: + """A version-1.1 state is retagged 1.2 with its gate squared. + + The gate is a CHILD parameter, so the migration has to rewrite the + incoming state: the descriptor's own buffers load before torch + descends into ``env_seed_embedding``, and a migration applied to the + live parameter would be overwritten by that descent. + """ + state = self._descriptor().state_dict() + state["version_tensor"] = torch.full_like(state["version_tensor"], 1.1) + state["env_seed_embedding.spin_scale"] = torch.full_like( + state["env_seed_embedding.spin_scale"], 2.0 + ) + model = DescrptSeZM(**self._kwargs()) + model.load_state_dict(state) + torch.testing.assert_close( + model.env_seed_embedding.spin_scale.detach(), + torch.full_like(model.env_seed_embedding.spin_scale, 4.0), + ) + self.assertEqual(model.version, 1.2) + self.assertEqual(float(model.version_tensor.item()), 1.2) + + def test_loading_legacy_spin_free_state_zeros_dormant_routes(self) -> None: + kwargs = {**self._kwargs(), "use_spin": [False, False]} + state = DescrptSeZM(**kwargs).state_dict() + state["version_tensor"] = torch.full_like(state["version_tensor"], 1.1) + dormant_keys = ( + "spin_embedding.mag_layer2.matrix", + "spin_embedding.adam_spin_vec_weight", + "spin_embedding.adam_spin_nbr_weight", + "env_seed_embedding.spin_scale", + ) + for key in dormant_keys: + state[key] = torch.full_like(state[key], 3.0) + state["spin_embedding.mag_layer1.matrix"] = torch.full_like( + state["spin_embedding.mag_layer1.matrix"], 5.0 + ) + + model = DescrptSeZM(**kwargs) + model.load_state_dict(state) + + migrated = model.state_dict() + for key in dormant_keys: + torch.testing.assert_close(migrated[key], torch.zeros_like(migrated[key])) + torch.testing.assert_close( + migrated["spin_embedding.mag_layer1.matrix"], + torch.full_like(migrated["spin_embedding.mag_layer1.matrix"], 5.0), + ) + self.assertEqual(model.version, 1.2) + + def test_a_migrated_state_is_migrated_only_once(self) -> None: + """Re-saving a migrated descriptor advertises 1.2, so a reload is inert.""" + state = self._descriptor().state_dict() + state["version_tensor"] = torch.full_like(state["version_tensor"], 1.1) + state["env_seed_embedding.spin_scale"] = torch.full_like( + state["env_seed_embedding.spin_scale"], 2.0 + ) + migrated = DescrptSeZM(**self._kwargs()) + migrated.load_state_dict(state) + reloaded = DescrptSeZM(**self._kwargs()) + reloaded.load_state_dict(migrated.state_dict()) + torch.testing.assert_close( + reloaded.env_seed_embedding.spin_scale.detach(), + torch.full_like(reloaded.env_seed_embedding.spin_scale, 4.0), + ) + self.assertEqual(reloaded.version, 1.2) + + def test_serialize_roundtrip_migrates_a_legacy_payload(self) -> None: + """The dp-format path shares the rule with the state-dict path.""" + data = self._descriptor().serialize() + data["@version"] = 1.1 + data["@variables"]["env_seed_embedding.spin_scale"] = ( + 0.0 * data["@variables"]["env_seed_embedding.spin_scale"] + 3.0 + ) + model = DescrptSeZM.deserialize(data) + torch.testing.assert_close( + model.env_seed_embedding.spin_scale.detach(), + torch.full_like(model.env_seed_embedding.spin_scale, 9.0), + ) + self.assertEqual(model.version, 1.2) + + class TestBuildEdgeQuaternion(_SeZMTestCase): """Test the stable edge-quaternion chart used by SeZM.""" diff --git a/source/tests/pt/model/test_dpa4_ptexpt_grad_parity.py b/source/tests/pt/model/test_dpa4_ptexpt_grad_parity.py index ce1ce7bb1e..e836bcf8b7 100644 --- a/source/tests/pt/model/test_dpa4_ptexpt_grad_parity.py +++ b/source/tests/pt/model/test_dpa4_ptexpt_grad_parity.py @@ -147,9 +147,28 @@ def _inputs(self, seed=2151): ntypes=self.ntypes, ) - @pytest.mark.parametrize("use_env_seed", [False, True]) # env FiLM (film_* params) - def test_descriptor_grad_parity(self, use_env_seed) -> None: - pt_mod, expt_mod = self._build_pair(use_env_seed=use_env_seed) + @pytest.mark.parametrize( + "overrides", + [ + pytest.param({"use_env_seed": False}, id="no_env_seed"), + pytest.param({"use_env_seed": True}, id="env_seed"), # film_* params + # The grid-product branches are off by default, yet the SO(3) ones + # carry FrameExpand/FrameContract -- the only descriptor weights + # dpmodel stores as bare numpy outside a NativeLayer, and therefore + # the ones most exposed to a missing promotion. The S2 branches + # share the rest of the grid-net structure without those two + # sub-modules. + pytest.param({"message_node_so3": True}, id="message_node_so3"), + pytest.param({"node_wise_so3": True}, id="node_wise_so3"), + pytest.param({"message_node_s2": True}, id="message_node_s2"), + pytest.param({"node_wise_s2": True}, id="node_wise_s2"), + # ReducedEquivariantRMSNorm is reachable only through so2_norm and + # is the one module that sizes its forward off a stored index array + pytest.param({"so2_norm": True}, id="so2_norm"), + ], + ) + def test_descriptor_grad_parity(self, overrides) -> None: + pt_mod, expt_mod = self._build_pair(**overrides) inp = self._inputs() coord = inp["coord"].reshape(self.nf, -1) atype_ext, nlist, mapping = inp["atype_ext"], inp["nlist"], inp["mapping"] diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 92a1cc14a6..8855fe7aa4 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -1613,12 +1613,15 @@ def _build_model( *, use_compile: bool = False, bridging_method: str = "none", + use_spin: list[bool] | None = None, + randomize: bool = True, ) -> SeZMNativeSpinModel: - """Build a tiny float64 native-spin model with randomized parameters.""" + """Build a tiny float64 native-spin model.""" + use_spin = [True, False] if use_spin is None else use_spin params = { "type": "dpa4", "type_map": ["Ni", "O"], - "spin": {"use_spin": [True, False], "scheme": "native"}, + "spin": {"use_spin": use_spin, "scheme": "native"}, "descriptor": { "type": "dpa4", "sel": [12, 12], @@ -1652,12 +1655,13 @@ def _build_model( "bridging_r_outer": 1.2, } model = get_model(params) - # Perturb away from the near-identity initialization so the spin - # embedding measurably shapes the output. - torch.manual_seed(1234) - with torch.no_grad(): - for p in model.parameters(): - p.copy_(torch.randn_like(p) * 0.1) + if randomize: + # Perturb away from the near-identity initialization so the spin + # embedding measurably shapes the output. + torch.manual_seed(1234) + with torch.no_grad(): + for p in model.parameters(): + p.copy_(torch.randn_like(p) * 0.1) model.eval() return model @@ -1692,6 +1696,60 @@ def _frame( ) return coord, atype, spin, box + def test_spin_routes_initialize_to_zero(self) -> None: + """Fresh native-spin models start from the spin-free function.""" + model = self._build_model(randomize=False) + descriptor = model.atomic_model.descriptor + spin_embedding = descriptor.spin_embedding + env_seed_embedding = descriptor.env_seed_embedding + + self.assertIsNotNone(spin_embedding) + self.assertIsNotNone(env_seed_embedding) + self.assertTrue(torch.all(spin_embedding.mag_layer2.matrix == 0.0)) + self.assertTrue(torch.all(spin_embedding.adam_spin_vec_weight == 0.0)) + self.assertTrue(torch.all(spin_embedding.adam_spin_nbr_weight == 0.0)) + self.assertTrue(torch.all(env_seed_embedding.spin_scale == 0.0)) + + def test_migrated_spin_routes_are_trainable_after_activation(self) -> None: + """Mirror the production two-stage native-spin fine-tune load. + + The trainer first rebuilds the pretrained all-false model and loads the + legacy checkpoint into it, allowing version migration to canonicalize + dormant spin routes. It then transfers that migrated state into the + magnetic target. Loading directly into the target would lose the source + configuration needed to distinguish dormant routes from trained ones. + """ + legacy = self._build_model(use_spin=[False, False], randomize=True) + legacy_state = legacy.state_dict() + version_key = "atomic_model.descriptor.version_tensor" + legacy_state[version_key] = torch.full_like(legacy_state[version_key], 1.1) + migrated = self._build_model(use_spin=[False, False], randomize=False) + migrated.load_state_dict(legacy_state) + + model = self._build_model(use_spin=[True, False], randomize=False) + model.load_state_dict(migrated.state_dict()) + coord, atype, spin, box = self._frame() + + model.train() + model.zero_grad(set_to_none=True) + model(coord, atype, spin, box=box)["energy"].sum().backward() + descriptor = model.atomic_model.descriptor + spin_embedding = descriptor.spin_embedding + env_seed_embedding = descriptor.env_seed_embedding + gradients = { + "magnitude": spin_embedding.mag_layer2.matrix.grad, + "onsite_l1": spin_embedding.adam_spin_vec_weight.grad, + "neighbor_l1": spin_embedding.adam_spin_nbr_weight.grad, + "env_seed": env_seed_embedding.spin_scale.grad, + } + for name, gradient in gradients.items(): + self.assertIsNotNone(gradient, f"{name} has no gradient") + self.assertGreater( + float(gradient.abs().max()), + 0.0, + f"{name} cannot leave zero initialization", + ) + def test_zbl_change_out_bias_is_invariant_for_self_labels(self) -> None: """Native-spin statistics consume spin and the complete ZBL energy.""" model = self._build_model(bridging_method="ZBL") diff --git a/source/tests/pt_expt/descriptor/test_dpa4.py b/source/tests/pt_expt/descriptor/test_dpa4.py index ca583ebe6a..ed405be916 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4.py +++ b/source/tests/pt_expt/descriptor/test_dpa4.py @@ -1,5 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import os +from contextlib import ( + nullcontext, +) from unittest import ( mock, ) @@ -103,6 +107,44 @@ def test_consistency(self, use_env_seed, use_mapping) -> None: err_msg=err_msg, ) + def test_train_and_eval_amp_switches_are_independent(self) -> None: + """Training follows ``use_amp``, evaluation follows ``DP_AMP_INFER``. + + The block implementation is stubbed so a CUDA-device test double can + exercise the policy on CPU without constructing or executing a CUDA + tensor. Neither switch may leak into the other's mode. + """ + block_input = mock.Mock(device=mock.Mock(type="cuda")) + block_output = object() + + for amp_infer in (False, True): + with mock.patch.dict( + os.environ, {"DP_AMP_INFER": "1" if amp_infer else "0"}, clear=False + ): + for use_amp in (False, True): + dd = make_descriptor( + self.nt, self.sel_mix, self.rcut, use_amp=use_amp + ) + for training in (False, True): + dd.train(training) + expected = use_amp if training else amp_infer + with ( + mock.patch.object( + DPDescrptDPA4, + "_forward_blocks", + return_value=block_output, + ), + mock.patch( + "torch.autocast", return_value=nullcontext() + ) as autocast_mock, + ): + actual = dd._forward_blocks(block_input) + assert actual is block_output + assert autocast_mock.called is expected, ( + f"amp_infer={amp_infer} use_amp={use_amp} " + f"training={training}" + ) + def test_random_gamma_train_eval_gate(self) -> None: """``random_gamma`` mirrors pt: rolled in train mode, fixed otherwise. diff --git a/source/tests/pt_expt/model/test_dpa2_graph_lower.py b/source/tests/pt_expt/model/test_dpa2_graph_lower.py index 6e1a08481d..a984fc74bf 100644 --- a/source/tests/pt_expt/model/test_dpa2_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa2_graph_lower.py @@ -601,7 +601,9 @@ def test_compiled_training_graph_smoke(self) -> None: model = self._make_model().to("cpu") model.eval() - compiled_lower, buf_order = _trace_and_compile_graph(model, None, None, None) + compiled_lower, buf_order = _trace_and_compile_graph( + model, None, None, None, None + ) assert isinstance(compiled_lower, torch.nn.Module) assert buf_order == () @@ -619,7 +621,7 @@ def test_compiled_training_graph_smoke(self) -> None: atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample compiled_out = compiled_lower( - atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs, None ) eager = model.forward_common_lower_graph( atype, @@ -892,7 +894,7 @@ def test_compiled_training_graph_small_sel(self) -> None: model = self._make_model(repinit_nsel=10, repformer_nsel=6).to("cpu") model.eval() - compiled_lower, _ = _trace_and_compile_graph(model, None, None, None) + compiled_lower, _ = _trace_and_compile_graph(model, None, None, None, None) sample = build_synthetic_graph_inputs( model, @@ -907,7 +909,7 @@ def test_compiled_training_graph_small_sel(self) -> None: ) atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample compiled_out = compiled_lower( - atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs, None ) eager = model.forward_common_lower_graph( atype, @@ -1011,9 +1013,9 @@ def test_graph_lower_fparam_symbolic_trace_and_compile(self) -> None: ) # (b) compiled-training path (fparam threaded through the compile) - compiled_lower, _ = _trace_and_compile_graph(model, fp, None, None) + compiled_lower, _ = _trace_and_compile_graph(model, fp, None, None, None) compiled_out = compiled_lower( - atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs, None ) ctol = {"rtol": 1e-10, "atol": 1e-10} torch.testing.assert_close(compiled_out["energy"], ref["energy_redu"], **ctol) diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py index e16a08172a..dd98f89059 100644 --- a/source/tests/pt_expt/model/test_dpa4_native_spin.py +++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py @@ -730,6 +730,70 @@ def test_graph_lower_exportable_torch_export(self) -> None: ) +class TestDPA4NativeSpinCompiledTraining: + """Native-spin data flow through the compiled training lower.""" + + def test_eager_eval_and_compiled_magnetic_force_backward( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Keep eval eager while retaining compiled magnetic gradients.""" + from deepmd.pt_expt.train.training import ( + _CompiledModel, + _get_model_structure_key, + ) + + # ``make_fx`` and the rebuilt graph exercise the compile boundary. The + # identity backend keeps this regression independent of platform + # toolchains while preserving the traced forward and double backward. + monkeypatch.setattr(torch, "compile", lambda model, **_: model) + + model = _build_native_spin_model_cpu() + compiled = _CompiledModel( + model, + _get_model_structure_key(model), + compile_eval=False, + ) + coord = torch.tensor( + [ + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 0.0], + [1.0, 0.0, 1.0], + ] + ], + dtype=torch.float64, + ) + atype = torch.tensor([[0, 0, 0, 1, 1, 1]], dtype=torch.int64) + spin = torch.arange(18, dtype=torch.float64).reshape(1, 6, 3) / 20.0 + 0.1 + box = (8.0 * torch.eye(3, dtype=torch.float64)).reshape(1, 9) + + eager_result = compiled.eval()(coord, atype, box=box, spin=spin) + assert compiled._compiled_lower_by_mode == {} + assert eager_result["force_mag"].shape == (1, 6, 3) + + result = compiled.train()(coord, atype, box=box, spin=spin) + + assert result["force"].shape == (1, 6, 3) + assert result["force_mag"].shape == (1, 6, 3) + assert result["mask_mag"].shape == (1, 6, 1) + assert set(compiled._compiled_lower_by_mode) == {True} + result["force_mag"].square().sum().backward() + + spin_parameters = [ + parameter + for name, parameter in model.named_parameters() + if "spin_embedding" in name + ] + assert spin_parameters + assert any( + parameter.grad is not None and torch.any(parameter.grad != 0) + for parameter in spin_parameters + ) + + # ============================================================================= # Task 11: training smoke -- native-spin DPA4 through the real pt_expt # trainer (data loading, ``ener_spin`` loss dispatch, ``ModelWrapper``, @@ -908,6 +972,182 @@ def test_training_smoke(self, tmp_path) -> None: os.chdir(old_cwd) +class TestDPA4NativeSpinFinetuneFromSpinFreePretrain: + """Fine-tune a native-spin DPA4 on top of a pretraining with no magnetic type. + + The production workflow is a large corpus that declares no magnetic species + followed by a small magnetic one. Both stages are native-spin, so the + parameter trees agree and transfer wholesale, while ``use_spin`` differs -- + and the per-type spin gate it produces must be rebuilt by the fine-tuned + model rather than inherited, since the pretraining gate is all zero. + Fine-tuning also recomputes the output bias from statistics, which evaluates + the model and therefore requires the spin input to reach it. + + The assertion is the symptom a dead gate produces: a magnetic-force error of + exactly zero, reported for every step of an otherwise healthy run. + """ + + def setup_method(self) -> None: + self.data_dir = os.path.join( + os.path.dirname(__file__), "..", "..", "pt", "NiO", "data", "single" + ) + if not os.path.isdir(self.data_dir): + pytest.skip(f"NiO spin data not found: {self.data_dir}") + + def _normalized(self, config: dict) -> dict: + return normalize(update_deepmd_input(copy.deepcopy(config), warning=False)) + + def test_magnetic_force_is_live_after_finetune(self, tmp_path) -> None: + from deepmd.pt_expt.utils.finetune import ( + get_finetune_rules, + ) + + old_cwd = os.getcwd() + os.chdir(tmp_path) + try: + # Phase 1: pretrain with no type declared magnetic, then checkpoint. + pretrain_config = _make_train_config(self.data_dir, numb_steps=1) + pretrain_config["model"]["spin"]["use_spin"] = [False, False] + get_trainer(self._normalized(pretrain_config)).run() + checkpoint = tmp_path / "model.ckpt.pt" + assert checkpoint.is_file(), "pretraining wrote no checkpoint" + checkpoint_data = torch.load( + checkpoint, + map_location="cpu", + weights_only=True, + ) + checkpoint_state = checkpoint_data["model"] + dormant_keys = [ + key + for key, value in checkpoint_state.items() + if torch.is_tensor(value) + and ( + "spin_embedding.mag_layer2." in key + or "spin_embedding.adam_spin_" in key + or "env_seed_embedding.spin_scale" in key + ) + ] + assert dormant_keys, "checkpoint exposes no dormant spin parameters" + for key in dormant_keys: + checkpoint_state[key] = torch.randn_like(checkpoint_state[key]) + version_keys = [ + key + for key in checkpoint_state + if key.endswith("descriptor.version_tensor") + ] + assert len(version_keys) == 1 + version_key = version_keys[0] + checkpoint_state[version_key] = torch.full_like( + checkpoint_state[version_key], 1.1 + ) + torch.save(checkpoint_data, checkpoint) + + # Phase 2: fine-tune with Ni magnetic, through the production rules. + finetune_config = self._normalized( + _make_train_config(self.data_dir, numb_steps=1) + ) + finetune_config["model"], finetune_links = get_finetune_rules( + str(checkpoint), finetune_config["model"] + ) + trainer = get_trainer( + finetune_config, + finetune_model=str(checkpoint), + finetune_links=finetune_links, + ) + descriptor = trainer.models["Default"].atomic_model.descriptor + assert torch.all(descriptor.spin_embedding.mag_layer2.w == 0.0) + assert torch.all(descriptor.spin_embedding.adam_spin_vec_weight == 0.0) + assert torch.all(descriptor.spin_embedding.adam_spin_nbr_weight == 0.0) + assert torch.all(descriptor.env_seed_embedding.spin_scale == 0.0) + + task = trainer.select_task(trainer._make_training_tasks()) + more_loss = trainer.train_step(task, 0).payload["more_loss"] + assert "rmse_fm" in more_loss + magnetic_error = float(torch.as_tensor(more_loss["rmse_fm"]).detach()) + assert magnetic_error > 0.0, ( + "the magnetic force error is exactly zero after fine-tuning: the " + "per-type spin gate was inherited from the spin-free pretraining " + "instead of being rebuilt from `use_spin`" + ) + finally: + os.chdir(old_cwd) + + +class TestDPA4DescriptorVersionPersistence: + """pt_expt persists the descriptor version so migrations still fire. + + The backend rebuilds a module from its config before loading, so a + version kept only as a python attribute would come back claiming the + semantics of the running code and every checkpoint would silently skip + ``_migrate_variables``. + """ + + def _descriptor(self, use_spin: list[bool] | None = None) -> DescrptDPA4: + config = copy.deepcopy(NATIVE_SPIN_CONFIG) + if use_spin is not None: + config["spin"]["use_spin"] = use_spin + return get_model(config).atomic_model.descriptor + + def test_version_rides_the_state_dict(self) -> None: + descriptor = self._descriptor() + assert "version_tensor" in descriptor.state_dict() + assert float(descriptor.version_tensor.item()) == DescrptDPA4.LATEST_VERSION + + def test_legacy_state_squares_the_env_seed_spin_gate(self) -> None: + """Version 1.2 stores the gate after the environment quadratic form. + + The gate is a child parameter, so the rewrite has to land on the + incoming state: torch restores a module's own buffers before + descending into its children. + """ + state = self._descriptor().state_dict() + state["version_tensor"] = torch.full_like(state["version_tensor"], 1.1) + gate_key = "env_seed_embedding.spin_scale" + state[gate_key] = torch.full_like(state[gate_key], 2.0) + + target = self._descriptor() + target.load_state_dict(state) + assert torch.all(target.env_seed_embedding.spin_scale == 4.0) + assert target.version == 1.2 + assert float(target.version_tensor.item()) == 1.2 + + def test_legacy_spin_free_state_zeros_dormant_routes(self) -> None: + state = self._descriptor([False, False]).state_dict() + state["version_tensor"] = torch.full_like(state["version_tensor"], 1.1) + dormant_keys = ( + "spin_embedding.mag_layer2.w", + "spin_embedding.adam_spin_vec_weight", + "spin_embedding.adam_spin_nbr_weight", + "env_seed_embedding.spin_scale", + ) + for key in dormant_keys: + state[key] = torch.full_like(state[key], 3.0) + state["spin_embedding.mag_layer1.w"] = torch.full_like( + state["spin_embedding.mag_layer1.w"], 5.0 + ) + + target = self._descriptor([False, False]) + target.load_state_dict(state) + + migrated = target.state_dict() + for key in dormant_keys: + assert torch.all(migrated[key] == 0.0) + assert torch.all(migrated["spin_embedding.mag_layer1.w"] == 5.0) + assert target.version == 1.2 + + def test_state_without_a_version_is_read_as_the_last_untagged_one(self) -> None: + """Checkpoints predating the buffer were written under version 1.1.""" + state = self._descriptor().state_dict() + del state["version_tensor"] + gate_key = "env_seed_embedding.spin_scale" + state[gate_key] = torch.full_like(state[gate_key], 3.0) + + target = self._descriptor() + target.load_state_dict(state) + assert torch.all(target.env_seed_embedding.spin_scale == 9.0) + assert target.version == 1.2 + + class TestNativeSpinConfigFormsPtExpt: """pt_expt twin of ``test_dpa4_native_spin_model.py::TestNativeSpinConfigForms``. diff --git a/source/tests/pt_expt/model/test_get_model_dpa4.py b/source/tests/pt_expt/model/test_get_model_dpa4.py index aa76fd7ecc..d0990c0aa1 100644 --- a/source/tests/pt_expt/model/test_get_model_dpa4.py +++ b/source/tests/pt_expt/model/test_get_model_dpa4.py @@ -2,10 +2,8 @@ """Tests for the DPA4/SeZM model-type dispatch in pt_expt ``get_model``.""" import copy -import logging import unittest -import pytest import torch from deepmd.pt_expt.model import ( @@ -195,12 +193,13 @@ def test_unsupported_keys_raise(self) -> None: """pt-only SeZM model-level features fail fast with NotImplementedError. ``bridging_method`` is no longer in this list: it is supported as an - atomic-model composition (see ``test_zbl_bridging.py``). + atomic-model composition (see ``test_zbl_bridging.py``). Neither is + ``use_compile``, which pt_expt relocated to ``training.enable_compile`` + and merely warns about. """ cases = { "spin": ({"use_spin": [True, False], "virtual_scale": [0.3]}, "Spin DPA4"), "lora": ({"rank": 4}, "`lora` is not supported"), - "use_compile": (True, "`use_compile` is not supported"), "preset_out_bias": ( {"energy": [None, 1.0]}, "`preset_out_bias` is not supported", @@ -271,53 +270,6 @@ def test_default_unsupported_values_pass(self) -> None: self.assertIsInstance(model, EnergyModel) -# `enable_tf32` toggles TF32 matmul precision in pt but is ignored by pt_expt -# (always "highest" precision); a truthy value must emit a warn-once message. -@pytest.mark.parametrize("enable_tf32", [True, False]) # truthy warns, falsy silent -def test_enable_tf32_warns_once(enable_tf32, monkeypatch) -> None: - import importlib - - # the package __init__ rebinds the name ``get_model`` to the function, so - # ``import ...get_model as`` would shadow the submodule; load it explicitly - gm_mod = importlib.import_module("deepmd.pt_expt.model.get_model") - - # reset the warn-once set so the assertion is deterministic regardless of - # test ordering (other get_sezm_model calls may have already warned) - monkeypatch.setattr(gm_mod, "_WARNED_ONCE", set()) - - # Count emissions on the EMITTING logger with our own handler rather than - # through caplog: caplog reads a root handler, so whatever global logging - # state earlier tests left behind (set_log_handles flips the ``deepmd`` - # logger's propagate off and installs its own handlers) changes how many - # records reach it -- zero when propagation is off, more than one when the - # record is seen through several attached handlers. A handler on the - # emitting logger sees exactly one record per ``log.warning`` call. - records: list[logging.LogRecord] = [] - - class _Collect(logging.Handler): - def emit(self, record: logging.LogRecord) -> None: - records.append(record) - - handler = _Collect(level=logging.WARNING) - old_level = gm_mod.log.level - gm_mod.log.setLevel(logging.WARNING) - gm_mod.log.addHandler(handler) - try: - gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32)) - matches = [r for r in records if "enable_tf32" in r.getMessage()] - if enable_tf32: - assert len(matches) == 1, [r.getMessage() for r in records] - # a second call must NOT warn again (warn-once per process) - records.clear() - gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32)) - assert not [r for r in records if "enable_tf32" in r.getMessage()] - else: - assert not matches, [r.getMessage() for r in records] - finally: - gm_mod.log.removeHandler(handler) - gm_mod.log.setLevel(old_level) - - class TestNativeSpinErrorTranslation(unittest.TestCase): """Only the unexpected-``use_spin`` TypeError becomes the capability error.""" diff --git a/source/tests/pt_expt/test_finetune.py b/source/tests/pt_expt/test_finetune.py index a97e8e208c..bdba008ec8 100644 --- a/source/tests/pt_expt/test_finetune.py +++ b/source/tests/pt_expt/test_finetune.py @@ -1004,5 +1004,106 @@ def test_finetune_from_pt2_use_pretrain_script(self) -> None: shutil.rmtree(tmpdir, ignore_errors=True) +# Native-spin descriptors, paired with the fitting types they require. Each +# descriptor entry is deliberately tiny: these tests inspect the parameter tree +# of a freshly built model, not its predictions. +_NATIVE_SPIN_DESCRIPTORS = { + "dpa4": ( + { + "type": "dpa4", + "rcut": 4.0, + "sel": 8, + "channels": 8, + "n_radial": 4, + "lmax": 1, + "mmax": 1, + "n_blocks": 1, + "precision": "float64", + "seed": 17, + }, + "dpa4_ener", + ), +} + + +def _native_spin_config(descriptor_key: str, use_spin: list[bool]) -> dict: + """Build a native-spin model config over the type map ``["Fe", "C"]``.""" + descriptor, fitting_type = _NATIVE_SPIN_DESCRIPTORS[descriptor_key] + config = { + "type_map": ["Fe", "C"], + "descriptor": deepcopy(descriptor), + "fitting_net": { + "type": fitting_type, + "neuron": [8], + "precision": "float64", + "seed": 19, + }, + "spin": {"use_spin": use_spin, "scheme": "native"}, + } + if descriptor_key == "dpa4": + config["type"] = "dpa4" + return config + + +class TestNativeSpinGateNotInherited(unittest.TestCase): + """The per-type spin gate is rebuilt from the configuration, never inherited. + + A native-spin model is commonly pretrained on a corpus that declares no + magnetic species and then fine-tuned on a magnetic one. Both models are + native-spin, so their parameter trees agree and transfer wholesale, but the + gate that ``use_spin`` produces is not model state: adopting the pretraining + copy leaves it all zero, which silences the spin channel and pins the + magnetic force at exactly zero for the whole fine-tuning run. + """ + + @staticmethod + def _gates(model: torch.nn.Module) -> dict[str, torch.Tensor]: + """Return every per-type spin gate of *model*, keyed by buffer name.""" + return { + name: buffer + for name, buffer in model.named_buffers() + if name.endswith("spin_mask") + } + + def test_gate_is_rebuilt_not_inherited(self) -> None: + for descriptor_key in _NATIVE_SPIN_DESCRIPTORS: + with self.subTest(descriptor=descriptor_key): + pretrained = get_model( + _native_spin_config(descriptor_key, [False, False]) + ) + finetuned = get_model( + _native_spin_config(descriptor_key, [True, False]) + ) + + # Being configuration-derived, the gate is not part of the state. + self.assertEqual( + [ + key + for key in finetuned.state_dict() + if key.endswith("spin_mask") + ], + [], + ) + expected = { + name: gate.clone() for name, gate in self._gates(finetuned).items() + } + self.assertTrue(expected, "the model declares no per-type spin gate") + for name, gate in expected.items(): + self.assertTrue( + bool((gate != 0).any()), f"{name} is all zero as built" + ) + + # A checkpoint that still archives the gate -- written before it + # became non-persistent -- must neither fail a strict load nor + # override the built value. + archived = dict(pretrained.state_dict()) + archived.update( + {name: torch.zeros_like(gate) for name, gate in expected.items()} + ) + finetuned.load_state_dict(archived) + for name, gate in self._gates(finetuned).items(): + torch.testing.assert_close(gate, expected[name], rtol=0, atol=0) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index 0569f21851..45eb76fa03 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -610,7 +610,7 @@ class TestCompiledModelGetattr(unittest.TestCase): These tests do not require example data or torch.compile — they use a lightweight mock original_model to verify that __getattr__ correctly forwards unknown attributes/methods to the wrapped original model. - Compilation is lazy, so no compiled_forward_lower is needed for construction. + Compilation is lazy, so no compiled graph is needed for construction. """ def _make_compiled_model(self): @@ -653,11 +653,11 @@ def test_own_attrs_not_delegated(self) -> None: """Attributes owned by _CompiledModel itself are NOT delegated.""" cm = self._make_compiled_model() # original_model is a registered submodule and must not fall through - # to delegation. compiled_forward_lower is None before the first - # forward call (lazy compile) — accessing it must return None, not - # delegate to original_model. + # to delegation. The compiled-graph cache is empty before the first + # forward call (lazy compile) — accessing it must return that empty + # dict, not delegate to original_model. self.assertIsInstance(cm.original_model, torch.nn.Module) - self.assertIsNone(cm.compiled_forward_lower) + self.assertEqual(cm._compiled_lower_by_mode, {}) def test_missing_attr_raises(self) -> None: """Accessing an attribute missing from both wrapper and original raises.""" @@ -703,8 +703,8 @@ def test_compiled_handles_varying_nall(self) -> None: # The wrapper.model should be a _CompiledModel compiled_model = trainer.wrapper.model["Default"] self.assertIsInstance(compiled_model, _CompiledModel) - # Lazy compile: compiled_forward_lower is None before any forward. - self.assertIsNone(compiled_model.compiled_forward_lower) + # Lazy compile: no graph exists before any forward. + self.assertEqual(compiled_model._compiled_lower_by_mode, {}) trainer.wrapper.train() for step in range(3): @@ -715,9 +715,10 @@ def test_compiled_handles_varying_nall(self) -> None: loss.backward() trainer.optimizer.step() - # After first forward, compiled_forward_lower must be set. + # After the first forward the training graph must exist, + # keyed by the mode it was traced in. if step == 0: - self.assertIsNotNone(compiled_model.compiled_forward_lower) + self.assertIn(True, compiled_model._compiled_lower_by_mode) # Loss should be a finite scalar at every step self.assertFalse(torch.isnan(loss)) @@ -759,6 +760,7 @@ def _build_config(enable_compile: bool) -> dict: config["model"]["fitting_net"]["activation_function"] = activation if enable_compile: config["training"]["enable_compile"] = True + config["validating"] = {"compiled_infer": True} config = update_deepmd_input(config, warning=False) return normalize(config)