diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 69349fc5b0..04a49b9c41 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -96,6 +96,10 @@ from deepmd.pt.utils.learning_rate import ( BaseLR, ) +from deepmd.pt.utils.spin_finetune import ( + compute_native_spin_input_stats, + get_native_spin_finetune_state, +) from deepmd.pt.utils.stat import ( make_stat_input, pair_margin_frame_mask, @@ -770,6 +774,24 @@ def collect_single_finetune_params( _pretrained_model_params["descriptor"], _model_key_from, ) + if _finetune_rule_single.convert_to_native_spin: + log.info( + "Initialize native-spin branch '%s' from spin-free branch '%s'.", + _model_key, + _model_key_from, + ) + converted = get_native_spin_finetune_state( + self.wrapper.model[_model_key], + pretrained_model_wrapper.model[_model_key_from], + random_fitting=_new_fitting, + ) + _new_state_dict.update( + { + f"model.{_model_key}.{key}": value + for key, value in converted.items() + } + ) + return target_keys = [ i for i in _random_state_dict.keys() @@ -846,11 +868,16 @@ def update_finetune_bias( model = ( self.model[_model_key] if self.multi_task else self.model ) - model = model_change_out_bias( - model, + sample_func = ( self.get_sample_func[_model_key] if self.multi_task - else self.get_sample_func, + else self.get_sample_func + ) + if _finetune_rule.convert_to_native_spin: + compute_native_spin_input_stats(model, sample_func) + model = model_change_out_bias( + model, + sample_func, _bias_adjust_mode="change-by-statistic" if not _finetune_rule.get_random_fitting() else "set-by-statistic", diff --git a/deepmd/pt/utils/finetune.py b/deepmd/pt/utils/finetune.py index bf1c328408..ad918701da 100644 --- a/deepmd/pt/utils/finetune.py +++ b/deepmd/pt/utils/finetune.py @@ -21,6 +21,33 @@ __all__ = ["get_finetune_rule_single", "get_finetune_rules"] +def _learned_part_path(model: torch.nn.Module) -> str | None: + """Return the state-dict path of the learned descriptor-fitting part of a model. + + A plain model holds that part as its atomic model and a composition with + one learned sub-model, such as a zone-bridging one, as that sub-model, so + fine-tuning between the two transfers the part across the paths. A model + without a single learned part returns ``None``. + """ + atomic_model = model.atomic_model + if getattr(atomic_model, "descriptor", None) is not None: + return "atomic_model." + learned = [ + index + for index, child in enumerate(getattr(atomic_model, "models", ())) + if getattr(child, "descriptor", None) is not None + ] + return f"atomic_model.models.{learned[0]}." if len(learned) == 1 else None + + +def _move_learned_key(key: str, path: str, source_path: str) -> str: + """Move a descriptor or fitting-net key between learned atomic-model paths.""" + for part in ("descriptor.", "fitting_net."): + if f".{path}{part}" in key: + return key.replace(f".{path}{part}", f".{source_path}{part}", 1) + return key + + def get_finetune_rules( finetune_model: str, model_config: dict[str, Any], diff --git a/deepmd/pt/utils/spin_finetune.py b/deepmd/pt/utils/spin_finetune.py new file mode 100644 index 0000000000..064eed85df --- /dev/null +++ b/deepmd/pt/utils/spin_finetune.py @@ -0,0 +1,342 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Spin-free checkpoint transfer for the PyTorch and pt_expt native-spin models.""" + +from collections.abc import ( + Callable, +) +from functools import ( + singledispatch, +) + +import torch + +from deepmd.dpmodel.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, +) +from deepmd.dpmodel.fitting.dpa4_ener import ( + GLUFittingNet, +) +from deepmd.dpmodel.utils.network import ( + FittingNet, + NativeLayer, +) +from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, +) +from deepmd.pt.utils.finetune import ( + _learned_part_path, +) + + +def get_native_spin_finetune_state( + model: torch.nn.Module, + pretrained_model: torch.nn.Module, + *, + random_fitting: bool = False, +) -> dict[str, torch.Tensor]: + """Build native-spin state from a spin-free learned atomic model. + + Type maps must already agree. The learned parts are matched independently + of analytical bridging. Output biases and scales belong to the target's + subsequent statistics pass; spin-only parameters keep their constructor + values, while target data calibrates spin-only input statistics afterward. + Descriptor-specific transfer rules preserve the geometric function while + leaving a nonzero gradient into the initially closed spin routes. + + Parameters + ---------- + model + Fresh native-spin target, with its input configuration and spin mask. + pretrained_model + Loaded spin-free source after type-map alignment. + random_fitting + Whether to retain the target fitting network instead of transferring it. + + Returns + ------- + dict[str, torch.Tensor] + Complete state dict for a strict load into the target model. + + Raises + ------ + ValueError + If the learned parts, descriptor or fitting families, type maps, + network layouts or tensor shapes are incompatible with spin-free transfer, + or widening the fitting input changes an implicit residual connection. + """ + path = _learned_part_path(model) + source_path = _learned_part_path(pretrained_model) + if path is None or source_path is None: + raise ValueError( + "Spin-free transfer requires one learned atomic model in each model." + ) + if model.get_type_map() != pretrained_model.get_type_map(): + raise ValueError("Spin-free transfer requires aligned type maps.") + learned = model.get_submodule(path.removesuffix(".")) + source = pretrained_model.get_submodule(source_path.removesuffix(".")) + descriptor = learned.descriptor + if type(descriptor) is not type(source.descriptor): + raise ValueError( + "Spin-free transfer requires the same descriptor family: " + f"source={type(source.descriptor).__name__}, " + f"target={type(descriptor).__name__}." + ) + if descriptor.use_spin is None or source.descriptor.use_spin is not None: + raise ValueError( + "Spin-free transfer requires a spin-free source descriptor and a " + "native-spin target descriptor." + ) + + expanded = _expanded_spin_state( + descriptor, source, learned, random_fitting=random_fitting + ) + if not random_fitting: + _validate_fitting_networks(source.fitting_net, learned.fitting_net) + state = model.state_dict() + for key, value in source.state_dict().items(): + target_key = path + key + if target_key not in state or key.rsplit(".", 1)[-1] in { + "out_bias", + "out_std", + "bias_atom_e", + }: + continue + if random_fitting and not key.startswith("descriptor."): + continue + if key in expanded: + state[target_key] = expanded[key] + else: + _check_shape(key, value.shape, state[target_key].shape) + state[target_key] = value.detach().clone() + return state + + +def compute_native_spin_input_stats( + model: torch.nn.Module, sampled_func: Callable[[], list[dict]] +) -> None: + """Calibrate newly introduced spin state after spin-free weight transfer. + + Geometric input statistics remain those of the pretrained descriptor. + Only DPA4C requires a data-dependent reference for its spin invariants. + + Parameters + ---------- + model + Native-spin target after a successful spin-free state transfer. + sampled_func + Sampled training systems containing atom types and magnetic moments. + """ + path = _learned_part_path(model) + descriptor = model.get_submodule(path.removesuffix(".")).descriptor + if isinstance(descriptor, DescrptDPA4C): + descriptor.spin.set_spin_reference( + descriptor._measure_spin_reference(sampled_func()) + ) + + +@singledispatch +def _expanded_spin_state( + descriptor: object, + source: torch.nn.Module, + target: torch.nn.Module, + *, + random_fitting: bool, +) -> dict[str, torch.Tensor]: + """Expand state under the descriptor's native-spin layout contract.""" + raise ValueError( + f"Spin-free transfer is not registered for {type(descriptor).__name__}." + ) + + +@_expanded_spin_state.register(DescrptSeZM) +@_expanded_spin_state.register(DescrptDPA4) +def _expanded_dpa4_state( + descriptor: object, + source: torch.nn.Module, + target: torch.nn.Module, + *, + random_fitting: bool, +) -> dict[str, torch.Tensor]: + # DPA4 injects zero-initialized additive spin routes without widening the + # geometric descriptor or its fitting input. + return {} + + +@_expanded_spin_state.register(DescrptDPA4C) +def _expanded_dpa4c_state( + descriptor: DescrptDPA4C, + source: torch.nn.Module, + target: torch.nn.Module, + *, + random_fitting: bool, +) -> dict[str, torch.Tensor]: + # The descriptor is [readout, spin, masses, type embedding]. Any frame or + # atomic fitting parameters follow this descriptor and shift with its tail. + start = descriptor.readout.get_dim_out() + width = descriptor.spin.get_dim_out() + expanded = ( + {} + if random_fitting + else _expanded_fitting_input_state( + source.fitting_net, target.fitting_net, start=start, width=width + ) + ) + source_state = source.state_dict() + target_state = target.state_dict() + for key in ("descriptor.mean", "descriptor.stddev"): + expanded[key] = _insert_spin_block( + key, + source_state[key], + target_state[key], + axis=0, + start=start, + width=width, + ) + # Ordered spin scale/shift columns follow every geometric FiLM head. + key = "descriptor.pair_film.network.layers.1.w" + expanded[key] = _insert_spin_block( + key, + source_state[key], + target_state[key], + axis=1, + start=source_state[key].shape[1], + width=2 * descriptor.spin_channels, + ) + return expanded + + +@singledispatch +def _fitting_input_projection(network: object) -> NativeLayer: + """Return the affine layer that reads the fitting input.""" + raise ValueError( + "Spin-free input expansion is not registered for fitting network " + f"{type(network).__name__}." + ) + + +@_fitting_input_projection.register(FittingNet) +def _mlp_input_projection(network: FittingNet) -> NativeLayer: + return network.layers[0] + + +@_fitting_input_projection.register(GLUFittingNet) +def _glu_input_projection(network: GLUFittingNet) -> NativeLayer: + return network.hidden_layers[0] if network.hidden_layers else network.output_layer + + +def _expanded_fitting_input_state( + source: torch.nn.Module, + target: torch.nn.Module, + *, + start: int, + width: int, +) -> dict[str, torch.Tensor]: + """Expand fitting input projections without changing their residual semantics.""" + if source.nets.ndim != target.nets.ndim: + raise ValueError("Spin-free transfer requires the same fitting type layout.") + expanded: dict[int, torch.Tensor] = {} + for index in range(target.nets.ntypes**target.nets.ndim): + source_net, target_net = source.nets[index], target.nets[index] + if type(source_net) is not type(target_net): + raise ValueError( + "Spin-free transfer requires matching fitting network types: " + f"source={type(source_net).__name__}, " + f"target={type(target_net).__name__}." + ) + source_layer = _fitting_input_projection(source_net) + target_layer = _fitting_input_projection(target_net) + # NativeLayer derives its residual and timestep from the input/output + # widths. An inserted input block changes an identity or duplicated + # residual, which cannot in general be absorbed into the nonlinear + # branch's weights. Reject that structural change before loading state. + for layer in (source_layer, target_layer): + if layer.resnet and layer.dim_out() in ( + layer.dim_in(), + 2 * layer.dim_in(), + ): + raise ValueError( + "Spin-free transfer cannot preserve the first fitting layer's " + "residual connection when its input expands from " + f"{source_layer.dim_in()} to {target_layer.dim_in()} " + f"with output width {target_layer.dim_out()}. " + "Use a compatible pretrained fitting architecture or select " + "finetune_head='RANDOM' to initialize a fresh fitting network." + ) + # A closed descriptor spin gate preserves the geometric function while + # fresh fitting rows give that gate a nonzero gradient. + expanded[id(target_layer)] = _insert_spin_block( + f"fitting_net.nets[{index}].input_projection.w", + source_layer.w, + target_layer.w, + axis=0, + start=start, + width=width, + ) + # Network collections expose multiple state-dict paths for the same layer. + # Every alias refers to the same expanded weight, computed once per network. + return { + f"fitting_net.{name}.w": expanded[id(module)] + for name, module in target.named_modules(remove_duplicate=False) + if id(module) in expanded + } + + +def _validate_fitting_networks( + source: torch.nn.Module, target: torch.nn.Module +) -> None: + """Require complete fitting-network transfer, including frozen weights.""" + if type(source) is not type(target): + raise ValueError( + "Spin-free transfer requires matching fitting network types: " + f"source={type(source).__name__}, target={type(target).__name__}." + ) + # PT and pt_expt register their networks under these respective paths. + # Root buffers hold output biases, statistics and runtime defaults, not + # the learned network layout. Frozen network weights remain in state_dict. + source_keys, target_keys = ( + { + key + for key in fitting.state_dict() + if key.startswith(("filter_layers.", "nets.")) + } + for fitting in (source, target) + ) + if source_keys != target_keys: + raise ValueError( + "Spin-free transfer requires matching fitting network state keys: " + f"source-only={sorted(source_keys - target_keys)}, " + f"target-only={sorted(target_keys - source_keys)}. " + "Use a compatible fitting architecture or select finetune_head='RANDOM'." + ) + + +def _check_shape(key: str, expected: tuple[int, ...], actual: tuple[int, ...]) -> None: + if actual != expected: + raise ValueError( + f"Incompatible spin-free transfer for {key}: " + f"expected target shape {tuple(expected)}, got {tuple(actual)}." + ) + + +def _insert_spin_block( + key: str, + source: torch.Tensor, + target: torch.Tensor, + *, + axis: int, + start: int, + width: int, +) -> torch.Tensor: + """Copy the source around an inserted block, retaining its target values.""" + expected = list(source.shape) + expected[axis] += width + _check_shape(key, tuple(expected), target.shape) + source = source.detach() + result = target.detach().clone() + result.narrow(axis, 0, start).copy_(source.narrow(axis, 0, start)) + tail = source.shape[axis] - start + result.narrow(axis, start + width, tail).copy_(source.narrow(axis, start, tail)) + return result diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index c231f428f1..051701dc96 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -110,6 +110,14 @@ strip_saved_tensor_detach as _strip_saved_tensor_detach, ) from deepmd.pt.utils.compile_compat import trace_pad_dim as _trace_pad_dim +from deepmd.pt.utils.finetune import ( + _learned_part_path, + _move_learned_key, +) +from deepmd.pt.utils.spin_finetune import ( + compute_native_spin_input_stats, + get_native_spin_finetune_state, +) from deepmd.pt_expt.loss import ( DOSLoss, EnergyLoss, @@ -2274,34 +2282,42 @@ def initialize_statistics( # Selective weight copy (per-branch key remapping) pretrained_state = pretrained_wrapper.state_dict() target_state = self._unwrapped.state_dict() - # A plain model and its zone-bridging composition hold the same - # learned descriptor and fitting net at different paths. - learned_paths = { - model_key: ( - _learned_part_path(self._unwrapped.model[model_key]), - _learned_part_path( - pretrained_wrapper.model[ - finetune_links[model_key].get_model_branch() - ] - ), - ) - for model_key in self.model_keys - } - new_state = {} - for key in target_state: - if key == "_extra_state": - new_state[key] = target_state[key] + new_state = target_state.copy() + for model_key in self.model_keys: + finetune_rule = finetune_links[model_key] + source_key = finetune_rule.get_model_branch() + model = self._unwrapped.model[model_key] + source_model = pretrained_wrapper.model[source_key] + prefix = f"model.{model_key}." + if finetune_rule.convert_to_native_spin: + log.info( + "Initialize native-spin branch '%s' from spin-free branch '%s'.", + model_key, + source_key, + ) + converted = get_native_spin_finetune_state( + model, + source_model, + random_fitting=finetune_rule.get_random_fitting(), + ) + new_state.update( + {prefix + key: value for key, value in converted.items()} + ) continue - # Find which model_key this key belongs to - matched = False - for model_key in self.model_keys: - if f".{model_key}." not in key: + + # Plain and bridged models store their learned part at + # different paths; the analytical part keeps target state. + path = _learned_part_path(model) + source_path = _learned_part_path(source_model) + for key in target_state: + if not key.startswith(prefix): continue - matched = True - finetune_rule = finetune_links[model_key] - _key_from = finetune_rule.get_model_branch() - pretrained_key = key.replace(f".{model_key}.", f".{_key_from}.") - path, source_path = learned_paths[model_key] + if ( + finetune_rule.get_random_fitting() + and ".descriptor." not in key + ): + continue + pretrained_key = f"model.{source_key}." + key[len(prefix) :] if ( pretrained_key not in pretrained_state and path is not None @@ -2310,19 +2326,8 @@ def initialize_statistics( pretrained_key = _move_learned_key( pretrained_key, path, source_path ) - use_random = ( - finetune_rule.get_random_fitting() - and ".descriptor." not in key - ) - if use_random: - new_state[key] = target_state[key] - elif pretrained_key in pretrained_state: + if pretrained_key in pretrained_state: new_state[key] = pretrained_state[pretrained_key] - else: - new_state[key] = target_state[key] - break - if not matched: - new_state[key] = target_state[key] self._unwrapped.load_state_dict(new_state) # Per-branch bias adjustment (rank 0 only, then broadcast) @@ -2342,7 +2347,12 @@ def initialize_statistics( def update_finetune_bias( _model_key: str = model_key, _bias_mode: str = bias_mode, + _convert_to_native_spin: bool = finetune_rule.convert_to_native_spin, ) -> None: + if _convert_to_native_spin: + compute_native_spin_input_stats( + self.models[_model_key], self._sample_funcs[_model_key] + ) self.models[_model_key] = model_change_out_bias( self.models[_model_key], self._sample_funcs[_model_key], @@ -3706,33 +3716,6 @@ def _more_loss_to_float(cls, more_loss: dict[str, Any]) -> dict[str, float]: } -def _learned_part_path(model: torch.nn.Module) -> str | None: - """Return the state-dict path of the learned descriptor-fitting part of a model. - - A plain model holds that part as its atomic model and a composition with - one learned sub-model, such as a zone-bridging one, as that sub-model, so - fine-tuning between the two transfers the part across the paths. A model - without a single learned part returns ``None``. - """ - atomic_model = model.atomic_model - if getattr(atomic_model, "descriptor", None) is not None: - return "atomic_model." - learned = [ - index - for index, child in enumerate(getattr(atomic_model, "models", ())) - if getattr(child, "descriptor", None) is not None - ] - return f"atomic_model.models.{learned[0]}." if len(learned) == 1 else None - - -def _move_learned_key(key: str, path: str, source_path: str) -> str: - """Rewrite a descriptor or fitting-net key from one learned-part path to another.""" - for part in ("descriptor.", "fitting_net."): - if f".{path}{part}" in key: - return key.replace(f".{path}{part}", f".{source_path}{part}", 1) - return key - - def model_change_out_bias( _model: Any, _sample_func: Any, diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py index 5dd3623cba..945805c592 100644 --- a/deepmd/utils/finetune.py +++ b/deepmd/utils/finetune.py @@ -221,17 +221,25 @@ def _native_spin_types( 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.""" +) -> bool: + """Validate spin transfer and report whether a spin-free source needs conversion.""" target_types = _native_spin_types(target_config) if target_types is None: - return + return False pretrained_types = _native_spin_types(pretrained_config) if pretrained_types is None: + if pretrained_config.get("spin") is None: + if target_config["type_map"] != pretrained_config["type_map"]: + raise ValueError( + "Spin-free to native-spin fine-tuning requires identical ordered " + f"type maps: pretrained={pretrained_config['type_map']}, " + f"target={target_config['type_map']}." + ) + return True raise ValueError( - "Native-spin fine-tuning requires a native-spin pretrained model; " - "automatic conversion from a spin-free model is not supported." + "Native-spin fine-tuning requires a spin-free or native-spin " + "pretrained model." ) if pretrained_types and pretrained_types != target_types: raise ValueError( @@ -239,6 +247,86 @@ def _validate_native_spin_finetune( f"fine-tuning is not supported: pretrained={sorted(pretrained_types)}, " f"target={sorted(target_types)}." ) + return False + + +def _learned_model_config( + model: Mapping[str, Any], +) -> Mapping[str, Any] | None: + """Find the unique descriptor-bearing part of a plain or composed model.""" + candidates = model["models"] if model.get("type") == "linear_ener" else (model,) + learned = [candidate for candidate in candidates if "descriptor" in candidate] + return learned[0] if len(learned) == 1 else None + + +def _inherit_native_spin_fitting_widths( + target: dict[str, Any], pretrained: Mapping[str, Any] +) -> None: + """Resolve automatic target widths against the pretrained fitting layout.""" + from deepmd.utils.argcheck import ( + descrpt_args_plugin, + fitting_args_plugin, + ) + + target_learned = _learned_model_config(target) + source_learned = _learned_model_config(pretrained) + if target_learned is None or source_learned is None: + return + descriptor_schema = descrpt_args_plugin.get_argument("dpa4c") + fitting_schema = fitting_args_plugin.get_argument("dpa4_ener") + descriptor_types = {descriptor_schema.name, *descriptor_schema.alias} + glu_types = {fitting_schema.name, *fitting_schema.alias} + fitting = target_learned.get("fitting_net", {}) + source_fitting = source_learned.get("fitting_net", {}) + if ( + target_learned["descriptor"].get("type") not in descriptor_types + or source_learned["descriptor"].get("type") not in descriptor_types + or fitting.get("type") not in glu_types + or source_fitting.get("type") not in glu_types + ): + return + neuron = fitting_schema.normalize_value(fitting).get("neuron") + target_widths = [0] if neuron is None else list(neuron) + if 0 not in target_widths: + return + source_fitting = fitting_schema.normalize_value(source_fitting) + source_neuron = source_fitting.get("neuron") + source_widths = [0] if source_neuron is None else list(source_neuron) + if len(target_widths) != len(source_widths): + raise ValueError( + "Automatic native-spin fitting widths require matching layer counts: " + f"pretrained={len(source_widths)}, target={len(target_widths)}." + ) + if 0 in source_widths: + from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, + ) + from deepmd.dpmodel.fitting.dpa4_ener import ( + _resolve_auto_neuron, + ) + + descriptor_params = descriptor_schema.normalize_value( + source_learned["descriptor"] + ) + descriptor_params.pop("type") + type_map = source_learned.get("type_map", pretrained["type_map"]) + descriptor_params.update(ntypes=len(type_map), type_map=list(type_map)) + descriptor = DescrptDPA4C(**descriptor_params) + source_widths = _resolve_auto_neuron( + source_widths, + dim_descrpt=descriptor.get_dim_out(), + numb_fparam=source_fitting.get("numb_fparam", 0), + numb_aparam=source_fitting.get("numb_aparam", 0), + dim_case_embd=source_fitting.get("dim_case_embd", 0), + case_film_embd=source_fitting.get("case_film_embd", False), + use_aparam_as_mask=source_fitting.get("use_aparam_as_mask", False), + ) + for alias in fitting_schema.sub_fields["neuron"].alias: + fitting.pop(alias, None) + fitting["neuron"] = [ + source_width if target_width == 0 else target_width + for source_width, target_width in zip(source_widths, target_widths, strict=True) + ] class FinetuneRuleItem: @@ -249,6 +337,7 @@ def __init__( model_branch: str = "Default", random_fitting: bool = False, resuming: bool = False, + convert_to_native_spin: bool = False, ) -> None: """ The rules for fine-tuning the model from pretrained model. @@ -265,12 +354,15 @@ def __init__( If true, the fitting net will be randomly initialized instead of inherit from the pretrained model. resuming If true, the model will just resume from model_branch without fine-tuning. + convert_to_native_spin + If true, initialize the native-spin model from a spin-free branch. """ self.p_type_map = p_type_map self.type_map = type_map self.model_branch = model_branch self.random_fitting = random_fitting self.resuming = resuming + self.convert_to_native_spin = convert_to_native_spin self.update_type = self.p_type_map != self.type_map def get_index_mapping(self) -> list[int]: @@ -398,7 +490,9 @@ def build(self) -> tuple[dict[str, Any], dict[str, FinetuneRuleItem]]: model_branch_from=model_branch_from, ) ) - finetune_rule.resuming = resuming + finetune_rule.resuming = ( + resuming and not finetune_rule.convert_to_native_spin + ) finetune_links[model_key] = finetune_rule return model_config, finetune_links @@ -444,7 +538,9 @@ 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) + convert_to_native_spin = _validate_native_spin_finetune( + single_config, single_config_chosen + ) old_type_map = single_config_chosen["type_map"] new_type_map = single_config["type_map"] @@ -453,6 +549,7 @@ def build_single_rule( type_map=new_type_map, model_branch=model_branch_chosen, random_fitting=new_fitting, + convert_to_native_spin=convert_to_native_spin, ) if self.change_model_params: self._apply_pretrained_model_params( @@ -462,6 +559,8 @@ def build_single_rule( model_branch_chosen=model_branch_chosen, random_fitting=new_fitting, ) + if convert_to_native_spin and not new_fitting: + _inherit_native_spin_fitting_widths(single_config, single_config_chosen) return single_config, finetune_rule def _apply_pretrained_model_params( diff --git a/doc/train/finetuning.md b/doc/train/finetuning.md index 5995c79836..273d6bf628 100644 --- a/doc/train/finetuning.md +++ b/doc/train/finetuning.md @@ -74,7 +74,8 @@ If you wish to conduct fine-tuning using a randomly initialized fitting net in t $ dp --pt train input.json --finetune pretrained.pt --model-branch RANDOM ``` -The model section in input.json **must be the same as that in the pretrained model**. +The model section in input.json **must be the same as that in the pretrained model**, +apart from the native-spin initialization described below. If you do not know the model params in the pretrained model, you can add `--use-pretrain-script` in the fine-tuning command: ```bash @@ -91,6 +92,26 @@ The model section will be overwritten (except the `type_map` subsection) by that } ``` +#### Initializing native spin from a spin-free checkpoint + +PyTorch DPA4/SeZM and PyTorch Exportable DPA4/DPA4C models support spin-free +pretraining with a native-spin target. The source and target must use the same +descriptor family and identical ordered `type_map` values. Set the target's +`spin.scheme` to `native` and configure `use_spin`; these settings are retained +even with `--use-pretrain-script`. The same rule applies to each selected branch +in multi-task fine-tuning. + +For DPA4C with `dpa4_ener` or `sezm_ener` fitting, automatic hidden widths +(an omitted `neuron` or zero entries) inherit the pretrained fitting widths. +The resolved widths are stored in the target configuration so that checkpoint +reconstruction uses the same architecture. Explicit widths are not rewritten +and must remain compatible with weight transfer. `RANDOM` instead initializes +the target fitting network and resolves its automatic widths from the target. + +The transfer retains pretrained geometric statistics and initializes new spin +routes with a closed gate. DPA4C's per-type spin-reference magnitudes are measured +from the fine-tuning data before adjusting the output bias. + #### Fine-tuning from a multi-task pre-trained model Additionally, within the PyTorch implementation and leveraging the flexibility offered by the framework and the multi-task training process proposed in DPA2 [paper](https://doi.org/10.1038/s41524-024-01493-2), diff --git a/source/tests/common/test_finetune_utils.py b/source/tests/common/test_finetune_utils.py index 33547b34d0..92ad767271 100644 --- a/source/tests/common/test_finetune_utils.py +++ b/source/tests/common/test_finetune_utils.py @@ -4,6 +4,8 @@ deepcopy, ) +import pytest + from deepmd.utils import ( finetune, ) @@ -298,22 +300,62 @@ def test_finetune_rule_builder_accepts_compatible_native_spin_transfers(): ), ) for pretrained, target in compatible_pairs: - finetune.FinetuneRuleBuilder( + _, rules = finetune.FinetuneRuleBuilder( pretrained, target, change_model_params=False, ).build() + assert not rules["Default"].convert_to_native_spin + + +def test_finetune_rule_builder_accepts_spin_free_pretraining() -> None: + pretrained = _model_config(["Fe", "C"]) + target = _native_spin_model_config(["Fe", "C"], ["Fe"]) + target["preset_out_bias"] = {"energy": {"Fe": -1.0, "C": -2.0}} + target["bridging_method"] = "zbl" + original = deepcopy(target) + for change_model_params in (False, True): + config, rules = finetune.FinetuneRuleBuilder( + pretrained, target, change_model_params=change_model_params + ).build() + assert rules["Default"].convert_to_native_spin + assert config["spin"] == original["spin"] + assert config["preset_out_bias"] == original["preset_out_bias"] + assert config["bridging_method"] == "zbl" + assert target == original -def test_finetune_rule_builder_rejects_non_native_spin_pretraining(): - try: +def test_spin_free_multitask_branch_requires_finetune_statistics() -> None: + pretrained = {"model_dict": {"magnetic": _model_config(["Fe", "C"])}} + target = { + "model_dict": {"magnetic": _native_spin_model_config(["Fe", "C"], ["Fe"])} + } + _, rules = finetune.FinetuneRuleBuilder(pretrained, target).build() + assert rules["magnetic"].convert_to_native_spin + assert not rules["magnetic"].get_resuming() + + +@pytest.mark.parametrize("type_map", [["Fe", "C", "Ni"], ["Fe"], ["C", "Fe"]]) +def test_spin_free_conversion_rejects_changed_type_map(type_map: list[str]) -> None: + with pytest.raises(ValueError, match="identical ordered type maps"): finetune.FinetuneRuleBuilder( _model_config(["Fe", "C"]), + _native_spin_model_config(type_map, ["Fe"]), + change_model_params=False, + ).build() + + +def test_finetune_rule_builder_rejects_virtual_spin_pretraining(): + pretrained = _model_config(["Fe", "C"]) + pretrained["spin"] = {"use_spin": [True, False], "scheme": "deepspin"} + try: + finetune.FinetuneRuleBuilder( + pretrained, _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) + assert "requires a spin-free or native-spin pretrained model" in str(exc) else: raise AssertionError("expected ValueError") diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index ecead9a672..a73674e48a 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -185,7 +185,7 @@ def test_fine_tune_locates_the_learned_part_by_structure() -> None: A descriptor exclusion gives the child its own pair exclusion, which keeps the composition off the fused route, but not its learned part off its path. """ - from deepmd.pt_expt.train.training import ( + from deepmd.pt.utils.finetune import ( _learned_part_path, ) diff --git a/source/tests/pt_expt/test_finetune_native_spin.py b/source/tests/pt_expt/test_finetune_native_spin.py new file mode 100644 index 0000000000..88b453f36a --- /dev/null +++ b/source/tests/pt_expt/test_finetune_native_spin.py @@ -0,0 +1,871 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Spin-free checkpoint transfer through both native-spin training backends.""" + +import json +from copy import ( + deepcopy, +) +from importlib import ( + import_module, +) +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import numpy as np +import pytest +import torch + +from deepmd.pt.entrypoints.main import main as pt_main +from deepmd.pt.model.model import get_model as get_pt_model +from deepmd.pt.train.wrapper import ModelWrapper as PTModelWrapper +from deepmd.pt.utils.finetune import ( + _learned_part_path, +) +from deepmd.pt.utils.spin_finetune import ( + compute_native_spin_input_stats, + get_native_spin_finetune_state, +) +from deepmd.pt_expt.entrypoints.main import ( + main, +) +from deepmd.pt_expt.model import ( + get_model, +) +from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, +) +from deepmd.pt_expt.utils.env import ( + DEVICE, +) +from deepmd.utils.argcheck import ( + normalize, +) +from deepmd.utils.bridging import ( + expand_bridging_method, +) +from deepmd.utils.compat import ( + update_deepmd_input, +) +from deepmd.utils.finetune import ( + FinetuneRuleBuilder, +) +from deepmd.utils.model_preset import ( + get_model_preset, +) + +from .test_finetune import ( + _make_config, + _native_spin_config, +) + +_BACKENDS = [("pt", "dpa4"), ("pt_expt", "dpa4"), ("pt_expt", "dpa4c")] + + +def _config(descriptor: str, native: bool, bridged: bool = False) -> dict: + config = _native_spin_config(descriptor, [True, False]) + if not native: + config.pop("spin") + if bridged: + config["bridging_method"] = "zbl" + config["data_stat_nbatch"] = 1 + return config + + +def _canonical_config(config: dict) -> dict: + config = expand_bridging_method(config) + config["models"].reverse() + for child in config["models"]: + child.pop("type_map", None) + return config + + +def _build(backend: str, config: dict) -> torch.nn.Module: + builder = get_pt_model if backend == "pt" else get_model + return builder(deepcopy(config)).to(DEVICE) + + +def _learned(model: torch.nn.Module) -> torch.nn.Module: + return model.get_submodule(_learned_part_path(model).removesuffix(".")) + + +def _pretrained(backend: str, config: dict) -> torch.nn.Module: + """Build a geometric model whose nonzero forces make transfer observable.""" + model = _build(backend, config) + generator = torch.Generator(device=DEVICE).manual_seed(31) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.add_( + 0.02 + * torch.randn( + parameter.shape, + dtype=parameter.dtype, + device=parameter.device, + generator=generator, + ) + ) + descriptor = _learned(model).descriptor + for name, start, stop in (("mean", -0.2, 0.2), ("stddev", 0.8, 1.2)): + array = getattr(descriptor, name) + array.copy_( + torch.linspace( + start, stop, array.numel(), dtype=array.dtype, device=array.device + ).reshape(array.shape) + ) + return model + + +def _inputs() -> dict[str, torch.Tensor]: + return { + "coord": torch.tensor( + [[[0.0, 0.0, 0.0], [1.3, 0.2, 0.1], [0.3, 1.5, 0.4]]], + dtype=torch.float64, + device=DEVICE, + ), + "atype": torch.tensor([[0, 1, 0]], device=DEVICE), + "box": (6.0 * torch.eye(3, dtype=torch.float64, device=DEVICE)).reshape(1, 9), + "spin": torch.tensor( + [[[0.6, 0.2, -0.7], [0.1, 0.5, 0.1], [-0.4, 0.8, 0.3]]], + dtype=torch.float64, + device=DEVICE, + ), + } + + +@pytest.mark.parametrize("backend, descriptor", _BACKENDS) +@pytest.mark.parametrize("bridged", [False, True]) +def test_spin_free_transfer_preserves_function_and_trains_spin( + backend: str, descriptor: str, bridged: bool +) -> None: + source = _pretrained(backend, _config(descriptor, False, bridged)) + target = _build(backend, _config(descriptor, True, bridged)) + target.load_state_dict(get_native_spin_finetune_state(target, source)) + inputs = _inputs() + spin = inputs.pop("spin") + expected = source(**inputs, do_atomic_virial=True) + assert expected["force"].abs().max() > 1e-8 + for scale in (0.0, 1.0, 2.0): + actual = target(**inputs, spin=scale * spin, do_atomic_virial=True) + for key in ("energy", "force", "virial", "atom_virial"): + torch.testing.assert_close( + actual[key], expected[key], rtol=1e-11, atol=1e-11 + ) + torch.testing.assert_close( + actual["force_mag"], torch.zeros_like(actual["force_mag"]), rtol=0, atol=0 + ) + + spin_params = [param for name, param in target.named_parameters() if "spin" in name] + optimizer = torch.optim.SGD(spin_params, lr=1e-3) + (actual["force_mag"] - 0.2).square().sum().backward() + assert any( + param.grad is not None and param.grad.abs().max() > 1e-8 + for param in spin_params + ) + optimizer.step() + trained = target(**inputs, spin=spin) + assert trained["force_mag"].abs().max() > 1e-10 + assert torch.count_nonzero(trained["force_mag"][:, 1]) == 0 + + +@pytest.mark.parametrize("backend, descriptor", _BACKENDS) +@pytest.mark.parametrize("random_fitting", [False, True]) +def test_spin_free_transfer_preserves_target_bias_and_bridging( + backend: str, descriptor: str, random_fitting: bool +) -> None: + source = _pretrained(backend, _config(descriptor, False)) + target = _build(backend, _config(descriptor, True, bridged=True)) + initial = {key: value.clone() for key, value in target.state_dict().items()} + with torch.no_grad(): + for key, value in source.state_dict().items(): + if key.rsplit(".", 1)[-1] in {"out_bias", "out_std", "bias_atom_e"}: + value.fill_(1234.0) + state = get_native_spin_finetune_state( + target, source, random_fitting=random_fitting + ) + target.load_state_dict(state) + for key, value in state.items(): + if key.rsplit(".", 1)[-1] in {"out_bias", "out_std", "bias_atom_e"} or ( + random_fitting and ".fitting_net." in key + ): + torch.testing.assert_close(value, initial[key], rtol=0, atol=0) + source_descriptor = _learned(source).descriptor.state_dict() + target_descriptor = _learned(target).descriptor.state_dict() + for key, value in source_descriptor.items(): + if value.shape == target_descriptor[key].shape: + torch.testing.assert_close(value, target_descriptor[key], rtol=0, atol=0) + np.testing.assert_array_equal(target.spin.use_spin, [True, False]) + + +@pytest.mark.parametrize("backend, descriptor", _BACKENDS) +@pytest.mark.parametrize("width", [4, 16]) +def test_spin_free_transfer_rejects_unrelated_shape_change( + backend: str, descriptor: str, width: int +) -> None: + source = _build(backend, _config(descriptor, False)) + config = _config(descriptor, True) + config["fitting_net"]["neuron"] = [width] + target = _build(backend, config) + with pytest.raises(ValueError, match="Incompatible spin-free transfer"): + get_native_spin_finetune_state(target, source) + + +@pytest.mark.parametrize( + "invalid, message", + [ + ("source_without_learned_part", "one learned atomic model"), + ("target_without_learned_part", "one learned atomic model"), + ("type_map", "aligned type maps"), + ("descriptor_family", "same descriptor family"), + ("native_source", "spin-free source descriptor"), + ("spin_free_target", "native-spin target descriptor"), + ], +) +def test_spin_free_transfer_rejects_incompatible_models( + invalid: str, message: str +) -> None: + source = _build("pt_expt", _config("dpa4", invalid == "native_source")) + target_config = _config( + "dpa4c" if invalid == "descriptor_family" else "dpa4", + invalid != "spin_free_target", + ) + if invalid == "type_map": + target_config["type_map"] = ["C", "Fe"] + target = _build("pt_expt", target_config) + if invalid.endswith("without_learned_part"): + analytical = torch.nn.Module() + analytical.add_module("atomic_model", torch.nn.Module()) + if invalid.startswith("source"): + source = analytical + else: + target = analytical + with pytest.raises(ValueError, match=message): + get_native_spin_finetune_state(target, source) + + +@pytest.mark.parametrize("magnitude", [0.0, 2.0]) +def test_native_spin_statistics_preserve_pretrained_geometry(magnitude: float) -> None: + source = _pretrained("pt_expt", _config("dpa4c", False)) + target = _build("pt_expt", _config("dpa4c", True)) + target.load_state_dict(get_native_spin_finetune_state(target, source)) + inputs = { + "coord": torch.tensor( + [[[0.0, 0.0, 0.0], [1.3, 0.0, 0.0]]], dtype=torch.float64, device=DEVICE + ), + "atype": torch.zeros((1, 2), dtype=torch.int64, device=DEVICE), + "spin": torch.tensor( + [[[0.0, 0.0, magnitude], [0.0, 0.0, magnitude]]], + dtype=torch.float64, + device=DEVICE, + ), + } + descriptor = _learned(target).descriptor + initial = {key: value.clone() for key, value in descriptor.state_dict().items()} + # Equivalent environments need no recalibration of pretrained geometry. + compute_native_spin_input_stats(target, lambda: [inputs]) + torch.testing.assert_close( + descriptor.spin.spin_reference, + torch.tensor([magnitude or 1.0, 1.0, 1.0], dtype=torch.float64, device=DEVICE), + rtol=0, + atol=0, + ) + for key, value in descriptor.state_dict().items(): + if key != "spin.spin_reference": + torch.testing.assert_close(value, initial[key], rtol=0, atol=0) + actual = target(**inputs) + expected = source(coord=inputs["coord"], atype=inputs["atype"]) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(actual[key], expected[key], rtol=1e-11, atol=1e-11) + assert torch.count_nonzero(actual["force_mag"]) == 0 + + +def test_fresh_dpa4c_spin_statistics_keep_closed_gate_scale() -> None: + descriptor = _learned(_build("pt_expt", _config("dpa4c", True))).descriptor + descriptor.compute_input_stats([_inputs()]) + start = descriptor.readout.get_dim_out() + stop = start + descriptor.spin.get_dim_out() + assert torch.count_nonzero(descriptor.spin.spin_gate) == 0 + torch.testing.assert_close( + descriptor.mean[start:stop], + torch.zeros_like(descriptor.mean[start:stop]), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + descriptor.stddev[start:stop], + torch.ones_like(descriptor.stddev[start:stop]), + rtol=0, + atol=0, + ) + + +@pytest.mark.parametrize( + "fitting_type, neuron, source_form, target_form", + [ + (fitting_type, neuron, "plain", "plain") + for fitting_type in ("dpa4_ener", "sezm_ener") + for neuron in ("omitted", None, [0], [0, 8], [8, 0]) + ] + + [ + ("dpa4_ener", "omitted", "concise", "canonical"), + ("sezm_ener", None, "concise", "canonical"), + ("sezm_ener", [0], "canonical", "concise"), + ("dpa4_ener", [0, 8], "canonical", "canonical"), + ], +) +def test_dpa4c_auto_width_inherits_source_architecture( + fitting_type: str, + neuron: str | list[int] | None, + source_form: str, + target_form: str, +) -> None: + source_config = _config("dpa4c", False, bridged=source_form != "plain") + source_config["fitting_net"]["type"] = fitting_type + if neuron == "omitted": + source_config["fitting_net"].pop("neuron") + else: + source_config["fitting_net"]["neuron"] = neuron + target_config = deepcopy(source_config) + target_config["spin"] = {"scheme": "native", "use_spin": [True, False]} + if source_form == "canonical": + source_config = _canonical_config(source_config) + if target_form == "canonical": + target_config = _canonical_config(target_config) + original_source = deepcopy(source_config) + original = deepcopy(target_config) + source = _pretrained("pt_expt", source_config) + config, _ = FinetuneRuleBuilder( + source_config, target_config, change_model_params=False + ).build() + learned_config = config["models"][-1] if target_form == "canonical" else config + assert ( + learned_config["fitting_net"]["neuron"] == _learned(source).fitting_net.neuron + ) + assert target_config == original + assert source_config == original_source + target = _build("pt_expt", config) + target.load_state_dict(get_native_spin_finetune_state(target, source)) + inputs = _inputs() + spin = inputs.pop("spin") + expected = source(**inputs) + actual = target(**inputs, spin=spin) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(actual[key], expected[key], rtol=1e-11, atol=1e-11) + + +@pytest.mark.parametrize("explicit_width", [False, True]) +@pytest.mark.parametrize("source_auto", [False, True]) +def test_dpa4c_width_aliases_preserve_explicit_config( + explicit_width: bool, source_auto: bool +) -> None: + source_config = _config("dpa4c", False) + source_config["fitting_net"].update( + type="dpa4_ener", neuron=[0] if source_auto else [64] + ) + source = _build("pt_expt", source_config) + source_script = deepcopy(source_config) + source_script["descriptor"]["type"] = "DPA4C" + source_script["fitting_net"]["n_neuron"] = source_script["fitting_net"].pop( + "neuron" + ) + target = deepcopy(source_script) + target["spin"] = {"scheme": "native", "use_spin": [True, False]} + requested = [64] if explicit_width else [0] + target["fitting_net"]["n_neuron"] = requested + original = deepcopy(target) + resolved, _ = FinetuneRuleBuilder( + source_script, target, change_model_params=False + ).build() + config = normalize( + update_deepmd_input(_make_config("unused", resolved), warning=False) + ) + model_config = config["model"] + expected = requested if explicit_width else _learned(source).fitting_net.neuron + assert model_config["fitting_net"]["neuron"] == expected + assert "n_neuron" not in model_config["fitting_net"] + assert target == original + model = _build("pt_expt", model_config) + if expected == _learned(source).fitting_net.neuron: + model.load_state_dict(get_native_spin_finetune_state(model, source)) + + +@pytest.mark.parametrize("resnet_dt", [False, True]) +@pytest.mark.parametrize("residual_side", ["source", "target"]) +@pytest.mark.parametrize("width_multiple", [1, 2]) +def test_dpa4c_transfer_rejects_changed_input_residual( + resnet_dt: bool, residual_side: str, width_multiple: int +) -> None: + source_config = _config("dpa4c", False) + target_config = _config("dpa4c", True) + width_config = source_config if residual_side == "source" else target_config + width = _learned(_build("pt_expt", width_config)).descriptor.get_dim_out() + for config in (source_config, target_config): + config["fitting_net"].update( + neuron=[width_multiple * width], resnet_dt=resnet_dt + ) + source = _pretrained("pt_expt", source_config) + target = _build("pt_expt", target_config) + initial = {key: value.clone() for key, value in target.state_dict().items()} + with pytest.raises(ValueError, match=r"residual.*input"): + get_native_spin_finetune_state(target, source) + for key, value in target.state_dict().items(): + torch.testing.assert_close(value, initial[key], rtol=0, atol=0) + + # A random fitting net does not inherit the source function or its residual. + state = get_native_spin_finetune_state(target, source, random_fitting=True) + target.load_state_dict(state) + for key, value in state.items(): + if ".fitting_net." in key: + torch.testing.assert_close(value, initial[key], rtol=0, atol=0) + + +@pytest.mark.parametrize("fitting_type", ["ener", "dpa4_ener", "sezm_ener"]) +@pytest.mark.parametrize("neuron", [[], [8], [8, 8]]) +def test_dpa4c_transfer_fitting_input_projection( + fitting_type: str, neuron: list[int] +) -> None: + source_config = _config("dpa4c", False) + source_config["fitting_net"].update(type=fitting_type, neuron=neuron) + target_config = deepcopy(source_config) + target_config["spin"] = {"scheme": "native", "use_spin": [True, False]} + source = _pretrained("pt_expt", source_config) + target = _build("pt_expt", target_config) + # DPA4C's descriptor selects a shared fitting network at model construction. + assert _learned(source).fitting_net.nets.ndim == 0 + assert _learned(target).fitting_net.nets.ndim == 0 + state = get_native_spin_finetune_state(target, source) + assert all(not value.requires_grad for value in state.values()) + target.load_state_dict(state) + inputs = _inputs() + spin = inputs.pop("spin") + expected = source(**inputs, do_atomic_virial=True) + assert expected["force"].abs().max() > 1e-8 + for scale in (0.0, 1.0, 2.0): + actual = target(**inputs, spin=scale * spin, do_atomic_virial=True) + for key in ("energy", "force", "virial", "atom_virial"): + torch.testing.assert_close( + actual[key], expected[key], rtol=1e-11, atol=1e-11 + ) + assert torch.count_nonzero(actual["force_mag"]) == 0 + + (actual["force_mag"] - 0.2).square().sum().backward() + gate = _learned(target).descriptor.spin.spin_gate + assert gate.grad is not None and gate.grad.abs().max() > 1e-8 + with torch.no_grad(): + gate.add_(-1e-3 * gate.grad) + trained = target(**inputs, spin=spin) + assert trained["force_mag"].abs().max() > 1e-10 + assert torch.count_nonzero(trained["force_mag"][:, 1]) == 0 + + restored = type(target).deserialize(target.serialize()).to(DEVICE) + reloaded = restored(**inputs, spin=spin) + for key in ("energy", "force", "virial", "force_mag"): + torch.testing.assert_close(reloaded[key], trained[key], rtol=1e-11, atol=1e-11) + + +@pytest.mark.parametrize( + "backend, descriptor, source_type, target_type", + [ + ("pt", "dpa4", "property", "dpa4_ener"), + ("pt_expt", "dpa4", "ener", "dpa4_ener"), + ("pt_expt", "dpa4c", "ener", "dpa4_ener"), + ("pt_expt", "dpa4c", "dpa4_ener", "ener"), + ], +) +def test_spin_free_transfer_rejects_changed_fitting_family( + backend: str, descriptor: str, source_type: str, target_type: str +) -> None: + source_config = _config(descriptor, False) + target_config = _config(descriptor, True) + if backend == "pt_expt": + source_config["type"] = target_config["type"] = "standard" + for config, fitting_type in ( + (source_config, source_type), + (target_config, target_type), + ): + config["fitting_net"].update( + type=fitting_type, neuron=[16 if fitting_type == "ener" else 8] + ) + if fitting_type == "property": + config["fitting_net"].update(task_dim=1, property_name="energy") + source = _build(backend, source_config) + target = _build(backend, target_config) + initial = {key: value.clone() for key, value in target.state_dict().items()} + with pytest.raises(ValueError, match="matching fitting network types"): + get_native_spin_finetune_state(target, source) + for key, value in target.state_dict().items(): + torch.testing.assert_close(value, initial[key], rtol=0, atol=0) + target.load_state_dict( + get_native_spin_finetune_state(target, source, random_fitting=True) + ) + for key, value in target.state_dict().items(): + if ".fitting_net." in key: + torch.testing.assert_close(value, initial[key], rtol=0, atol=0) + + +@pytest.mark.parametrize( + "backend, descriptor, source_neuron, target_neuron, trainable", + [ + ("pt", "dpa4", [8], [8, 8], True), + ("pt", "dpa4", [8, 8], [8], True), + ("pt_expt", "dpa4", [8], [8, 8], True), + ("pt_expt", "dpa4c", [8], [8, 8], True), + ("pt_expt", "dpa4c", [8, 8], [8], True), + ("pt_expt", "dpa4c", [8], [8, 8], False), + ], +) +def test_spin_free_transfer_rejects_changed_fitting_layers( + backend: str, + descriptor: str, + source_neuron: list[int], + target_neuron: list[int], + trainable: bool, +) -> None: + source_config = _config(descriptor, False) + target_config = _config(descriptor, True) + for config, neuron in ( + (source_config, source_neuron), + (target_config, target_neuron), + ): + config["fitting_net"].update( + type="dpa4_ener", neuron=neuron, trainable=trainable + ) + source = _build(backend, source_config) + target = _build(backend, target_config) + initial = {key: value.clone() for key, value in target.state_dict().items()} + with pytest.raises(ValueError, match="matching fitting network state keys"): + get_native_spin_finetune_state(target, source) + for key, value in target.state_dict().items(): + torch.testing.assert_close(value, initial[key], rtol=0, atol=0) + target.load_state_dict( + get_native_spin_finetune_state(target, source, random_fitting=True) + ) + for key, value in target.state_dict().items(): + if ".fitting_net." in key: + torch.testing.assert_close(value, initial[key], rtol=0, atol=0) + + +def test_spin_free_transfer_preserves_target_default_fparam() -> None: + source_config = _config("dpa4", False) + source_config["fitting_net"]["numb_fparam"] = 1 + target_config = deepcopy(source_config) + target_config["spin"] = {"scheme": "native", "use_spin": [True, False]} + target_config["fitting_net"]["default_fparam"] = [0.25] + source = _build("pt", source_config) + target = _build("pt", target_config) + fitting = _learned(target).fitting_net + expected = fitting.default_fparam_tensor.clone() + target.load_state_dict(get_native_spin_finetune_state(target, source)) + torch.testing.assert_close(fitting.default_fparam_tensor, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("fitting_type", ["ener", "dpa4_ener"]) +def test_dpa4c_spin_transfer_preserves_fitting_parameter_tail( + fitting_type: str, +) -> None: + source_config = _config("dpa4c", False) + source_config["fitting_net"].update(type=fitting_type, numb_fparam=2, numb_aparam=1) + target_config = deepcopy(source_config) + target_config["spin"] = {"scheme": "native", "use_spin": [True, False]} + source = _pretrained("pt_expt", source_config) + target = _build("pt_expt", target_config) + target.load_state_dict(get_native_spin_finetune_state(target, source)) + inputs = _inputs() + spin = inputs.pop("spin") + inputs["fparam"] = torch.tensor([[0.2, -0.1]], dtype=torch.float64, device=DEVICE) + inputs["aparam"] = torch.tensor( + [[[0.3], [0.5], [0.7]]], dtype=torch.float64, device=DEVICE + ) + expected = source(**inputs) + actual = target(**inputs, spin=spin) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(actual[key], expected[key], rtol=1e-11, atol=1e-11) + + +@pytest.mark.parametrize("grade", ["nano", "mini", "neo", "air", "plus"]) +def test_dpa4c_preset_spin_free_transfer(grade: str) -> None: + source_config = get_model_preset(f"dpa4c-{grade}-v20260911") + source_config["descriptor"]["seed"] = 17 + source_config["fitting_net"]["seed"] = 19 + target_config = deepcopy(source_config) + target_config["spin"] = {"scheme": "native", "use_spin": ["Fe"]} + source = _pretrained("pt_expt", source_config) + target = _build("pt_expt", target_config) + target.load_state_dict(get_native_spin_finetune_state(target, source)) + inputs = _inputs() + spin = inputs.pop("spin") + inputs["atype"] = torch.tensor( + [[source.get_type_map().index(name) for name in ("Fe", "C", "Fe")]], + device=DEVICE, + ) + expected = source(**inputs) + actual = target(**inputs, spin=spin) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(actual[key], expected[key], rtol=2e-5, atol=2e-5) + assert torch.count_nonzero(actual["force_mag"]) == 0 + + +@pytest.mark.parametrize( + "backend, descriptor, fitting_type", + [ + ("pt", "dpa4", "dpa4_ener"), + ("pt_expt", "dpa4", "dpa4_ener"), + ("pt_expt", "dpa4c", "ener"), + ("pt_expt", "dpa4c", "dpa4_ener"), + ], +) +@pytest.mark.parametrize("bridged", [False, True]) +@pytest.mark.parametrize("multitask", [False, True]) +def test_spin_free_native_finetune_cli( + backend: str, + descriptor: str, + fitting_type: str, + bridged: bool, + multitask: bool, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run_spin_free_native_finetune_cli( + backend, + descriptor, + fitting_type, + bridged, + multitask, + "transfer", + tmp_path, + monkeypatch, + ) + + +@pytest.mark.parametrize("backend, descriptor", [("pt", "dpa4"), ("pt_expt", "dpa4c")]) +@pytest.mark.parametrize("mode", ["random", "pretrain_script"]) +@pytest.mark.parametrize("multitask", [False, True]) +def test_spin_free_native_finetune_cli_options( + backend: str, + descriptor: str, + mode: str, + multitask: bool, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run_spin_free_native_finetune_cli( + backend, descriptor, "dpa4_ener", False, multitask, mode, tmp_path, monkeypatch + ) + + +@pytest.mark.parametrize("mode", ["transfer", "random", "pretrain_script"]) +@pytest.mark.parametrize("multitask", [False, True]) +def test_dpa4c_auto_width_finetune_cli( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mode: str, multitask: bool +) -> None: + _run_spin_free_native_finetune_cli( + "pt_expt", + "dpa4c", + "dpa4_ener", + False, + multitask, + mode, + tmp_path, + monkeypatch, + auto_width=True, + ) + + +def test_dpa4c_canonical_auto_width_finetune_cli( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _run_spin_free_native_finetune_cli( + "pt_expt", + "dpa4c", + "dpa4_ener", + True, + False, + "transfer", + tmp_path, + monkeypatch, + auto_width=True, + canonical=True, + ) + + +def _run_spin_free_native_finetune_cli( + backend: str, + descriptor: str, + fitting_type: str, + bridged: bool, + multitask: bool, + mode: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + *, + auto_width: bool = False, + canonical: bool = False, +) -> None: + source_config = _config(descriptor, False) + source_config["fitting_net"]["type"] = fitting_type + if auto_width: + source_config["fitting_net"].pop("neuron") + source = _pretrained(backend, source_config) + wrapper_cls = PTModelWrapper if backend == "pt" else ModelWrapper + wrapper = wrapper_cls( + {"magnetic": source} if multitask else source, + model_params={"model_dict": {"magnetic": source_config}} + if multitask + else source_config, + ) + wrapper.train_infos["step"] = 100 + checkpoint = tmp_path / "pretrained.pt" + torch.save( + { + "model": wrapper.state_dict(), + "optimizer": {"source_only": True}, + "ema": {"source_only": True}, + }, + checkpoint, + ) + + inputs = _inputs() + spin = inputs.pop("spin") + expected = source(**inputs) + system = tmp_path / "data" / "data_0" + arrays = system / "set.000" + arrays.mkdir(parents=True) + (system / "type.raw").write_text("0\n1\n0\n") + (system / "type_map.raw").write_text("Fe\nC\n") + labels = { + "coord": inputs["coord"], + "box": inputs["box"], + "spin": spin, + "energy": expected["energy"] + 1.75, + "force": expected["force"], + "force_mag": torch.full_like(spin, 0.2), + } + for key, value in labels.items(): + data = value.detach().cpu().numpy().reshape(1, -1) + np.save(arrays / f"{key}.npy", np.repeat(data, 2, axis=0)) + + target_config = _config(descriptor, True, bridged) + target_config["fitting_net"]["type"] = fitting_type + if auto_width: + target_config["fitting_net"].pop("neuron") + options = [] + if mode == "random": + initial_fitting = dict( + _learned(_build(backend, target_config)).fitting_net.named_parameters() + ) + if multitask: + target_config["finetune_head"] = "RANDOM" + else: + options = ["--model-branch", "RANDOM"] + elif mode == "pretrain_script": + target_config["descriptor"]["channels"] *= 2 + options = ["--use-pretrain-script"] + if bridged: + target_config["preset_out_bias"] = {"energy": {"Fe": 1.25, "C": -0.75}} + if canonical: + target_config = _canonical_config(target_config) + # A zero-step run exposes RANDOM initialization before optimization. + steps = 0 if mode == "random" else 2 + config = _make_config(str(system.parent), target_config, numb_steps=steps) + config["loss"] = { + "type": "ener_spin", + "start_pref_e": 1.0, + "limit_pref_e": 1.0, + "start_pref_fr": 1.0, + "limit_pref_fr": 1.0, + "start_pref_fm": 1.0, + "limit_pref_fm": 1.0, + } + config["training"].update(enable_ema=True, ema_decay=0.9) + if multitask: + config["model"] = {"model_dict": {"magnetic": target_config}} + config["loss_dict"] = {"magnetic": config.pop("loss")} + config["training"]["data_dict"] = { + "magnetic": { + key: config["training"].pop(key) + for key in ("training_data", "validation_data") + } + } + input_path = tmp_path / "input.json" + input_path.write_text(json.dumps(config)) + monkeypatch.chdir(tmp_path) + entrypoint = pt_main if backend == "pt" else main + trainers = [] + if mode == "random": + entrypoints = import_module(f"deepmd.{backend}.entrypoints.main") + get_trainer = entrypoints.get_trainer + + def capture_trainer(*args: Any, **kwargs: Any) -> Any: + trainer = get_trainer(*args, **kwargs) + trainers.append(trainer) + return trainer + + monkeypatch.setattr(entrypoints, "get_trainer", capture_trainer) + entrypoint( + [ + "train", + str(input_path), + "--finetune", + str(checkpoint), + "--skip-neighbor-stat", + *options, + ] + ) + if mode == "random": + assert len(trainers) == 1 + model = trainers[0].model + trained = model["magnetic"] if multitask else model + for key, value in _learned(trained).fitting_net.named_parameters(): + torch.testing.assert_close(value, initial_fitting[key], rtol=0, atol=0) + assert torch.count_nonzero(trained(**inputs, spin=spin)["force_mag"]) == 0 + return + saved = torch.load( + tmp_path / "model.ckpt.pt", map_location=DEVICE, weights_only=True + ) + params = saved["model"]["_extra_state"]["model_params"] + model_params = params["model_dict"]["magnetic"] if multitask else params + learned_params = model_params["models"][-1] if canonical else model_params + if auto_width: + assert ( + learned_params["fitting_net"]["neuron"] + == _learned(source).fitting_net.neuron + ) + assert model_params["spin"]["scheme"] == "native" + assert model_params["spin"]["use_spin"] == [True, False] + assert ( + learned_params["descriptor"]["channels"] + == source_config["descriptor"]["channels"] + ) + # PT checkpoints record the zero-based step; pt_expt records the step count. + assert saved["model"]["_extra_state"]["train_infos"]["step"] == ( + max(steps - 1, 0) if backend == "pt" else steps + ) + assert "source_only" not in saved["optimizer"] + assert "source_only" not in saved["ema"] + trained = _build(backend, model_params) + loaded = wrapper_cls( + {"magnetic": trained} if multitask else trained, model_params=params + ) + loaded.load_state_dict(saved["model"]) + if descriptor == "dpa4c": + reference = _learned(trained).descriptor.spin.spin_reference + expected_reference = spin[:, [0, 2]].square().sum(dim=-1).mean().sqrt() + torch.testing.assert_close(reference[0], expected_reference, rtol=1e-12, atol=0) + bias = trained.get_out_bias().reshape(-1) + if bridged: + torch.testing.assert_close( + bias, + torch.tensor([1.25, -0.75], dtype=torch.float64, device=DEVICE), + rtol=0, + atol=0, + ) + else: + torch.testing.assert_close( + 2.0 * bias[0] + bias[1], + torch.tensor(1.75, dtype=torch.float64, device=DEVICE), + rtol=1e-11, + atol=1e-11, + ) + assert trained(**inputs, spin=spin)["force_mag"].abs().max() > 1e-10