diff --git a/deepmd/dpmodel/utils/multi_task.py b/deepmd/dpmodel/utils/multi_task.py new file mode 100644 index 0000000000..be2b962193 --- /dev/null +++ b/deepmd/dpmodel/utils/multi_task.py @@ -0,0 +1,342 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-neutral helpers for multi-task shared-parameter wiring.""" + +import logging +from collections.abc import ( + Callable, +) +from copy import ( + deepcopy, +) +from typing import ( + Any, +) + +log = logging.getLogger(__name__) + + +def cascade_top_level_defaults(model_config: dict[str, Any]) -> None: + """Lower model-wide entries into each multi-task branch in-place.""" + reserved_top_level = ("model_dict", "shared_dict") + top_level_defaults = { + k: deepcopy(v) for k, v in model_config.items() if k not in reserved_top_level + } + for branch in model_config["model_dict"].values(): + for k, v in top_level_defaults.items(): + branch.setdefault(k, deepcopy(v)) + for k in top_level_defaults: + model_config.pop(k, None) + + +def preprocess_shared_params( + model_config: dict[str, Any], + get_class_name: Callable[[str, dict[str, Any]], type], + *, + require_shared_type_map: bool = True, + cascade_defaults: bool = False, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Expand ``shared_dict`` references and generate runtime sharing links. + + Parameters + ---------- + model_config : dict[str, Any] + Multi-task model configuration containing ``model_dict`` and an optional + ``shared_dict``. String references to shared descriptors, fitting + networks, and type maps are replaced in-place by their configurations. + get_class_name : Callable[[str, dict[str, Any]], type] + Backend-specific callback that resolves the class for a shared + descriptor or fitting-network configuration. + require_shared_type_map : bool, default=True + Whether exactly one shared type-map reference is required. If false, at + most one shared type-map reference is allowed. + cascade_defaults : bool, default=False + Whether to copy non-reserved top-level model options into each model + branch that does not define them, then remove those options from the top + level. + + Returns + ------- + model_config : dict[str, Any] + The input configuration with shared references expanded. + shared_links : dict[str, Any] + Sharing metadata keyed by entries in ``shared_dict``. Each entry records + the resolved class and the model components linked to it, sorted by + sharing level. + """ + assert "model_dict" in model_config, "only multi-task model can use this method!" + if cascade_defaults: + cascade_top_level_defaults(model_config) + + supported_types = ["type_map", "descriptor", "fitting_net"] + shared_dict = model_config.get("shared_dict", {}) + shared_links: dict[str, Any] = {} + type_map_keys: list[str] = [] + + def replace_one_item( + params_dict: dict[str, Any] | list[Any], + model_key: str, + key_type: str, + key_in_dict: str, + suffix: str = "", + index: int | None = None, + ) -> None: + shared_type = key_type + shared_key = key_in_dict + shared_level = 0 + if ":" in key_in_dict: + shared_key = key_in_dict.split(":")[0] + shared_level = int(key_in_dict.split(":")[1]) + assert shared_key in shared_dict, ( + f"Appointed {shared_type} {shared_key} are not in the shared_dict! " + "Please check the input params." + ) + if index is None: + assert isinstance(params_dict, dict) + params_dict[shared_type] = deepcopy(shared_dict[shared_key]) + else: + params_dict[index] = deepcopy(shared_dict[shared_key]) + if shared_type == "type_map": + if key_in_dict not in type_map_keys: + type_map_keys.append(key_in_dict) + else: + if shared_key not in shared_links: + shared_links[shared_key] = { + "type": get_class_name(shared_type, shared_dict[shared_key]), + "links": [], + } + shared_links[shared_key]["links"].append( + { + "model_key": model_key, + "shared_type": shared_type + suffix, + "shared_level": shared_level, + } + ) + + for model_key in model_config["model_dict"]: + model_params_item = model_config["model_dict"][model_key] + for item_key in model_params_item: + if item_key not in supported_types: + continue + item_params = model_params_item[item_key] + if isinstance(item_params, str): + replace_one_item(model_params_item, model_key, item_key, item_params) + elif ( + isinstance(item_params, dict) + and item_params.get("type", "") == "hybrid" + ): + for ii, hybrid_item in enumerate(item_params["list"]): + if isinstance(hybrid_item, str): + replace_one_item( + model_params_item[item_key]["list"], + model_key, + item_key, + hybrid_item, + suffix=f"_hybrid_{ii}", + index=ii, + ) + + for shared_key in shared_links: + shared_links[shared_key]["links"] = sorted( + shared_links[shared_key]["links"], + key=lambda x: ( + x["shared_level"] + - ("spin" in model_config["model_dict"][x["model_key"]]) * 100 + ), + ) + if require_shared_type_map: + assert len(type_map_keys) == 1, "Multitask model must have only one type_map!" + else: + assert len(type_map_keys) <= 1, "Shared params must have at most one type_map!" + return model_config, shared_links + + +def apply_shared_links( + models: dict[str, Any], + shared_links: dict[str, Any] | None, + *, + share_descriptor: Callable[..., None], + share_fitting: Callable[..., None], + model_key_prob_map: dict[str, float] | None = None, + data_stat_protect: float = 1e-2, + resume: bool = False, + get_descriptor: Callable[[Any, str], Any] | None = None, + get_fitting: Callable[[Any, str], Any | None] | None = None, + logger: logging.Logger | None = None, +) -> None: + """Apply sharing links to constructed models with backend-specific hooks. + + Parameters + ---------- + models : dict[str, Any] + Constructed models keyed by the model names used in ``shared_links``. + shared_links : dict[str, Any] or None + Sharing metadata generated by :func:`preprocess_shared_params`. No work + is performed when it is empty or ``None``. + share_descriptor : Callable[..., None] + Backend-specific callback that shares a linked descriptor with its base + descriptor. + share_fitting : Callable[..., None] + Backend-specific callback that shares a linked fitting network with its + base fitting network. + model_key_prob_map : dict[str, float] or None, default=None + Sampling probabilities keyed by model name. The linked-to-base + probability ratio is passed to the sharing callbacks to weight merged + statistics. Equal probabilities are used when omitted. + data_stat_protect : float, default=1e-2 + Protection value passed to the fitting-network sharing callback when + statistics are merged. + resume : bool, default=False + Whether sharing is being restored from a checkpoint. Backend callbacks + use this flag to skip merging statistics again. + get_descriptor : Callable[[Any, str], Any] or None, default=None + Optional accessor for a descriptor or hybrid sub-descriptor. The common + model accessor is used when omitted. + get_fitting : Callable[[Any, str], Any or None] or None, default=None + Optional accessor for a fitting component. The common atomic-model + accessor is used when omitted. + logger : logging.Logger or None, default=None + Logger used to report parameter-sharing operations. + + Returns + ------- + None + The components in ``models`` are shared in-place. + """ + if not shared_links: + return + if model_key_prob_map is None: + model_key_prob_map = dict.fromkeys(models, 1.0) + if get_descriptor is None: + get_descriptor = get_descriptor_component + if get_fitting is None: + get_fitting = get_atomic_model_attr + if logger is None: + logger = log + + for shared_info in shared_links.values(): + links = shared_info.get("links", []) + if len(links) < 2: + continue + shared_base = links[0] + class_type_base = shared_base["shared_type"] + model_key_base = shared_base["model_key"] + shared_level_base = int(shared_base["shared_level"]) + + if "descriptor" in class_type_base: + base_class = get_descriptor(models[model_key_base], class_type_base) + for link_item in links[1:]: + class_type_link = link_item["shared_type"] + model_key_link = link_item["model_key"] + shared_level_link = int(link_item["shared_level"]) + assert shared_level_link >= shared_level_base, ( + "The shared_links must be sorted by shared_level!" + ) + assert "descriptor" in class_type_link, ( + f"Class type mismatched: {class_type_base} vs {class_type_link}!" + ) + link_model = models[model_key_link] + link_class = get_descriptor(link_model, class_type_link) + share_descriptor( + link_model, + class_type_link, + link_class, + base_class, + shared_level_link, + _model_prob_ratio( + model_key_prob_map, model_key_base, model_key_link + ), + resume=resume, + ) + logger.warning( + "Shared params of %s.%s and %s.%s!", + model_key_base, + class_type_base, + model_key_link, + class_type_link, + ) + continue + + base_class = get_fitting(models[model_key_base], class_type_base) + if base_class is None: + continue + for link_item in links[1:]: + class_type_link = link_item["shared_type"] + model_key_link = link_item["model_key"] + shared_level_link = int(link_item["shared_level"]) + assert shared_level_link >= shared_level_base, ( + "The shared_links must be sorted by shared_level!" + ) + assert class_type_base == class_type_link, ( + f"Class type mismatched: {class_type_base} vs {class_type_link}!" + ) + link_class = get_fitting(models[model_key_link], class_type_link) + if link_class is None: + continue + share_fitting( + link_class, + base_class, + shared_level_link, + _model_prob_ratio(model_key_prob_map, model_key_base, model_key_link), + protection=data_stat_protect, + resume=resume, + ) + logger.warning( + "Shared params of %s.%s and %s.%s!", + model_key_base, + class_type_base, + model_key_link, + class_type_link, + ) + + +def get_descriptor_component(model: Any, shared_type: str) -> Any: + """Get a model descriptor or a hybrid sub-descriptor by shared type.""" + if shared_type == "descriptor": + return model.get_descriptor() + if "hybrid" in shared_type: + hybrid_index = int(shared_type.split("_")[-1]) + return model.get_descriptor().descrpt_list[hybrid_index] + raise RuntimeError(f"Unknown class_type {shared_type}!") + + +def set_descriptor_component(model: Any, shared_type: str, descriptor: Any) -> None: + """Set a model descriptor or a hybrid sub-descriptor by shared type.""" + if shared_type == "descriptor": + model.atomic_model.descriptor = descriptor + return + if "hybrid" in shared_type: + hybrid_index = int(shared_type.split("_")[-1]) + model.get_descriptor().descrpt_list[hybrid_index] = descriptor + return + raise RuntimeError(f"Unknown class_type {shared_type}!") + + +def get_atomic_model_attr(model: Any, attr: str) -> Any | None: + """Get an attribute from ``model.atomic_model`` if it exists.""" + if hasattr(model.atomic_model, attr): + return getattr(model.atomic_model, attr) + return None + + +def sanitize_shared_links(shared_links: dict[str, Any] | None) -> dict[str, Any] | None: + """Return a JSON-safe copy of ``shared_links``.""" + if shared_links is None: + return None + sanitized: dict[str, Any] = {} + for shared_key, shared_info in shared_links.items(): + class_type = shared_info.get("type") + sanitized[shared_key] = { + "type": getattr(class_type, "__name__", str(class_type)), + "links": deepcopy(shared_info.get("links", [])), + } + return sanitized + + +def _model_prob_ratio( + model_key_prob_map: dict[str, float], + model_key_base: str, + model_key_link: str, +) -> float: + return float(model_key_prob_map[model_key_link]) / float( + model_key_prob_map[model_key_base] + ) diff --git a/deepmd/jax/entrypoints/train.py b/deepmd/jax/entrypoints/train.py index fdc695548d..61d415a340 100644 --- a/deepmd/jax/entrypoints/train.py +++ b/deepmd/jax/entrypoints/train.py @@ -85,6 +85,7 @@ class JAXTrainEntrypoint(AbstractTrainEntrypoint): def __init__(self) -> None: self.finetune_links: dict[str, Any] | None = None + self.shared_links: dict[str, Any] | None = None def validate_options( self, @@ -96,10 +97,6 @@ def validate_options( raise NotImplementedError( "JAX training does not support init_frz_model yet" ) - if self.is_multi_task(config) and config["model"].get("shared_dict"): - raise NotImplementedError( - "JAX multi-task training does not support shared_dict yet" - ) def preprocess_config( self, @@ -108,6 +105,8 @@ def preprocess_config( ) -> dict[str, Any]: """Apply JAX fine-tuning and pretrained-script preprocessing.""" self.finetune_links = None + self.shared_links = None + if options.finetune is not None: from deepmd.jax.utils.finetune import ( get_finetune_rules, @@ -122,6 +121,17 @@ def preprocess_config( elif options.init_model is not None and options.use_pretrain_script: model_data = serialize_from_file(options.init_model) config["model"] = model_data["model_def_script"] + if self.is_multi_task(config): + if config["model"].get("shared_dict"): + from deepmd.jax.utils.multi_task import ( + preprocess_shared_params, + ) + + config["model"], self.shared_links = preprocess_shared_params( + config["model"] + ) + if "RANDOM" in config["model"]["model_dict"]: + raise ValueError("Model name can not be 'RANDOM' in multi-task mode!") return config def update_neighbor_stat( @@ -154,6 +164,7 @@ def run_training( restart=options.restart, finetune_model=options.finetune, finetune_links=self.finetune_links, + shared_links=self.shared_links, ) if neighbor_stat is not None: model.set_min_nbor_dist(neighbor_stat) diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index c77bc944b5..684befada1 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -48,6 +48,10 @@ from deepmd.dpmodel.utils.learning_rate import ( LearningRateExp, ) +from deepmd.dpmodel.utils.multi_task import ( + apply_shared_links, + set_descriptor_component, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -70,6 +74,9 @@ from deepmd.jax.model.model import ( get_model, ) +from deepmd.jax.utils.multi_task import ( + preprocess_shared_params, +) from deepmd.jax.utils.serialization import ( serialize_from_file, ) @@ -102,6 +109,7 @@ def __init__( restart: str | None = None, finetune_model: str | None = None, finetune_links: dict[str, Any] | None = None, + shared_links: dict[str, Any] | None = None, ) -> None: """Initialize the trainer from input data and optional checkpoints.""" if finetune_model is not None and ( @@ -114,6 +122,7 @@ def __init__( self.restart = restart self.finetune_model = finetune_model self.finetune_links = finetune_links + self.shared_links = shared_links self.restart_training = restart is not None self.training_param = jdata["training"] self.validating_param = jdata.get("validating", {}) or {} @@ -122,6 +131,14 @@ def __init__( self.model_def_script = jdata["model"] self.multi_task = "model_dict" in self.model_def_script + if ( + self.multi_task + and self.shared_links is None + and self.model_def_script.get("shared_dict") + ): + self.model_def_script, self.shared_links = preprocess_shared_params( + self.model_def_script + ) self.model_keys = ( list(self.model_def_script["model_dict"]) if self.multi_task @@ -460,6 +477,12 @@ def _setup_training( if self.finetune_model is not None: self._apply_finetune() + self._share_model_params( + resume=self.init_model is not None + or self.restart is not None + or self.finetune_model is not None + ) + for model_key in self.model_keys: tx = optax.chain( optax.scale_by_adam(), @@ -544,6 +567,34 @@ def _apply_finetune(self) -> None: bias_adjust_mode=bias_mode, ) + def _share_model_params(self, *, resume: bool = False) -> None: + """Apply multi-task shared_dict links to JAX model branches.""" + if not self.multi_task or not self.shared_links: + return + data_stat_protect = np.array( + [ + self.model_params_by_task[model_key].get("data_stat_protect", 1e-2) + for model_key in self.model_keys + ] + ) + if not np.allclose(data_stat_protect, data_stat_protect[0]): + raise ValueError( + "Model key 'data_stat_protect' must be the same in each branch when multitask!" + ) + if self.model_prob is None: + model_prob = np.ones(len(self.model_keys), dtype=float) / len( + self.model_keys + ) + else: + model_prob = self.model_prob + share_jax_model_params( + self.models, + self.shared_links, + model_key_prob_map=dict(zip(self.model_keys, model_prob, strict=True)), + data_stat_protect=float(data_stat_protect[0]), + resume=resume, + ) + def _warn_finetune_config_mismatch( self, model_key: str, @@ -893,6 +944,236 @@ def _evaluate_model_dict( return model_dict +def share_jax_model_params( + models: dict[str, BaseModel], + shared_links: dict[str, Any], + *, + model_key_prob_map: dict[str, float], + data_stat_protect: float = 1e-2, + resume: bool = False, +) -> None: + """Share JAX model parameters following ``preprocess_shared_params`` links.""" + apply_shared_links( + models, + shared_links, + share_descriptor=_share_jax_descriptor, + share_fitting=_share_jax_fitting, + model_key_prob_map=model_key_prob_map, + data_stat_protect=data_stat_protect, + resume=resume, + logger=log, + ) + + +def _share_jax_descriptor( + link_model: BaseModel, + link_type: str, + link_class: Any, + base_class: Any, + shared_level: int, + model_prob: float, + *, + resume: bool, +) -> None: + _share_descriptor_component( + base_class, + link_class, + shared_level, + model_prob=model_prob, + resume=resume, + ) + if shared_level == 0: + set_descriptor_component(link_model, link_type, base_class) + + +def _share_jax_fitting( + link_class: Any, + base_class: Any, + shared_level: int, + model_prob: float, + *, + protection: float, + resume: bool, +) -> None: + _share_fitting_component( + base_class, + link_class, + shared_level, + model_prob=model_prob, + protection=protection, + resume=resume, + ) + + +def _share_descriptor_component( + base_class: Any, + link_class: Any, + shared_level: int, + *, + model_prob: float, + resume: bool, +) -> None: + if type(link_class) is not type(base_class): + raise AssertionError("Only descriptors of the same type can share params!") + if shared_level == 0: + if not resume: + _merge_descriptor_stats(base_class, link_class, model_prob) + return + if shared_level == 1 and hasattr(base_class, "type_embedding"): + link_class.type_embedding = base_class.type_embedding + return + raise NotImplementedError( + f"JAX shared_dict does not support descriptor shared_level={shared_level} " + f"for {type(base_class).__name__}." + ) + + +def _merge_descriptor_stats( + base_class: Any, + link_class: Any, + model_prob: float, +) -> None: + from deepmd.dpmodel.utils.env_mat_stat import ( + merge_env_stat, + ) + + merge_env_stat(base_class, link_class, model_prob) + for attr in ( + "se_atten", + "seat", + "se_ttebd", + "repinit", + "repinit_three_body", + "repformers", + "repflows", + ): + if hasattr(base_class, attr) and hasattr(link_class, attr): + _merge_descriptor_stats( + getattr(base_class, attr), + getattr(link_class, attr), + model_prob, + ) + if hasattr(base_class, "descrpt_list") and hasattr(link_class, "descrpt_list"): + for base_item, link_item in zip( + base_class.descrpt_list, + link_class.descrpt_list, + strict=True, + ): + _merge_descriptor_stats(base_item, link_item, model_prob) + + +def _share_fitting_component( + base_class: Any, + link_class: Any, + shared_level: int, + *, + model_prob: float, + protection: float, + resume: bool, +) -> None: + if type(link_class) is not type(base_class): + raise AssertionError("Only fitting nets of the same type can share params!") + if shared_level != 0: + raise NotImplementedError( + f"JAX shared_dict does not support fitting_net shared_level={shared_level}." + ) + if not resume: + _merge_fitting_param_stats( + base_class, + link_class, + model_prob=model_prob, + protection=protection, + ) + link_class.nets = base_class.nets + for attr in ( + "fparam_avg", + "fparam_inv_std", + "aparam_avg", + "aparam_inv_std", + "default_fparam_tensor", + ): + if getattr(base_class, attr, None) is not None: + setattr(link_class, attr, getattr(base_class, attr)) + + +def _merge_fitting_param_stats( + base_class: Any, + link_class: Any, + *, + model_prob: float, + protection: float, +) -> None: + _merge_one_fitting_stat( + base_class, + link_class, + name="fparam", + avg_attr="fparam_avg", + inv_std_attr="fparam_inv_std", + numb_attr="numb_fparam", + model_prob=model_prob, + protection=protection, + ) + _merge_one_fitting_stat( + base_class, + link_class, + name="aparam", + avg_attr="aparam_avg", + inv_std_attr="aparam_inv_std", + numb_attr="numb_aparam", + model_prob=model_prob, + protection=protection, + ) + + +def _merge_one_fitting_stat( + base_class: Any, + link_class: Any, + *, + name: str, + avg_attr: str, + inv_std_attr: str, + numb_attr: str, + model_prob: float, + protection: float, +) -> None: + if getattr(base_class, numb_attr, 0) <= 0: + return + base_stats = base_class.get_param_stats().get(name, []) + link_stats = link_class.get_param_stats().get(name, []) + if not base_stats or not link_stats: + return + if len(base_stats) != getattr(base_class, numb_attr): + raise AssertionError(f"{name} statistics length mismatch!") + merged = [ + base_stats[ii] + link_stats[ii] * model_prob + for ii in range(getattr(base_class, numb_attr)) + ] + avg = np.array([stat.compute_avg() for stat in merged], dtype=np.float64) + inv_std = 1.0 / np.array( + [stat.compute_std(protection=protection) for stat in merged], + dtype=np.float64, + ) + setattr(base_class, avg_attr, _as_backend_array(getattr(base_class, avg_attr), avg)) + setattr( + base_class, + inv_std_attr, + _as_backend_array(getattr(base_class, inv_std_attr), inv_std), + ) + base_class._param_stats[name] = merged + + +def _as_backend_array(reference: Any, value: np.ndarray) -> Any: + import array_api_compat + + ref_value = getattr(reference, "value", reference) + xp = array_api_compat.array_namespace(ref_value) + return xp.asarray( + value, + dtype=ref_value.dtype, + device=array_api_compat.device(ref_value), + ) + + def _init_empty_state(params: Any) -> optax.EmptyState: """Initialize an empty Optax state without requiring optax.init_empty_state. diff --git a/deepmd/jax/utils/multi_task.py b/deepmd/jax/utils/multi_task.py new file mode 100644 index 0000000000..05f745bb11 --- /dev/null +++ b/deepmd/jax/utils/multi_task.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from importlib import ( + import_module, +) +from typing import ( + Any, +) + +from deepmd.dpmodel.utils.multi_task import ( + preprocess_shared_params as preprocess_shared_params_common, +) +from deepmd.jax.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.jax.fitting.base_fitting import ( + BaseFitting, +) + +# Populate JAX descriptor and fitting registries before class lookup. +import_module("deepmd.jax.descriptor") +import_module("deepmd.jax.fitting.fitting") + + +def preprocess_shared_params( + model_config: dict[str, Any], + require_shared_type_map: bool = True, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Preprocess JAX shared model params and generate sharing links.""" + return preprocess_shared_params_common( + model_config, + get_class_name, + require_shared_type_map=require_shared_type_map, + cascade_defaults=True, + ) + + +def get_class_name(item_key: str, item_params: dict[str, Any]) -> type: + if item_key == "descriptor": + return BaseDescriptor.get_class_by_type(item_params.get("type", "se_e2_a")) + if item_key == "fitting_net": + return BaseFitting.get_class_by_type(item_params.get("type", "ener")) + raise RuntimeError(f"Unknown class_name type {item_key}") diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py index 46cff69515..1509924e93 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/deepmd/pt_expt/train/wrapper.py @@ -12,9 +12,50 @@ import torch +from deepmd.dpmodel.utils.multi_task import ( + apply_shared_links, +) + log = logging.getLogger(__name__) +def _share_descriptor( + link_model: torch.nn.Module, + link_type: str, + link_class: Any, + base_class: Any, + shared_level: int, + model_prob: float, + *, + resume: bool, +) -> None: + del link_model, link_type + link_class.share_params( + base_class, + shared_level, + model_prob=model_prob, + resume=resume, + ) + + +def _share_fitting( + link_class: Any, + base_class: Any, + shared_level: int, + model_prob: float, + *, + protection: float, + resume: bool, +) -> None: + link_class.share_params( + base_class, + shared_level, + model_prob=model_prob, + protection=protection, + resume=resume, + ) + + class ModelWrapper(torch.nn.Module): """Model wrapper that bundles model(s) and loss(es). @@ -83,91 +124,16 @@ def share_params( resume : bool Whether resuming from checkpoint. """ - for shared_item in shared_links: - shared_base = shared_links[shared_item]["links"][0] - class_type_base = shared_base["shared_type"] - model_key_base = shared_base["model_key"] - shared_level_base = shared_base["shared_level"] - if "descriptor" in class_type_base: - if class_type_base == "descriptor": - base_class = self.model[model_key_base].get_descriptor() - elif "hybrid" in class_type_base: - hybrid_index = int(class_type_base.split("_")[-1]) - base_class = ( - self.model[model_key_base] - .get_descriptor() - .descrpt_list[hybrid_index] - ) - else: - raise RuntimeError(f"Unknown class_type {class_type_base}!") - for link_item in shared_links[shared_item]["links"][1:]: - class_type_link = link_item["shared_type"] - model_key_link = link_item["model_key"] - shared_level_link = int(link_item["shared_level"]) - assert shared_level_link >= shared_level_base, ( - "The shared_links must be sorted by shared_level!" - ) - assert "descriptor" in class_type_link, ( - f"Class type mismatched: {class_type_base} vs {class_type_link}!" - ) - if class_type_link == "descriptor": - link_class = self.model[model_key_link].get_descriptor() - elif "hybrid" in class_type_link: - hybrid_index = int(class_type_link.split("_")[-1]) - link_class = ( - self.model[model_key_link] - .get_descriptor() - .descrpt_list[hybrid_index] - ) - else: - raise RuntimeError(f"Unknown class_type {class_type_link}!") - frac_prob = ( - model_key_prob_map[model_key_link] - / model_key_prob_map[model_key_base] - ) - link_class.share_params( - base_class, - shared_level_link, - model_prob=frac_prob, - resume=resume, - ) - log.warning( - f"Shared params of {model_key_base}.{class_type_base} " - f"and {model_key_link}.{class_type_link}!" - ) - else: - if hasattr(self.model[model_key_base].atomic_model, class_type_base): - base_class = self.model[model_key_base].atomic_model.__getattr__( - class_type_base - ) - for link_item in shared_links[shared_item]["links"][1:]: - class_type_link = link_item["shared_type"] - model_key_link = link_item["model_key"] - shared_level_link = int(link_item["shared_level"]) - assert shared_level_link >= shared_level_base, ( - "The shared_links must be sorted by shared_level!" - ) - assert class_type_base == class_type_link, ( - f"Class type mismatched: {class_type_base} vs {class_type_link}!" - ) - link_class = self.model[ - model_key_link - ].atomic_model.__getattr__(class_type_link) - frac_prob = ( - model_key_prob_map[model_key_link] - / model_key_prob_map[model_key_base] - ) - link_class.share_params( - base_class, - shared_level_link, - model_prob=frac_prob, - protection=data_stat_protect, - resume=resume, - ) - log.warning( - f"Shared params of {model_key_base}.{class_type_base} " - f"and {model_key_link}.{class_type_link}!" - ) + apply_shared_links( + self.model, + shared_links, + share_descriptor=_share_descriptor, + share_fitting=_share_fitting, + model_key_prob_map=model_key_prob_map, + data_stat_protect=data_stat_protect, + resume=resume, + logger=log, + ) def forward( self, diff --git a/deepmd/pt_expt/utils/multi_task.py b/deepmd/pt_expt/utils/multi_task.py index a4600d5ebb..54d761b35d 100644 --- a/deepmd/pt_expt/utils/multi_task.py +++ b/deepmd/pt_expt/utils/multi_task.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from copy import ( - deepcopy, -) from typing import ( Any, ) +from deepmd.dpmodel.utils.multi_task import ( + preprocess_shared_params as preprocess_shared_params_common, +) from deepmd.pt_expt.descriptor.base_descriptor import ( BaseDescriptor, ) @@ -36,75 +36,7 @@ def preprocess_shared_params( Lower for more params to share, 0 means to share all params in this item. This list are sorted by "shared_level". """ - assert "model_dict" in model_config, "only multi-task model can use this method!" - supported_types = ["type_map", "descriptor", "fitting_net"] - shared_dict = model_config.get("shared_dict", {}) - shared_links = {} - type_map_keys = [] - - def replace_one_item( - params_dict: dict[str, Any], - key_type: str, - key_in_dict: str, - suffix: str = "", - index: int | None = None, - ) -> None: - shared_type = key_type - shared_key = key_in_dict - shared_level = 0 - if ":" in key_in_dict: - shared_key = key_in_dict.split(":")[0] - shared_level = int(key_in_dict.split(":")[1]) - assert shared_key in shared_dict, ( - f"Appointed {shared_type} {shared_key} are not in the shared_dict! Please check the input params." - ) - if index is None: - params_dict[shared_type] = deepcopy(shared_dict[shared_key]) - else: - params_dict[index] = deepcopy(shared_dict[shared_key]) - if shared_type == "type_map": - if key_in_dict not in type_map_keys: - type_map_keys.append(key_in_dict) - else: - if shared_key not in shared_links: - class_name = get_class_name(shared_type, shared_dict[shared_key]) - shared_links[shared_key] = {"type": class_name, "links": []} - link_item = { - "model_key": model_key, - "shared_type": shared_type + suffix, - "shared_level": shared_level, - } - shared_links[shared_key]["links"].append(link_item) - - for model_key in model_config["model_dict"]: - model_params_item = model_config["model_dict"][model_key] - for item_key in model_params_item: - if item_key in supported_types: - item_params = model_params_item[item_key] - if isinstance(item_params, str): - replace_one_item(model_params_item, item_key, item_params) - elif item_params.get("type", "") == "hybrid": - for ii, hybrid_item in enumerate(item_params["list"]): - if isinstance(hybrid_item, str): - replace_one_item( - model_params_item[item_key]["list"], - item_key, - hybrid_item, - suffix=f"_hybrid_{ii}", - index=ii, - ) - for shared_key in shared_links: - shared_links[shared_key]["links"] = sorted( - shared_links[shared_key]["links"], - key=lambda x: ( - x["shared_level"] - - ("spin" in model_config["model_dict"][x["model_key"]]) * 100 - ), - ) - # little trick to make spin models in the front to be the base models, - # because its type embeddings are more general. - assert len(type_map_keys) == 1, "Multitask model must have only one type_map!" - return model_config, shared_links + return preprocess_shared_params_common(model_config, get_class_name) def get_class_name(item_key: str, item_params: dict[str, Any]) -> type: diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index 25d3ccdc49..1f63f3cb34 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -11,6 +11,9 @@ import tempfile import textwrap import unittest +from copy import ( + deepcopy, +) from pathlib import ( Path, ) @@ -43,6 +46,8 @@ from deepmd.jax.train.trainer import ( DPTrainer, _copy_matching_state_tree, + _merge_descriptor_stats, + _merge_fitting_param_stats, _scale_by_global_learning_rate, ) from deepmd.jax.utils.finetune import ( @@ -54,6 +59,9 @@ from deepmd.utils.compat import ( convert_optimizer_v31_to_v32, ) +from deepmd.utils.env_mat_stat import ( + StatItem, +) MODEL_SE_E2_A = { "type_map": ["O", "H", "B"], @@ -157,6 +165,59 @@ def _minimal_jax_config(model_params: dict) -> dict: } +def _minimal_jax_multitask_config(model_params: dict) -> dict: + return { + "model": model_params, + "training": { + "numb_steps": 1, + "data_dict": { + "task_a": {"training_data": {}}, + "task_b": {"training_data": {}}, + }, + }, + "learning_rate": { + "type": "exp", + "start_lr": 0.001, + "stop_lr": 1e-8, + "decay_steps": 1, + }, + "loss_dict": { + "task_a": {}, + "task_b": {}, + }, + } + + +def _shared_jax_model_config(*, share_fitting: bool = True) -> dict: + shared_dict: dict = { + "shared_type_map": ["O", "H", "B"], + "shared_descriptor": deepcopy(MODEL_SE_E2_A["descriptor"]), + } + fitting_ref_a: dict | str = deepcopy(MODEL_SE_E2_A["fitting_net"]) + fitting_ref_b: dict | str = deepcopy(MODEL_SE_E2_A["fitting_net"]) + if share_fitting: + shared_dict["shared_fitting"] = deepcopy(MODEL_SE_E2_A["fitting_net"]) + fitting_ref_a = "shared_fitting" + fitting_ref_b = "shared_fitting" + return { + "shared_dict": shared_dict, + "model_dict": { + "task_a": { + "type_map": "shared_type_map", + "descriptor": "shared_descriptor", + "fitting_net": fitting_ref_a, + "data_stat_nbatch": 1, + }, + "task_b": { + "type_map": "shared_type_map", + "descriptor": "shared_descriptor", + "fitting_net": fitting_ref_b, + "data_stat_nbatch": 1, + }, + }, + } + + @patch("deepmd.jax.train.trainer.DPTrainer._build_losses") @patch("deepmd.jax.train.trainer.DPTrainer._deserialize_models") @patch("deepmd.jax.train.trainer.serialize_from_file") @@ -214,6 +275,211 @@ def test_jax_restart_uses_checkpoint_model_script( assert trainer.start_step == 7 +def test_jax_train_entrypoint_preprocesses_shared_dict() -> None: + """JAX multi-task preprocessing expands shared_dict references.""" + entrypoint = JAXTrainEntrypoint() + config = { + "model": _shared_jax_model_config(), + "training": {}, + } + + updated = entrypoint.preprocess_config( + config, + TrainEntrypointOptions(input_file="input.json"), + ) + + model_dict = updated["model"]["model_dict"] + assert model_dict["task_a"]["type_map"] == ["O", "H", "B"] + assert model_dict["task_b"]["descriptor"]["type"] == "se_e2_a" + assert entrypoint.shared_links is not None + assert set(entrypoint.shared_links) == {"shared_descriptor", "shared_fitting"} + + +def test_jax_train_entrypoint_keeps_multitask_without_shared_dict() -> None: + """JAX multi-task configs without shared_dict keep the existing path.""" + entrypoint = JAXTrainEntrypoint() + config = { + "model": { + "model_dict": { + "task_a": deepcopy(MODEL_SE_E2_A), + "task_b": deepcopy(MODEL_SE_E2_A), + }, + }, + "training": {}, + } + + updated = entrypoint.preprocess_config( + config, + TrainEntrypointOptions(input_file="input.json"), + ) + + assert updated["model"]["model_dict"]["task_a"]["type_map"] == ["O", "H", "B"] + assert entrypoint.shared_links is None + + +@patch("deepmd.jax.entrypoints.train.serialize_from_file") +def test_jax_train_entrypoint_preprocesses_shared_dict_after_model_replacement( + serialize_from_file, +) -> None: + """JAX shared links match the final model after pretrain-script replacement.""" + input_model = { + "shared_dict": { + "input_type_map": ["input"], + "input_descriptor": deepcopy(MODEL_SE_E2_A["descriptor"]), + }, + "model_dict": { + "task_a": { + "type_map": "input_type_map", + "descriptor": "input_descriptor", + "fitting_net": deepcopy(MODEL_SE_E2_A["fitting_net"]), + }, + "task_b": { + "type_map": "input_type_map", + "descriptor": "input_descriptor", + "fitting_net": deepcopy(MODEL_SE_E2_A["fitting_net"]), + }, + }, + } + checkpoint_model = { + "shared_dict": { + "checkpoint_type_map": ["O", "H", "B"], + "checkpoint_descriptor": deepcopy(MODEL_SE_E2_A["descriptor"]), + }, + "model_dict": { + "task_a": { + "type_map": "checkpoint_type_map", + "descriptor": "checkpoint_descriptor", + "fitting_net": deepcopy(MODEL_SE_E2_A["fitting_net"]), + }, + "task_b": { + "type_map": "checkpoint_type_map", + "descriptor": "checkpoint_descriptor", + "fitting_net": deepcopy(MODEL_SE_E2_A["fitting_net"]), + }, + }, + } + serialize_from_file.return_value = { + "model_def_script": checkpoint_model, + } + entrypoint = JAXTrainEntrypoint() + config = {"model": input_model, "training": {}} + + updated = entrypoint.preprocess_config( + config, + TrainEntrypointOptions( + input_file="input.json", + init_model="model.jax", + use_pretrain_script=True, + ), + ) + + assert updated["model"]["model_dict"]["task_a"]["type_map"] == ["O", "H", "B"] + assert entrypoint.shared_links is not None + assert set(entrypoint.shared_links) == {"checkpoint_descriptor"} + assert "input_descriptor" not in entrypoint.shared_links + + +def test_jax_trainer_applies_shared_dict_links() -> None: + """Trainer-level sharing links descriptor and fitting-net parameters.""" + trainer = DPTrainer( + _minimal_jax_multitask_config(_shared_jax_model_config()), + ) + + trainer._share_model_params(resume=True) + + model_a = trainer.models["task_a"] + model_b = trainer.models["task_b"] + assert model_a.get_descriptor() is model_b.get_descriptor() + assert model_a.get_fitting_net() is not model_b.get_fitting_net() + assert model_a.get_fitting_net().nets is model_b.get_fitting_net().nets + + +class _FakeEnvMatStatSe: + def __init__(self, descriptor) -> None: + self.descriptor = descriptor + self.stats = {} + + def __call__(self) -> tuple[np.ndarray, np.ndarray]: + stat = self.stats["env"] + return ( + np.asarray([stat.compute_avg()], dtype=np.float64), + np.asarray([stat.compute_std()], dtype=np.float64), + ) + + +class _DescriptorWithStats: + def __init__(self, stats: dict[str, StatItem]) -> None: + self.stats = stats + self.davg = np.asarray([0.0], dtype=np.float64) + self.dstd = np.asarray([1.0], dtype=np.float64) + + +def test_jax_shared_descriptor_stats_merge_weighted_values() -> None: + """Shared descriptor merge recomputes weighted avg/std for nested stats.""" + base = _DescriptorWithStats({"env": StatItem(number=2, sum=4, squared_sum=10)}) + link = _DescriptorWithStats({"env": StatItem(number=4, sum=20, squared_sum=104)}) + base.se_atten = _DescriptorWithStats( + {"env": StatItem(number=2, sum=6, squared_sum=18)} + ) + link.se_atten = _DescriptorWithStats( + {"env": StatItem(number=4, sum=28, squared_sum=200)} + ) + + with patch("deepmd.dpmodel.utils.env_mat_stat.EnvMatStatSe", _FakeEnvMatStatSe): + _merge_descriptor_stats(base, link, model_prob=0.5) + + np.testing.assert_allclose(base.davg, [3.5]) + np.testing.assert_allclose(base.dstd, [np.sqrt(3.25)]) + np.testing.assert_allclose(base.se_atten.davg, [5.0]) + np.testing.assert_allclose(base.se_atten.dstd, [np.sqrt(4.5)]) + assert base.stats["env"].number == 4 + assert base.se_atten.stats["env"].number == 4 + + +class _FittingWithStats: + def __init__(self, param_stats: dict[str, list[StatItem]]) -> None: + self.numb_fparam = len(param_stats.get("fparam", [])) + self.numb_aparam = len(param_stats.get("aparam", [])) + self.fparam_avg = jnp.asarray(np.zeros(self.numb_fparam)) + self.fparam_inv_std = jnp.asarray(np.ones(self.numb_fparam)) + self.aparam_avg = jnp.asarray(np.zeros(self.numb_aparam)) + self.aparam_inv_std = jnp.asarray(np.ones(self.numb_aparam)) + self._param_stats = param_stats + + def get_param_stats(self) -> dict[str, list[StatItem]]: + return self._param_stats + + +def test_jax_shared_fitting_stats_merge_weighted_values() -> None: + """Shared fitting merge recomputes avg and protected inverse std.""" + base = _FittingWithStats( + { + "fparam": [StatItem(number=2, sum=4, squared_sum=10)], + "aparam": [StatItem(number=2, sum=4, squared_sum=8)], + } + ) + link = _FittingWithStats( + { + "fparam": [StatItem(number=4, sum=20, squared_sum=104)], + "aparam": [StatItem(number=2, sum=4, squared_sum=8)], + } + ) + + _merge_fitting_param_stats( + base, + link, + model_prob=0.5, + protection=0.25, + ) + + np.testing.assert_allclose(np.asarray(base.fparam_avg), [3.5]) + np.testing.assert_allclose(np.asarray(base.fparam_inv_std), [1.0 / np.sqrt(3.25)]) + np.testing.assert_allclose(np.asarray(base.aparam_avg), [2.0]) + np.testing.assert_allclose(np.asarray(base.aparam_inv_std), [4.0]) + assert base._param_stats["fparam"][0].number == 4 + assert base._param_stats["aparam"][0].number == 3 + + def test_jax_full_validator_saves_directory_best_checkpoint(tmp_path: Path) -> None: """JAX full validation uses .jax directory checkpoints.""" from deepmd.jax.train.validation import ( @@ -430,17 +696,6 @@ def test_train_entrypoint_rejects_remaining_unsupported_features(self) -> None: ), "init_frz_model", ), - ( - { - "model": { - "model_dict": {"task": {}}, - "shared_dict": {"shared": {}}, - }, - "training": {}, - }, - TrainEntrypointOptions(input_file="input.json"), - "shared_dict", - ), ] for config, options, message in cases: