diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index 00d7337254..5e72b98882 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -59,7 +59,7 @@ jobs: run: curl --proto '=https' --tlsv1.2 -LsSf https://github.com/astral-sh/uv/releases/download/0.2.24/uv-installer.sh | sh if: runner.os != 'Linux' - name: Build wheels - uses: pypa/cibuildwheel@v3.3 + uses: pypa/cibuildwheel@v3.4 env: CIBW_BUILD_VERBOSITY: 1 CIBW_ARCHS: all diff --git a/.github/workflows/package_c.yml b/.github/workflows/package_c.yml index 8589faab16..ae1b7d261b 100644 --- a/.github/workflows/package_c.yml +++ b/.github/workflows/package_c.yml @@ -21,8 +21,8 @@ jobs: strategy: matrix: include: - - tensorflow_build_version: "2.18" - tensorflow_version: "" + - tensorflow_build_version: "2.20" + tensorflow_version: "==2.20.*" filename: libdeepmd_c.tar.gz steps: - name: Free Disk Space (Ubuntu) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ea1e4817ce..07391696ea 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -39,14 +39,6 @@ repos: - id: ruff-format exclude: ^source/3rdparty types_or: [python, pyi, jupyter] - - repo: https://github.com/pycqa/flake8 - # flake8 cannot autofix - rev: "7.3.0" - hooks: - - id: flake8 - additional_dependencies: - - torchfix==0.7.0 - - flake8-pyproject==1.2.3 # numpydoc - repo: https://github.com/Carreau/velin rev: 0.0.12 diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index e745b28f94..c45bab0bc9 100644 --- a/deepmd/dpmodel/array_api.py +++ b/deepmd/dpmodel/array_api.py @@ -15,6 +15,21 @@ Array = np.ndarray | Any # Any to support JAX, PyTorch, etc. arrays +def _xp_arange( + xp: Any, + size: int, + dtype: Any, + like: Array, +) -> Array: + if array_api_compat.is_jax_array(like): + return xp.arange(size, dtype=dtype) # pylint: disable=no-explicit-device + return xp.arange( + size, + dtype=dtype, + device=array_api_compat.device(like), + ) + + # array api adds take_along_axis in https://github.com/data-apis/array-api/pull/816 # but it hasn't been released yet # below is a pure Python implementation of take_along_axis @@ -51,8 +66,7 @@ def xp_take_along_axis(arr: Array, indices: Array, axis: int) -> Array: else: indices = xp.reshape(indices, (0, 0)) - dev = array_api_compat.device(indices) - offset = (xp.arange(indices.shape[0], dtype=indices.dtype, device=dev) * m)[ + offset = (_xp_arange(xp, indices.shape[0], indices.dtype, indices) * m)[ :, xp.newaxis ] indices = xp.reshape(offset + indices, (-1,)) @@ -78,7 +92,7 @@ def xp_scatter_sum(input: Array, dim: int, index: Array, src: Array) -> Array: xp = array_api_compat.array_namespace(input) # Create flat index array matching input shape - idx = xp.arange(input.size, dtype=xp.int64, device=array_api_compat.device(input)) + idx = _xp_arange(xp, input.size, xp.int64, input) idx = xp.reshape(idx, input.shape) # Get flat indices where we want to add values diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index 47a4fb1478..5cdeafc34d 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -487,8 +487,14 @@ def change_type_map( ) repflow.ntypes = self.ntypes repflow.reinit_exclude(self.exclude_types) - repflow["davg"] = repflow["davg"][remap_index] - repflow["dstd"] = repflow["dstd"][remap_index] + xp = array_api_compat.array_namespace(repflow["davg"]) + remap_index_array = xp.asarray( + remap_index, + dtype=xp.int32, + device=array_api_compat.device(repflow["davg"]), + ) + repflow["davg"] = repflow["davg"][remap_index_array] + repflow["dstd"] = repflow["dstd"][remap_index_array] @property def dim_out(self) -> int: diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 3188bbfee5..0f6a953a87 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -58,6 +58,28 @@ ) +def _maybe_apply_jax_placeholder_sharding(array: Array) -> Array: + """Keep static placeholder tensors replicated when a mesh context exists. + + JAX training runs under an explicit mesh and benefits from pinning these + dummy index tensors to replicated sharding. Export/freeze does not create + a mesh context, so the same sharding constraint must be skipped there. + """ + from deepmd.jax.env import ( + jax, + ) + + try: + return jax.lax.with_sharding_constraint( + array, jax.sharding.PartitionSpec(None, None) + ) + except RuntimeError as exc: + message = str(exc) + if "non-empty mesh" not in message: + raise + return array + + @DescriptorBlock.register("se_repflow") class DescrptBlockRepflows(NativeOP, DescriptorBlock): r""" @@ -595,12 +617,25 @@ def call( # n_angle x 1 a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask] else: + dummy_edge_count = max(1, self.nnei) + dummy_angle_count = max(1, self.a_sel * self.a_sel) edge_index = xp.zeros( - [2, 1], dtype=nlist.dtype, device=array_api_compat.device(nlist) + [2, dummy_edge_count], + dtype=nlist.dtype, + device=array_api_compat.device(nlist), ) angle_index = xp.zeros( - [3, 1], dtype=nlist.dtype, device=array_api_compat.device(nlist) + [3, dummy_angle_count], + dtype=nlist.dtype, + device=array_api_compat.device(nlist), ) + if array_api_compat.is_jax_namespace(xp): + # These are placeholders in the static-selection path and should + # stay replicated instead of inheriting the active natoms mesh. + # Freeze/export does not run under a mesh, so skip the + # constraint there and keep the placeholders as-is. + edge_index = _maybe_apply_jax_placeholder_sharding(edge_index) + angle_index = _maybe_apply_jax_placeholder_sharding(angle_index) # get edge and angle embedding # nb x nloc x nnei x e_dim [OR] n_edge x e_dim diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 4710987f54..49f9966f8d 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -461,38 +461,21 @@ def call( gg = self.cal_g(ss, (tt,)) gr += xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1) else: - # Sort atoms by center type so each type forms a contiguous block. - # Slice indexing (arr[s:e]) is array-api compatible and lets us - # run cal_g only on atoms of the matching center type, keeping the - # same O(nf*nloc) total embedding cost as the original numpy code. atype_loc = xp.reshape(atype_ext[:, :nloc], (nf * nloc,)) - sort_idx = xp.argsort(atype_loc) - unsort_idx = xp.argsort(sort_idx) - rr_s = xp.take(rr, sort_idx, axis=0) - mask_s = xp.take(exclude_mask, sort_idx, axis=0) - dev = array_api_compat.device(coord_ext) - gr_s = xp.zeros([nf * nloc, ng, 4], dtype=input_dtype, device=dev) - # Per-type boundaries in sorted order - type_ends = [] - offset = 0 for ti in range(self.ntypes): - offset += int(xp.sum(xp.astype(atype_loc == ti, xp.int32))) - type_ends.append(offset) - type_starts = [0, *type_ends[:-1]] - for ti in range(self.ntypes): - s, e = type_starts[ti], type_ends[ti] - if s == e: - continue + center_mask = xp.astype(atype_loc == ti, input_dtype) + center_mask = xp.reshape(center_mask, (nf * nloc, 1, 1)) for tt in range(self.ntypes): - mm = mask_s[s:e, sec[tt] : sec[tt + 1]] - tr = rr_s[s:e, sec[tt] : sec[tt + 1], :] + mm = exclude_mask[:, sec[tt] : sec[tt + 1]] + tr = rr[:, sec[tt] : sec[tt + 1], :] tr = tr * xp.astype(mm[:, :, None], tr.dtype) ss = tr[..., 0:1] gg = self.cal_g(ss, (ti, tt)) - gr_s[s:e] = gr_s[s:e] + xp.sum( - gg[:, :, :, None] * tr[:, :, None, :], axis=1 + gr = ( + gr + + xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1) + * center_mask ) - gr = xp.take(gr_s, unsort_idx, axis=0) gr = xp.reshape(gr, (nf, nloc, ng, 4)) # nf x nloc x ng x 4 gr /= self.nnei diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 260be619fd..33401c40fd 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -368,7 +368,13 @@ def change_type_map( self.bias_atom_e = np.concatenate( [self.bias_atom_e, extend_bias_atom_e], axis=0 ) - self.bias_atom_e = self.bias_atom_e[remap_index] + xp = array_api_compat.array_namespace(self.bias_atom_e) + remap_index_array = xp.asarray( + remap_index, + dtype=xp.int32, + device=array_api_compat.device(self.bias_atom_e), + ) + self.bias_atom_e = self.bias_atom_e[remap_index_array] def __setitem__(self, key: str, value: Any) -> None: if key in ["bias_atom_e"]: diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index cbe782e91c..bebf75e12b 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -119,14 +119,15 @@ def build_neighbor_list( device = array_api_compat.device(diff) if array_api_compat.is_jax_namespace(xp): # fix jax sharding "list index out of range" - from jax.sharding import PartitionSpec as P, NamedSharding + from jax.sharding import ( + NamedSharding, + ) + from jax.sharding import PartitionSpec as P if isinstance(device, NamedSharding): device = NamedSharding(device.mesh, P()) # if central atom has two zero distances, sorting sometimes can not exclude itself - rr -= xp.eye(nloc, nall, dtype=diff.dtype, device=device)[ - xp.newaxis, :, : - ] + rr -= xp.eye(nloc, nall, dtype=diff.dtype, device=device)[xp.newaxis, :, :] nlist = xp.argsort(rr, axis=-1) rr = xp.sort(rr, axis=-1) rr = rr[:, :, 1:] @@ -309,16 +310,20 @@ def extend_coord_with_ghosts( if array_api_compat.is_jax_namespace(xp): # fix jax: Sharding is only valid for values of rank at least 2, # but was applied to a value of rank 1. - from jax.sharding import PartitionSpec as P, NamedSharding + from jax.sharding import ( + NamedSharding, + ) + from jax.sharding import PartitionSpec as P if isinstance(device, NamedSharding): - device_nloc = NamedSharding(device.mesh, P(device.spec[1])) + if len(device.spec) > 1: + device_nloc = NamedSharding(device.mesh, P(device.spec[1])) + else: + device_nloc = NamedSharding(device.mesh, P()) device_none = NamedSharding(device.mesh, P()) # int64 for index aidx = xp.tile( - xp.arange(nloc, dtype=xp.int64, device=device_nloc)[ - xp.newaxis, : - ], + xp.arange(nloc, dtype=xp.int64, device=device_nloc)[xp.newaxis, :], (nf, 1), ) if cell is None: diff --git a/deepmd/dpmodel/utils/region.py b/deepmd/dpmodel/utils/region.py index 3a0339ff79..e596fcdce3 100644 --- a/deepmd/dpmodel/utils/region.py +++ b/deepmd/dpmodel/utils/region.py @@ -77,13 +77,14 @@ def normalize_coord( device = array_api_compat.device(icoord) if array_api_compat.is_jax_namespace(xp): # fix jax sharding "list index out of range" - from jax.sharding import PartitionSpec as P, NamedSharding + from jax.sharding import ( + NamedSharding, + ) + from jax.sharding import PartitionSpec as P if isinstance(device, NamedSharding): device = NamedSharding(device.mesh, P()) - icoord = xp.remainder( - icoord, xp.ones((), dtype=icoord.dtype, device=device) - ) + icoord = xp.remainder(icoord, xp.ones((), dtype=icoord.dtype, device=device)) return inter2phys(icoord, cell) diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py index eea7a55bc3..ba70fd5114 100644 --- a/deepmd/dpmodel/utils/type_embed.py +++ b/deepmd/dpmodel/utils/type_embed.py @@ -226,7 +226,12 @@ def change_type_map( [first_layer_matrix, extend_type_params], axis=0 ) - first_layer_matrix = first_layer_matrix[remap_index] + remap_index_array = xp.asarray( + remap_index, + dtype=xp.int32, + device=array_api_compat.device(first_layer_matrix), + ) + first_layer_matrix = first_layer_matrix[remap_index_array] new_ntypes = len(type_map) eye_vector = xp.eye( new_ntypes, diff --git a/deepmd/jax/descriptor/hybrid.py b/deepmd/jax/descriptor/hybrid.py index b76e515c54..7e080c4c78 100644 --- a/deepmd/jax/descriptor/hybrid.py +++ b/deepmd/jax/descriptor/hybrid.py @@ -3,11 +3,15 @@ Any, ) +import array_api_compat from packaging.version import ( Version, ) from deepmd.dpmodel.descriptor.hybrid import DescrptHybrid as DescrptHybridDP +from deepmd.dpmodel.utils.nlist import ( + nlist_distinguish_types, +) from deepmd.jax.common import ( ArrayAPIVariable, flax_module, @@ -36,3 +40,44 @@ def __setattr__(self, name: str, value: Any) -> None: value = nnx.List([nnx.data(item) for item in value]) return super().__setattr__(name, value) + + def call( + self, *args: Any, **kwargs: Any + ) -> tuple[Any, Any | None, Any | None, Any | None, Any | None]: + if len(args) < 3: + return super().call(*args, **kwargs) + if len(args) > 4: + return super().call(*args, **kwargs) + if kwargs and set(kwargs) != {"mapping"}: + return super().call(*args, **kwargs) + coord_ext, atype_ext, nlist = args[:3] + mapping = kwargs.pop("mapping", args[3] if len(args) == 4 else None) + xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) + out_descriptor = [] + out_gr = [] + out_g2 = None + out_h2 = None + out_sw = None + if self.sel_no_mixed_types is not None: + nl_distinguish_types = nlist_distinguish_types( + nlist, + atype_ext, + self.sel_no_mixed_types, + ) + else: + nl_distinguish_types = None + for descrpt, nci in zip(self.descrpt_list, self.nlist_cut_idx, strict=True): + nci_value = getattr(nci, "value", nci) + if self.mixed_types() == descrpt.mixed_types(): + nl = xp.take(nlist, nci_value, axis=2) + else: + assert nl_distinguish_types is not None + nl = nl_distinguish_types[:, :, nci_value] + odescriptor, gr, g2, h2, sw = descrpt(coord_ext, atype_ext, nl, mapping) + out_descriptor.append(odescriptor) + if gr is not None: + out_gr.append(gr) + + out_descriptor = xp.concat(out_descriptor, axis=-1) + out_gr = xp.concat(out_gr, axis=-2) if out_gr else None + return out_descriptor, out_gr, out_g2, out_h2, out_sw diff --git a/deepmd/jax/entrypoints/freeze.py b/deepmd/jax/entrypoints/freeze.py index 345b2690ae..25d744e969 100644 --- a/deepmd/jax/entrypoints/freeze.py +++ b/deepmd/jax/entrypoints/freeze.py @@ -8,6 +8,7 @@ from deepmd.jax.utils.serialization import ( deserialize_to_file, + select_model_branch, serialize_from_file, ) @@ -16,6 +17,8 @@ def freeze( *, checkpoint_folder: str, output: str, + head: str | None = None, + model_branch: str | None = None, hessian: bool = False, **kwargs: Any, ) -> None: @@ -37,6 +40,9 @@ def freeze( checkpoint_folder = checkpoint_meta.read_text().strip() if Path(checkpoint_folder).is_dir(): data = serialize_from_file(checkpoint_folder) + selected_branch = model_branch or head + if selected_branch and "model_dict" in data["model_def_script"]: + data = select_model_branch(data, selected_branch) deserialize_to_file(output, data, hessian=hessian) else: raise FileNotFoundError(f"Checkpoint {checkpoint_folder} does not exist.") diff --git a/deepmd/jax/entrypoints/main.py b/deepmd/jax/entrypoints/main.py index 6bbb9f08f7..94f48d14c7 100644 --- a/deepmd/jax/entrypoints/main.py +++ b/deepmd/jax/entrypoints/main.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, - Union, -) from deepmd.backend.suffix import ( format_model_suffix, @@ -29,7 +25,7 @@ __all__ = ["main"] -def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None: +def main(args: list[str] | argparse.Namespace | None = None) -> None: """DeePMD-Kit entry point. Parameters diff --git a/deepmd/jax/entrypoints/train.py b/deepmd/jax/entrypoints/train.py index 7be5b3182a..e13c847f83 100644 --- a/deepmd/jax/entrypoints/train.py +++ b/deepmd/jax/entrypoints/train.py @@ -10,10 +10,8 @@ import time from typing import ( Any, - Optional, ) - from deepmd.common import ( j_loader, ) @@ -24,6 +22,15 @@ from deepmd.jax.train.trainer import ( DPTrainer, ) +from deepmd.jax.utils.finetune import ( + get_finetune_rules, +) +from deepmd.jax.utils.multi_task import ( + preprocess_shared_params, +) +from deepmd.jax.utils.serialization import ( + serialize_from_file, +) from deepmd.utils import random as dp_random from deepmd.utils.argcheck import ( normalize, @@ -45,87 +52,47 @@ class SummaryPrinter(BaseSummaryPrinter): """Summary printer for JAX.""" def is_built_with_cuda(self) -> bool: - """Check if the backend is built with CUDA.""" return jax_export.default_export_platform() == "cuda" def is_built_with_rocm(self) -> bool: - """Check if the backend is built with ROCm.""" return jax_export.default_export_platform() == "rocm" def get_compute_device(self) -> str: - """Get Compute device.""" return jax.default_backend() def get_ngpus(self) -> int: - """Get the number of GPUs.""" return jax.device_count() def get_backend_info(self) -> dict: - """Get backend information.""" return { "Backend": "JAX", "JAX ver": jax.__version__, } def get_device_name(self) -> str: - """Get the name of the device.""" devices = jax.devices() if devices: return devices[0].device_kind - else: - return "Unknown" + return "Unknown" def train( *, INPUT: str, - init_model: Optional[str], - restart: Optional[str], + init_model: str | None, + restart: str | None, output: str, init_frz_model: str, mpi_log: str, log_level: int, - log_path: Optional[str], + log_path: str | None, skip_neighbor_stat: bool = False, - finetune: Optional[str] = None, + finetune: str | None = None, use_pretrain_script: bool = False, + force_load: bool = False, + model_branch: str = "", **kwargs: Any, ) -> None: - """Run DeePMD model training. - - Parameters - ---------- - INPUT : str - json/yaml control file - init_model : Optional[str] - path prefix of checkpoint files or None - restart : Optional[str] - path prefix of checkpoint files or None - output : str - path for dump file with arguments - init_frz_model : str - path to frozen model or None - mpi_log : str - mpi logging mode - log_level : int - logging level defined by int 0-3 - log_path : Optional[str] - logging file path or None if logs are to be output only to stdout - skip_neighbor_stat : bool, default=False - skip checking neighbor statistics - finetune : Optional[str] - path to pretrained model or None - use_pretrain_script : bool - Whether to use model script in pretrained model when doing init-model or init-frz-model. - Note that this option is true and unchangeable for fine-tuning. - **kwargs - additional arguments - - Raises - ------ - RuntimeError - if distributed training job name is wrong - """ if int(os.environ.get("DP_JAX_MULTI_NPROC", "0")) > 1: multi_nproc = int(os.environ.get("DP_JAX_MULTI_NPROC", "0")) if multi_nproc <= 0: @@ -142,71 +109,108 @@ def train( process_id=multi_iproc, ) - # load json database jdata = j_loader(INPUT) - origin_type_map = None + multi_task = "model_dict" in jdata["model"] + shared_links = None + if multi_task: + jdata["model"], shared_links = preprocess_shared_params(jdata["model"]) + if "RANDOM" in jdata["model"]["model_dict"]: + raise ValueError("Model name can not be 'RANDOM' in multi-task mode!") + + finetune_links = None + finetune_data = None + if finetune is not None: + jdata["model"], finetune_links, finetune_data = get_finetune_rules( + finetune, + jdata["model"], + model_branch=model_branch, + change_model_params=use_pretrain_script, + ) + if (init_model is not None or init_frz_model) and use_pretrain_script: + source_model = init_model if init_model is not None else init_frz_model + source_model_data = serialize_from_file(source_model) + jdata["model"] = source_model_data["model_def_script"] jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") - - jdata = normalize(jdata) - jdata = update_sel(jdata) + jdata = normalize(jdata, multi_task=multi_task) + jdata = update_sel(jdata, multi_task=multi_task) with open(output, "w") as fp: json.dump(jdata, fp, indent=4) SummaryPrinter()() - # make necessary checks - assert "training" in jdata - - # init the model - model = DPTrainer( jdata, init_model=init_model, restart=restart, + init_frz_model=init_frz_model or None, + finetune_model=finetune, + force_load=force_load, + shared_links=shared_links, + finetune_links=finetune_links, + finetune_model_data=finetune_data, ) - rcut = model.model.get_rcut() - type_map = model.model.get_type_map() - if len(type_map) == 0: - ipt_type_map = None - else: - ipt_type_map = type_map - # init random seed of data systems seed = jdata["training"].get("seed", None) if seed is not None: seed += jax.process_index() seed = seed % (2**32) dp_random.seed(seed) - # init data - train_data = get_data(jdata["training"]["training_data"], rcut, ipt_type_map, None) - train_data.add_data_requirements(model.data_requirements) - train_data.print_summary("training") - if jdata["training"].get("validation_data", None) is not None: - valid_data = get_data( - jdata["training"]["validation_data"], - rcut, - train_data.type_map, - None, + if not multi_task: + rcut = model.model.get_rcut() + type_map = model.model.get_type_map() + ipt_type_map = None if len(type_map) == 0 else type_map + train_data = get_data( + jdata["training"]["training_data"], rcut, ipt_type_map, None ) - valid_data.add_data_requirements(model.data_requirements) - valid_data.print_summary("validation") + train_data.add_data_requirements(model.data_requirements) + train_data.print_summary("training") + if jdata["training"].get("validation_data", None) is not None: + valid_data = get_data( + jdata["training"]["validation_data"], + rcut, + train_data.type_map, + None, + ) + valid_data.add_data_requirements(model.data_requirements) + valid_data.print_summary("validation") + else: + valid_data = None else: - valid_data = None - - # get training info - stop_batch = jdata["training"]["numb_steps"] - origin_type_map = jdata["model"].get("origin_type_map", None) - if ( - origin_type_map is not None and not origin_type_map - ): # get the type_map from data if not provided - origin_type_map = get_data( - jdata["training"]["training_data"], rcut, None, None - ).get_type_map() - - # train the model with the provided systems in a cyclic way + train_data = {} + valid_data = {} + for model_key in model.model_keys: + branch_model = model.model[model_key] + rcut = branch_model.get_rcut() + type_map = branch_model.get_type_map() + ipt_type_map = None if len(type_map) == 0 else type_map + branch_train = get_data( + jdata["training"]["data_dict"][model_key]["training_data"], + rcut, + ipt_type_map, + None, + ) + branch_train.add_data_requirements(model.data_requirements[model_key]) + branch_train.print_summary(f"training in {model_key}") + train_data[model_key] = branch_train + if ( + jdata["training"]["data_dict"][model_key].get("validation_data", None) + is not None + ): + branch_valid = get_data( + jdata["training"]["data_dict"][model_key]["validation_data"], + rcut, + branch_train.type_map, + None, + ) + branch_valid.add_data_requirements(model.data_requirements[model_key]) + branch_valid.print_summary(f"validation in {model_key}") + valid_data[model_key] = branch_valid + else: + valid_data[model_key] = None + start_time = time.time() model.train(train_data, valid_data) end_time = time.time() @@ -214,20 +218,28 @@ def train( log.info(f"wall time: {(end_time - start_time):.3f} s") -def update_sel(jdata: dict) -> dict: +def update_sel(jdata: dict, *, multi_task: bool = False) -> dict: log.info( "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)" ) jdata_cpy = jdata.copy() - type_map = jdata["model"].get("type_map") - train_data = get_data( - jdata["training"]["training_data"], - 0, # not used - type_map, - None, # not used - ) - # TODO: OOM, need debug - # jdata_cpy["model"], min_nbor_dist = BaseModel.update_sel( - # train_data, type_map, jdata["model"] - # ) + if not multi_task: + type_map = jdata["model"].get("type_map") + train_data = get_data( + jdata["training"]["training_data"], + 0, + type_map, + None, + ) + del train_data + else: + for model_key in jdata["model"]["model_dict"]: + type_map = jdata["model"]["model_dict"][model_key].get("type_map") + train_data = get_data( + jdata["training"]["data_dict"][model_key]["training_data"], + 0, + type_map, + None, + ) + del train_data return jdata_cpy diff --git a/deepmd/jax/env.py b/deepmd/jax/env.py index d425d5a0e8..30a13f740a 100644 --- a/deepmd/jax/env.py +++ b/deepmd/jax/env.py @@ -19,6 +19,9 @@ if os.environ.get("DP_DTYPE_PROMOTION_STRICT") == "1": jax.config.update("jax_numpy_dtype_promotion", "strict") +if not hasattr(jax.monitoring, "record_scalar"): + jax.monitoring.record_scalar = lambda *args, **kwargs: None + __all__ = [ "flax_version", "jax", diff --git a/deepmd/jax/model/model.py b/deepmd/jax/model/model.py index 321f33b315..8014860acb 100644 --- a/deepmd/jax/model/model.py +++ b/deepmd/jax/model/model.py @@ -24,6 +24,12 @@ from deepmd.jax.model.dp_zbl_model import ( DPZBLModel, ) +from deepmd.jax.model.multitask import ( + ModelWrapper, +) +from deepmd.jax.utils.multi_task import ( + get_case_embd_config, +) def get_standard_model(data: dict) -> BaseModel: @@ -121,3 +127,23 @@ def get_model(data: dict) -> BaseModel: return get_standard_model(data) else: return BaseModel.get_class_by_type(model_type).get_model(data) + + +def get_model_for_wrapper( + model_params: dict, + *, + shared_links: dict | None = None, +) -> BaseModel | ModelWrapper: + if "model_dict" not in model_params: + return get_model(model_params) + + model_dict = { + model_key: get_model(model_params["model_dict"][model_key]) + for model_key in model_params["model_dict"] + } + do_case_embd, case_embd_index = get_case_embd_config(model_params) + return ModelWrapper( + model_dict, + shared_links=shared_links, + case_embd_index=case_embd_index if do_case_embd else None, + ) diff --git a/deepmd/jax/model/multitask.py b/deepmd/jax/model/multitask.py new file mode 100644 index 0000000000..8fad344e83 --- /dev/null +++ b/deepmd/jax/model/multitask.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +from deepmd.jax.common import ( + flax_module, +) +from deepmd.jax.model.base_model import ( + BaseModel, +) + + +def _branch_attr(branch: str) -> str: + return f"branch__{branch}" + + +def _get_component(model: BaseModel, shared_type: str) -> Any: + if shared_type == "descriptor": + return model.atomic_model.descriptor + if shared_type == "fitting_net": + return model.atomic_model.fitting + if shared_type.startswith("descriptor_hybrid_"): + idx = int(shared_type.rsplit("_", 1)[1]) + return model.atomic_model.descriptor.descrpt_list[idx] + raise NotImplementedError(f"Unsupported shared_type {shared_type}") + + +def _set_component(model: BaseModel, shared_type: str, value: Any) -> None: + if shared_type == "descriptor": + model.atomic_model.descriptor = value + return + if shared_type == "fitting_net": + model.atomic_model.fitting = value + return + if shared_type.startswith("descriptor_hybrid_"): + idx = int(shared_type.rsplit("_", 1)[1]) + model.atomic_model.descriptor.descrpt_list[idx] = value + return + raise NotImplementedError(f"Unsupported shared_type {shared_type}") + + +def _share_fitting_component(base_model: BaseModel, link_model: BaseModel) -> None: + """Mirror PT fitting sharing semantics without aliasing branch-local state. + + PT shared fitting keeps `bias_atom_e` and `case_embd` per branch, while + sharing the inner network modules plus the fparam/aparam normalization + buffers. JAX needs the same behavior for multitask parity. + """ + base_fitting = base_model.atomic_model.fitting + link_fitting = link_model.atomic_model.fitting + if base_fitting.__class__ is not link_fitting.__class__: + raise TypeError("Only fitting nets of the same type can share params!") + object.__setattr__(link_fitting, "nets", base_fitting.nets) + if getattr(base_fitting, "numb_fparam", 0) > 0: + object.__setattr__(link_fitting, "fparam_avg", base_fitting.fparam_avg) + object.__setattr__( + link_fitting, + "fparam_inv_std", + base_fitting.fparam_inv_std, + ) + if getattr(base_fitting, "numb_aparam", 0) > 0: + object.__setattr__(link_fitting, "aparam_avg", base_fitting.aparam_avg) + object.__setattr__( + link_fitting, + "aparam_inv_std", + base_fitting.aparam_inv_std, + ) + + +def _share_partial_descriptor_component( + base_component: Any, + link_component: Any, + shared_level: int, +) -> None: + if base_component.__class__ is not link_component.__class__: + raise TypeError("Only descriptors of the same type can share params!") + + descriptor_name = base_component.__class__.__name__ + if shared_level == 1 and descriptor_name in { + "DescrptDPA1", + "DescrptDPA2", + "DescrptDPA3", + "DescrptSeTTebd", + }: + object.__setattr__( + link_component, + "type_embedding", + base_component.type_embedding, + ) + return + + raise NotImplementedError( + f"Unsupported descriptor partial sharing for {descriptor_name} " + f"with shared_level={shared_level}." + ) + + +def _check_supported_share(link_info: dict[str, Any]) -> None: + links = link_info.get("links", []) + for link in links: + shared_type = link["shared_type"] + shared_level = int(link.get("shared_level", 0)) + if shared_type == "fitting_net" and shared_level != 0: + raise NotImplementedError( + "JAX multitask fitting_net sharing only supports shared_level=0, " + f"but got {shared_type}:{shared_level}." + ) + + +@flax_module +class ModelWrapper: + def __init__( + self, + model_dict: dict[str, BaseModel], + *, + shared_links: dict[str, Any] | None = None, + case_embd_index: dict[str, int] | None = None, + ) -> None: + self.model_keys = list(model_dict.keys()) + self.shared_links = shared_links or {} + self.case_embd_index = case_embd_index or {} + for key, model in model_dict.items(): + setattr(self, _branch_attr(key), model) + if self.shared_links: + self.share_params(self.shared_links) + for key in self.model_keys: + self.set_case_embd(key) + + def keys(self) -> list[str]: + return list(self.model_keys) + + def items(self) -> list[tuple[str, BaseModel]]: + return [(key, self[key]) for key in self.model_keys] + + def __getitem__(self, key: str) -> BaseModel: + return getattr(self, _branch_attr(key)) + + def __setitem__(self, key: str, model: BaseModel) -> None: + setattr(self, _branch_attr(key), model) + if self.shared_links: + self.share_params(self.shared_links) + self.set_case_embd(key) + + def get_type_map(self, key: str) -> list[str]: + return self[key].get_type_map() + + def share_params(self, shared_links: dict[str, Any] | None = None) -> None: + shared_links = shared_links or self.shared_links + for _, link_info in shared_links.items(): + _check_supported_share(link_info) + links = link_info.get("links", []) + if not links: + continue + base_link = links[0] + base_model = self[base_link["model_key"]] + base_shared_type = base_link["shared_type"] + if base_shared_type == "fitting_net": + for link in links[1:]: + _share_fitting_component(base_model, self[link["model_key"]]) + continue + base_component = _get_component(base_model, base_shared_type) + for link in links[1:]: + shared_level = int(link.get("shared_level", 0)) + link_model = self[link["model_key"]] + if shared_level == 0: + _set_component( + link_model, + link["shared_type"], + base_component, + ) + continue + link_component = _get_component(link_model, link["shared_type"]) + _share_partial_descriptor_component( + base_component, + link_component, + shared_level, + ) + + def set_case_embd(self, key: str) -> None: + if key in self.case_embd_index: + self[key].set_case_embd(self.case_embd_index[key]) + + def serialize(self) -> dict[str, Any]: + return {"model_dict": {key: self[key].serialize() for key in self.model_keys}} + + @classmethod + def deserialize( + cls, + data: dict[str, Any], + *, + shared_links: dict[str, Any] | None = None, + case_embd_index: dict[str, int] | None = None, + ) -> ModelWrapper: + model_dict = { + key: BaseModel.deserialize(value) + for key, value in data["model_dict"].items() + } + return cls( + model_dict, + shared_links=shared_links, + case_embd_index=case_embd_index, + ) diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index 36ecfc9429..62fe3562a2 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -4,26 +4,39 @@ import os import shutil import time +from collections.abc import ( + Callable, +) +from copy import ( + deepcopy, +) from pathlib import ( Path, ) from typing import ( Any, - Optional, TextIO, ) +import array_api_compat import numpy as np import optax import orbax.checkpoint as ocp +from jax.sharding import ( + Mesh, + NamedSharding, +) +from jax.sharding import PartitionSpec as P from packaging.version import ( Version, ) -from jax.sharding import PartitionSpec as P, NamedSharding from deepmd.common import ( symlink_prefix_files, ) +from deepmd.dpmodel.common import ( + to_numpy_array, +) from deepmd.dpmodel.loss.ener import ( EnergyHessianLoss, EnergyLoss, @@ -31,6 +44,10 @@ from deepmd.dpmodel.model.transform_output import ( communicate_extended_output, ) +from deepmd.dpmodel.utils import ( + compute_total_numb_batch, + resolve_model_prob_from_epochs, +) from deepmd.dpmodel.utils.learning_rate import ( LearningRateExp, ) @@ -51,79 +68,624 @@ BaseModel, ) from deepmd.jax.model.model import ( - get_model, + get_model_for_wrapper, +) +from deepmd.jax.model.multitask import ( + ModelWrapper, +) +from deepmd.jax.utils.finetune import ( + merge_finetune_model_data, +) +from deepmd.jax.utils.multi_task import ( + get_case_embd_config, ) from deepmd.jax.utils.serialization import ( + select_model_branch, serialize_from_file, ) from deepmd.loggers.training import ( format_training_message, format_training_message_per_task, ) +from deepmd.utils import random as dp_random from deepmd.utils.data import ( DataRequirementItem, ) from deepmd.utils.data_system import ( DeepmdDataSystem, ) +from deepmd.utils.finetune import ( + FinetuneRuleItem, +) from deepmd.utils.model_stat import ( make_stat_input, ) log = logging.getLogger(__name__) +_ACTIVE_JAX_MESH_CONTEXT: list[Any] = [] + +DefModel = BaseModel | ModelWrapper +DefLoss = EnergyLoss | EnergyHessianLoss + + +def _merge_init_frz_model_data( + target_node: Any, + source_node: Any, + *, + path: tuple[Any, ...] = (), + missing: list[tuple[Any, ...]] | None = None, + unexpected: list[tuple[Any, ...]] | None = None, +) -> Any: + """Merge overlapping frozen-model tensor leaves into a target model tree.""" + if missing is None: + missing = [] + if unexpected is None: + unexpected = [] + + if isinstance(target_node, dict): + merged = deepcopy(target_node) + if not isinstance(source_node, dict): + missing.append(path) + return merged + for key in source_node: + if key not in target_node: + unexpected.append((*path, key)) + for key, value in target_node.items(): + if key not in source_node: + missing.append((*path, key)) + continue + merged[key] = _merge_init_frz_model_data( + value, + source_node[key], + path=(*path, key), + missing=missing, + unexpected=unexpected, + ) + return merged + + if isinstance(target_node, list): + merged = deepcopy(target_node) + if not isinstance(source_node, list): + missing.append(path) + return merged + if len(target_node) != len(source_node): + raise ValueError( + f"Shape mismatch at {path}: target list len {len(target_node)}, " + f"source list len {len(source_node)}" + ) + for idx, value in enumerate(target_node): + merged[idx] = _merge_init_frz_model_data( + value, + source_node[idx], + path=(*path, idx), + missing=missing, + unexpected=unexpected, + ) + return merged + + if isinstance(target_node, tuple): + if not isinstance(source_node, tuple): + missing.append(path) + return deepcopy(target_node) + if len(target_node) != len(source_node): + raise ValueError( + f"Shape mismatch at {path}: target tuple len {len(target_node)}, " + f"source tuple len {len(source_node)}" + ) + return tuple( + _merge_init_frz_model_data( + tv, + sv, + path=(*path, idx), + missing=missing, + unexpected=unexpected, + ) + for idx, (tv, sv) in enumerate(zip(target_node, source_node, strict=True)) + ) + + if isinstance(target_node, np.ndarray): + if not isinstance(source_node, np.ndarray): + missing.append(path) + return np.array(target_node, copy=True) + if target_node.shape == source_node.shape: + return np.array(source_node, copy=True) + if target_node.size == source_node.size: + return np.array(source_node, copy=True).reshape(target_node.shape) + if target_node.shape != source_node.shape: + raise ValueError( + f"Shape mismatch at {path}: target {target_node.shape}, " + f"source {source_node.shape}" + ) + return np.array(source_node, copy=True) + + return deepcopy(target_node) + + +def _clear_jax_mesh_for_host_ops() -> None: + _set_nnx_eager_sharding(False) + _set_jax_mesh(Mesh(np.empty((), dtype=object), ())) + + +def _set_nnx_eager_sharding(enabled: bool) -> None: + use_eager_sharding = getattr(nnx, "use_eager_sharding", None) + if use_eager_sharding is not None: + use_eager_sharding(enabled) + + +def _set_jax_mesh(mesh: Mesh) -> None: + """Set the global mesh across JAX versions used by CI and users.""" + if _ACTIVE_JAX_MESH_CONTEXT: + _ACTIVE_JAX_MESH_CONTEXT.pop().__exit__(None, None, None) + + set_mesh = getattr(jax, "set_mesh", None) + if set_mesh is not None: + set_mesh(mesh) + return + + enter = getattr(mesh, "__enter__", None) + if enter is None or getattr(mesh, "__exit__", None) is None: + raise AttributeError("This JAX version cannot set a global mesh.") + + enter() + _ACTIVE_JAX_MESH_CONTEXT.append(mesh) + + +def _merge_batches_for_bias( + batch_list: list[np.ndarray], key: str +) -> np.ndarray | float: + arrays = [np.asarray(item) for item in batch_list] + if key.startswith("find_"): + return float(np.max(arrays)) + if key in {"natoms", "natoms_vec"}: + return np.stack(arrays, axis=0) + if arrays and arrays[0].ndim == 0: + return np.asarray(arrays) + return np.concatenate(arrays, axis=0) + + +def _infer_bias_frame_count(sampled: dict[str, np.ndarray | None]) -> int | None: + for key in ("coord", "atype", "energy", "force", "real_natoms_vec"): + value = sampled.get(key) + if value is None: + continue + array = np.asarray(value) + if array.ndim > 0: + return int(array.shape[0]) + return None + + +def _expand_bias_natoms( + value: np.ndarray | None, + nframes: int | None, +) -> np.ndarray | None: + if value is None or nframes is None: + return value + array = np.asarray(value) + if array.ndim == 0: + return array + if array.ndim == 1: + array = array.reshape(1, -1) + if array.shape[0] == nframes: + return array + if nframes % array.shape[0] != 0: + raise ValueError( + "Cannot expand natoms statistics from " + f"{array.shape[0]} rows to {nframes} frames." + ) + return np.repeat(array, nframes // array.shape[0], axis=0) + + +def _pack_data_for_bias_adjust( + train_data: DeepmdDataSystem, + nbatches: int, +) -> list[dict[str, np.ndarray | None]]: + all_stat = make_stat_input(train_data, nbatches, merge_sys=False) + all_stat["atype"] = all_stat.pop("type") + if "natoms_vec" in all_stat: + all_stat["natoms"] = all_stat["natoms_vec"] + sampled = [ + {kk: _merge_batches_for_bias(vv[ii], kk) for kk, vv in all_stat.items()} + for ii in range(train_data.get_nsystems()) + ] + for ii, single_data in enumerate(sampled): + for key, value in list(single_data.items()): + single_data[key] = to_numpy_array(value) + if not train_data.data_systems[ii].pbc: + single_data["box"] = None + nframes = _infer_bias_frame_count(single_data) + if "real_natoms_vec" in single_data: + single_data["real_natoms_vec"] = _expand_bias_natoms( + single_data["real_natoms_vec"], + nframes, + ) + single_data["natoms"] = single_data["real_natoms_vec"] + else: + single_data["natoms"] = _expand_bias_natoms( + single_data.get("natoms"), + nframes, + ) + if "natoms_vec" in single_data: + single_data["natoms_vec"] = _expand_bias_natoms( + single_data["natoms_vec"], + nframes, + ) + return sampled + + +def model_change_out_bias( + model: BaseModel, + sampled: list[dict[str, np.ndarray | None]], + bias_adjust_mode: str = "change-by-statistic", +) -> BaseModel: + old_bias = deepcopy(to_numpy_array(model.get_out_bias())) + try: + model.change_out_bias( + sampled, + bias_adjust_mode=bias_adjust_mode, + ) + except np.linalg.LinAlgError: + if bias_adjust_mode != "change-by-statistic": + raise + log.warning( + "change-by-statistic failed during JAX finetune bias adjustment; " + "falling back to set-by-statistic." + ) + model.change_out_bias( + sampled, + bias_adjust_mode="set-by-statistic", + ) + new_bias = deepcopy(to_numpy_array(model.get_out_bias())) + log.info( + "Change output bias of %s from %s to %s.", + model.get_type_map(), + old_bias.reshape(-1), + new_bias.reshape(-1), + ) + return model + + +def _build_loss( + loss_param: dict[str, Any], + starter_learning_rate: float, +) -> tuple[DefLoss, bool]: + loss_cfg = deepcopy(loss_param) + loss_cfg["starter_learning_rate"] = starter_learning_rate + loss_type = loss_cfg.get("type", "ener") + if loss_type != "ener": + raise RuntimeError("unknown loss type " + loss_type) + if loss_cfg.get("start_pref_h", 0.0) > 0.0: + return EnergyHessianLoss.get_loss(loss_cfg), True + return EnergyLoss.get_loss(loss_cfg), False + + +def _compute_single_data_stat(model: BaseModel, train_data: DeepmdDataSystem) -> None: + descriptor_stat, fitting_stat = _build_single_data_stat(train_data) + model.atomic_model.descriptor.compute_input_stats(descriptor_stat) + model.atomic_model.fitting.compute_output_stats( + fitting_stat, mixed_type=train_data.mixed_type + ) + + +def _build_single_data_stat( + train_data: DeepmdDataSystem, +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + data_stat_nbatch = 10 + all_stat = make_stat_input(train_data, data_stat_nbatch, merge_sys=False) + all_stat["atype"] = all_stat.pop("type") + all_stat_sys = [ + { + kk: jnp.asarray(np.concatenate(vv[ii], axis=0)) + for kk, vv in all_stat.items() + if not kk.startswith("find_") + } + for ii in range(train_data.get_nsystems()) + ] + for ii, single_data in enumerate(all_stat_sys): + if not train_data.data_systems[ii].pbc: + single_data["box"] = None + return all_stat_sys, all_stat + + +def _apply_weighted_shared_fitting_input_stats( + component: Any, + sampled_stats_by_branch: list[list[dict[str, Any]]], + branch_weights: list[float], + protection: float, +) -> None: + if getattr(component, "numb_fparam", 0) > 0: + weighted_sum = np.zeros(component.numb_fparam, dtype=np.float64) + weighted_sum_sq = np.zeros(component.numb_fparam, dtype=np.float64) + weighted_count = 0.0 + for sampled_stats, branch_weight in zip( + sampled_stats_by_branch, branch_weights, strict=True + ): + cat_data = np.concatenate( + [np.asarray(frame["fparam"]) for frame in sampled_stats], axis=0 + ) + cat_data = np.reshape(cat_data, [-1, component.numb_fparam]).astype( + np.float64 + ) + weighted_sum += branch_weight * np.sum(cat_data, axis=0) + weighted_sum_sq += branch_weight * np.sum(cat_data * cat_data, axis=0) + weighted_count += branch_weight * cat_data.shape[0] + fparam_avg = weighted_sum / weighted_count + fparam_std = np.sqrt( + np.maximum(weighted_sum_sq / weighted_count - fparam_avg**2, 0.0) + ) + fparam_std = np.where( + fparam_std < protection, + np.array(protection, dtype=fparam_std.dtype), + fparam_std, + ) + xp = array_api_compat.array_namespace(component.fparam_avg) + component.fparam_avg = xp.asarray( + fparam_avg, + dtype=component.fparam_avg.dtype, + device=array_api_compat.device(component.fparam_avg), + ) + component.fparam_inv_std = xp.asarray( + 1.0 / fparam_std, + dtype=component.fparam_inv_std.dtype, + device=array_api_compat.device(component.fparam_inv_std), + ) + + if getattr(component, "numb_aparam", 0) > 0: + weighted_sum = np.zeros(component.numb_aparam, dtype=np.float64) + weighted_sum_sq = np.zeros(component.numb_aparam, dtype=np.float64) + weighted_count = 0.0 + for sampled_stats, branch_weight in zip( + sampled_stats_by_branch, branch_weights, strict=True + ): + cat_data = np.concatenate( + [np.asarray(frame["aparam"]) for frame in sampled_stats], axis=0 + ) + cat_data = np.reshape(cat_data, [-1, component.numb_aparam]).astype( + np.float64 + ) + weighted_sum += branch_weight * np.sum(cat_data, axis=0) + weighted_sum_sq += branch_weight * np.sum(cat_data * cat_data, axis=0) + weighted_count += branch_weight * cat_data.shape[0] + aparam_avg = weighted_sum / weighted_count + aparam_std = np.sqrt( + np.maximum(weighted_sum_sq / weighted_count - aparam_avg**2, 0.0) + ) + aparam_std = np.where( + aparam_std < protection, + np.array(protection, dtype=aparam_std.dtype), + aparam_std, + ) + xp = array_api_compat.array_namespace(component.aparam_avg) + component.aparam_avg = xp.asarray( + aparam_avg, + dtype=component.aparam_avg.dtype, + device=array_api_compat.device(component.aparam_avg), + ) + component.aparam_inv_std = xp.asarray( + 1.0 / aparam_std, + dtype=component.aparam_inv_std.dtype, + device=array_api_compat.device(component.aparam_inv_std), + ) + + +def _compute_multitask_data_stat( + model: ModelWrapper, + train_data: dict[str, DeepmdDataSystem], + model_key_prob_map: dict[str, float], + data_stat_protect_map: dict[str, float], +) -> None: + stat_cache = { + model_key: _build_single_data_stat(train_data[model_key]) + for model_key in model.keys() + } + + descriptor_groups: dict[int, dict[str, Any]] = {} + fitting_input_groups: dict[tuple[int, int | None, int | None], dict[str, Any]] = {} + for model_key in model.keys(): + branch_model = model[model_key] + branch_fitting = branch_model.atomic_model.fitting + descriptor_id = id(branch_model.atomic_model.descriptor) + descriptor_groups.setdefault( + descriptor_id, + { + "component": branch_model.atomic_model.descriptor, + "stats": [], + }, + )["stats"].append(stat_cache[model_key][0]) + fitting_group_id = ( + id(branch_fitting.nets), + id(branch_fitting.fparam_avg) + if getattr(branch_fitting, "numb_fparam", 0) > 0 + else None, + id(branch_fitting.aparam_avg) + if getattr(branch_fitting, "numb_aparam", 0) > 0 + else None, + ) + fitting_input_groups.setdefault( + fitting_group_id, + { + "component": branch_fitting, + "mixed_type": train_data[model_key].mixed_type, + "weights": [], + "samples": [], + "data_stat_protect": data_stat_protect_map[model_key], + }, + ) + if ( + fitting_input_groups[fitting_group_id]["mixed_type"] + != train_data[model_key].mixed_type + ): + raise ValueError( + "All branches sharing a fitting_net must use the same mixed_type setting." + ) + if not np.isclose( + fitting_input_groups[fitting_group_id]["data_stat_protect"], + data_stat_protect_map[model_key], + ): + raise ValueError( + "Model key 'data_stat_protect' must be the same in each branch when multitask!" + ) + fitting_input_groups[fitting_group_id]["weights"].append( + model_key_prob_map[model_key] + ) + fitting_input_groups[fitting_group_id]["samples"].append( + stat_cache[model_key][0] + ) + + for descriptor_group in descriptor_groups.values(): + merged_descriptor_stat = [] + for single_stat in descriptor_group["stats"]: + merged_descriptor_stat.extend(single_stat) + descriptor_group["component"].compute_input_stats(merged_descriptor_stat) + + for fitting_group in fitting_input_groups.values(): + _apply_weighted_shared_fitting_input_stats( + fitting_group["component"], + fitting_group["samples"], + fitting_group["weights"], + protection=fitting_group["data_stat_protect"], + ) + + for model_key in model.keys(): + branch_fitting = model[model_key].atomic_model.fitting + branch_fitting.compute_output_stats( + stat_cache[model_key][1], + mixed_type=train_data[model_key].mixed_type, + ) + + +def _resolve_model_prob_multi( + model_keys: list[str], + training_params: dict[str, Any], + train_data: dict[str, DeepmdDataSystem], +) -> tuple[np.ndarray, int]: + num_steps = training_params.get("numb_steps") + num_epoch_dict = training_params.get("num_epoch_dict", {}) + if num_epoch_dict: + if num_steps is not None: + raise ValueError( + "training.numb_steps and training.num_epoch_dict are mutually exclusive." + ) + per_task_total = [] + for model_key in model_keys: + sampler_weights = np.asarray( + train_data[model_key].sys_probs, dtype=np.float64 + ) + per_task_total.append( + compute_total_numb_batch( + train_data[model_key].nbatches, sampler_weights + ) + ) + model_prob, resolved_num_steps, _ = resolve_model_prob_from_epochs( + model_keys, + num_epoch_dict, + np.asarray(per_task_total, dtype=np.float64), + ) + return model_prob, resolved_num_steps + + if num_steps is None: + raise ValueError( + "Either training.numb_steps (multi-task only) or training.num_epoch_dict must be set." + ) + model_prob_config = training_params.get("model_prob", {}) + if model_prob_config: + missing = [k for k in model_keys if k not in model_prob_config] + if missing: + raise ValueError( + f"training.model_prob must specify all tasks; missing: {missing}" + ) + model_prob = np.asarray( + [float(model_prob_config[k]) for k in model_keys], dtype=np.float64 + ) + else: + model_prob = np.asarray( + [float(train_data[k].get_nsystems()) for k in model_keys], dtype=np.float64 + ) + if np.any(model_prob < 0.0) or not np.all(np.isfinite(model_prob)): + raise ValueError("training.model_prob must be non-negative and finite.") + prob_sum = float(np.sum(model_prob)) + if prob_sum <= 0.0: + raise ValueError("training.model_prob must sum to a positive value.") + return model_prob / prob_sum, int(num_steps) class DPTrainer: def __init__( self, jdata: dict, - init_model: Optional[str] = None, - restart: Optional[str] = None, + init_model: str | None = None, + restart: str | None = None, + init_frz_model: str | None = None, + finetune_model: str | None = None, + force_load: bool = False, + shared_links: dict[str, Any] | None = None, + finetune_links: dict[str, FinetuneRuleItem] | None = None, + finetune_model_data: dict[str, Any] | None = None, ) -> None: self.init_model = init_model self.restart = restart - self.model_def_script = jdata["model"] + self.init_frz_model = init_frz_model + self.finetune_model = finetune_model + self.force_load = force_load + self.shared_links = shared_links or {} + self.finetune_links = finetune_links or {} + self.finetune_model_data = finetune_model_data + self.model_def_script = deepcopy(jdata["model"]) + if "model_dict" in self.model_def_script and self.shared_links: + self.model_def_script["shared_links"] = deepcopy(self.shared_links) + self.training_param = jdata["training"] + self.num_steps = self.training_param.get("numb_steps") self.start_step = 0 + self.multi_task = "model_dict" in jdata["model"] + self.model_keys = ( + list(jdata["model"]["model_dict"]) if self.multi_task else ["Default"] + ) + if self.multi_task: + _, self.case_embd_index = get_case_embd_config(jdata["model"]) + else: + self.case_embd_index = {} + + learning_rate_param = deepcopy(jdata["learning_rate"]) + self.learning_rate_param = learning_rate_param + self.lr = LearningRateExp( + **learning_rate_param, + num_steps=self.num_steps or 1, + ) + + if self.multi_task: + self.loss: dict[str, DefLoss] = {} + self.branch_has_hessian: dict[str, bool] = {} + for model_key in self.model_keys: + loss_param = jdata["loss_dict"][model_key] + self.loss[model_key], self.branch_has_hessian[model_key] = _build_loss( + loss_param, + learning_rate_param["start_lr"], + ) + else: + self.loss, has_hessian = _build_loss( + jdata.get("loss", {}), + learning_rate_param["start_lr"], + ) + self.branch_has_hessian = {"Default": has_hessian} + + self.model: DefModel = get_model_for_wrapper( + jdata["model"], + shared_links=self.shared_links, + ) + self._apply_hessian_flags(self.model) + if self.init_model is not None: model_dict = serialize_from_file(self.init_model) - self.model = BaseModel.deserialize(model_dict["model"]) + self._load_model_data(model_dict) elif self.restart is not None: model_dict = serialize_from_file(self.restart) - self.model = BaseModel.deserialize(model_dict["model"]) + self._load_model_data(model_dict) + self.model_def_script = deepcopy(model_dict["model_def_script"]) self.start_step = model_dict["@variables"].get("current_step", 0) - else: - # from scratch - self.model = get_model(jdata["model"]) - self.training_param = jdata["training"] - self.num_steps = self.training_param["numb_steps"] - - def get_lr_and_coef(lr_param: dict) -> LearningRateExp: - lr_type = lr_param.get("type", "exp") - if lr_type == "exp": - lr = LearningRateExp( - **lr_param, - num_steps=self.num_steps, - ) - else: - raise RuntimeError("unknown learning_rate type " + lr_type) - return lr - - learning_rate_param = jdata["learning_rate"] - self.lr = get_lr_and_coef(learning_rate_param) - loss_param = jdata.get("loss", {}) - loss_param["starter_learning_rate"] = learning_rate_param["start_lr"] - - loss_type = loss_param.get("type", "ener") - if loss_type == "ener" and loss_param.get("start_pref_h", 0.0) > 0.0: - self.loss = EnergyHessianLoss.get_loss(loss_param) - self.model.enable_hessian() - elif loss_type == "ener": - self.loss = EnergyLoss.get_loss(loss_param) - else: - raise RuntimeError("unknown loss type " + loss_type) + if self.init_frz_model is not None: + frozen_model_data = serialize_from_file(self.init_frz_model) + self._load_frozen_model_data(frozen_model_data) - # training - tr_data = jdata["training"] + tr_data = self.training_param self.disp_file = tr_data.get("disp_file", "lcurve.out") self.disp_freq = tr_data.get("disp_freq", 1000) self.save_freq = tr_data.get("save_freq", 1000) @@ -141,70 +703,394 @@ def get_lr_and_coef(lr_param: dict) -> LearningRateExp: self.change_bias_after_training = tr_data.get( "change_bias_after_training", False ) - self.numb_fparam = self.model.get_dim_fparam() - - if tr_data.get("validation_data", None) is not None: - self.valid_numb_batch = tr_data["validation_data"].get("numb_btch", 1) + if self.multi_task: + self.data_bias_nsample = { + model_key: jdata["model"]["model_dict"][model_key].get( + "data_bias_nsample", 10 + ) + for model_key in self.model_keys + } else: - self.valid_numb_batch = 1 - - # if init the graph with the frozen model - self.frz_model = None + self.data_bias_nsample = self.model_def_script.get("data_bias_nsample", 10) + self.model_prob = None self.ckpt_meta = None self.model_type = None + def _apply_hessian_flags(self, model: DefModel) -> None: + if self.multi_task: + assert isinstance(model, ModelWrapper) + for model_key in self.model_keys: + if self.branch_has_hessian[model_key]: + model[model_key].enable_hessian() + else: + if isinstance(model, ModelWrapper): + raise TypeError("single-task JAX trainer expected a single-task model.") + if self.branch_has_hessian["Default"]: + model.enable_hessian() + + def _load_model_data(self, model_data: dict[str, Any]) -> None: + serialized_model = model_data["model"] + if self.force_load: + missing: list[tuple[Any, ...]] = [] + unexpected: list[tuple[Any, ...]] = [] + serialized_model = _merge_init_frz_model_data( + self.model.serialize(), + serialized_model, + missing=missing, + unexpected=unexpected, + ) + if missing or unexpected: + log.warning( + "Checkpoint loaded in force_load mode. Missing keys reinitialized: %s; Unexpected keys ignored: %s", + [".".join(map(str, item)) for item in missing[:20]], + [".".join(map(str, item)) for item in unexpected[:20]], + ) + effective_shared_links = self.shared_links or model_data.get( + "model_def_script", {} + ).get("shared_links", {}) + if self.multi_task: + if "model_dict" not in serialized_model: + raise ValueError( + "init_model/restart for JAX multitask target requires a multitask checkpoint." + ) + self.model = ModelWrapper.deserialize( + serialized_model, + shared_links=effective_shared_links, + case_embd_index=self.case_embd_index, + ) + else: + if "model_dict" in serialized_model: + raise ValueError( + "init_model/restart for single-task JAX target does not accept a multitask checkpoint." + ) + self.model = BaseModel.deserialize(serialized_model) + self._apply_hessian_flags(self.model) + + def _validate_shared_finetune_rules(self) -> None: + if not self.multi_task or not self.shared_links or not self.finetune_links: + return + for shared_key, link_info in self.shared_links.items(): + for link in link_info.get("links", []): + shared_level = int(link.get("shared_level", 0)) + if link["shared_type"] == "fitting_net" and shared_level != 0: + raise NotImplementedError( + "JAX multitask finetune fitting_net sharing only supports shared_level=0, " + f"but got '{shared_key}' at level {shared_level}." + ) + + def _load_frozen_model_data(self, model_data: dict[str, Any]) -> None: + missing: list[tuple[Any, ...]] = [] + unexpected: list[tuple[Any, ...]] = [] + merged_model_data = _merge_init_frz_model_data( + self.model.serialize(), + model_data["model"], + missing=missing, + unexpected=unexpected, + ) + if self.multi_task: + self.model = ModelWrapper.deserialize( + merged_model_data, + shared_links=self.shared_links, + case_embd_index=self.case_embd_index, + ) + else: + self.model = BaseModel.deserialize(merged_model_data) + self._apply_hessian_flags(self.model) + if missing or unexpected: + log.warning( + "Frozen model loaded non-strictly. Missing keys: %s, Unexpected keys: %s", + [".".join(map(str, item)) for item in missing[:20]], + [".".join(map(str, item)) for item in unexpected[:20]], + ) + + @staticmethod + def _shared_type_to_serialized_paths( + shared_type: str, + shared_level: int, + source_model_data: dict[str, Any], + ) -> list[tuple[Any, ...]]: + def descriptor_paths( + path_prefix: tuple[Any, ...], + descriptor_data: dict[str, Any], + ) -> list[tuple[Any, ...]]: + descriptor_type = descriptor_data.get("type") + if shared_level == 0: + return [path_prefix] + if shared_level == 1 and descriptor_type in { + "dpa1", + "dpa2", + "dpa3", + "se_e3_tebd", + }: + return [(*path_prefix, "type_embedding")] + raise NotImplementedError( + "JAX multitask finetune does not support shared override for " + f"{shared_type}:{shared_level} with descriptor type {descriptor_type}." + ) + + if shared_type == "descriptor": + return descriptor_paths(("descriptor",), source_model_data["descriptor"]) + if shared_type.startswith("descriptor_hybrid_"): + idx = int(shared_type.rsplit("_", 1)[1]) + return descriptor_paths( + ("descriptor", "list", idx), + source_model_data["descriptor"]["list"][idx], + ) + if shared_type == "fitting_net": + if shared_level != 0: + raise NotImplementedError( + "JAX multitask finetune fitting_net sharing only supports shared_level=0, " + f"but got {shared_level}." + ) + fitting_paths = [("fitting", "nets")] + fitting_vars = source_model_data.get("fitting", {}).get("@variables", {}) + for key in fitting_vars: + if key not in {"bias_atom_e", "case_embd"}: + fitting_paths.append(("fitting", "@variables", key)) + return fitting_paths + return [] + + def _collect_shared_source_overrides( + self, + source_multi: bool, + ) -> dict[str, dict[tuple[Any, ...], dict[str, Any]]]: + overrides: dict[str, dict[tuple[Any, ...], dict[str, Any]]] = { + model_key: {} for model_key in self.model_keys + } + if ( + not self.shared_links + or not self.finetune_links + or self.finetune_model_data is None + ): + return overrides + + for _, link_info in self.shared_links.items(): + shareable_links = [ + link + for link in link_info.get("links", []) + if ( + link["shared_type"].startswith("descriptor") + or link["shared_type"] == "fitting_net" + ) + ] + if not shareable_links: + continue + base_link = shareable_links[0] + canonical_source_key = self.finetune_links[ + base_link["model_key"] + ].get_model_branch() + canonical_source_model_data = ( + self.finetune_model_data["model"]["model_dict"][canonical_source_key] + if source_multi + else self.finetune_model_data["model"] + ) + for link in shareable_links[1:]: + current_source_key = self.finetune_links[ + link["model_key"] + ].get_model_branch() + if current_source_key == canonical_source_key: + continue + paths = self._shared_type_to_serialized_paths( + link["shared_type"], + int(link.get("shared_level", 0)), + canonical_source_model_data, + ) + for path in paths: + overrides[link["model_key"]][path] = canonical_source_model_data + return overrides + @property - def data_requirements(self) -> list[DataRequirementItem]: + def data_requirements( + self, + ) -> list[DataRequirementItem] | dict[str, list[DataRequirementItem]]: + if self.multi_task: + return { + model_key: self.loss[model_key].label_requirement + for model_key in self.model_keys + } return self.loss.label_requirement + def _apply_single_finetune( + self, + target_model: BaseModel, + pretrained_model_data: dict[str, Any], + finetune_rule: FinetuneRuleItem, + *, + source_overrides: dict[tuple[Any, ...], dict[str, Any]] | None = None, + ) -> BaseModel: + if self.force_load: + missing: list[tuple[Any, ...]] = [] + unexpected: list[tuple[Any, ...]] = [] + pretrained_model_data = _merge_init_frz_model_data( + target_model.serialize(), + pretrained_model_data, + missing=missing, + unexpected=unexpected, + ) + if missing or unexpected: + log.warning( + "Finetune checkpoint loaded in force_load mode. Missing keys reinitialized: %s; Unexpected keys ignored: %s", + [".".join(map(str, item)) for item in missing[:20]], + [".".join(map(str, item)) for item in unexpected[:20]], + ) + pretrained_model = BaseModel.deserialize(pretrained_model_data) + if finetune_rule.get_update_type(): + pretrained_model.change_type_map( + target_model.get_type_map(), + model_with_new_type_stat=target_model.atomic_model, + ) + merged_model_data = merge_finetune_model_data( + target_model.serialize(), + pretrained_model.serialize(), + finetune_rule, + source_overrides=source_overrides, + ) + return BaseModel.deserialize(merged_model_data) + + def _finetune_single(self, train_data: DeepmdDataSystem) -> None: + if isinstance(self.model, ModelWrapper): + raise TypeError("single-task JAX finetune expected a single-task model.") + if self.finetune_model_data is None: + self.finetune_model_data = serialize_from_file(self.finetune_model) + finetune_rule = self.finetune_links["Default"] + pretrained_data = self.finetune_model_data + if "model_dict" in pretrained_data.get("model_def_script", {}): + pretrained_data = select_model_branch( + pretrained_data, + finetune_rule.get_model_branch(), + ) + self.model = self._apply_single_finetune( + self.model, + pretrained_data["model"], + finetune_rule, + ) + self._apply_hessian_flags(self.model) + self.model = model_change_out_bias( + self.model, + _pack_data_for_bias_adjust(train_data, self.data_bias_nsample), + bias_adjust_mode=( + "set-by-statistic" + if finetune_rule.get_random_fitting() + else "change-by-statistic" + ), + ) + + def _finetune_multi(self, train_data: dict[str, DeepmdDataSystem]) -> None: + if not isinstance(self.model, ModelWrapper): + raise TypeError("multitask JAX finetune expected a multitask model.") + self._validate_shared_finetune_rules() + if self.finetune_model_data is None: + self.finetune_model_data = serialize_from_file(self.finetune_model) + source_multi = "model_dict" in self.finetune_model_data.get("model", {}) + shared_source_overrides = self._collect_shared_source_overrides(source_multi) + merged_branch_models: dict[str, Any] = {} + for model_key in self.model_keys: + branch_model = self.model[model_key] + finetune_rule = self.finetune_links[model_key] + source_key = finetune_rule.get_model_branch() + source_model_data = ( + self.finetune_model_data["model"]["model_dict"][source_key] + if source_multi + else self.finetune_model_data["model"] + ) + merged_branch_model = self._apply_single_finetune( + branch_model, + source_model_data, + finetune_rule, + source_overrides=shared_source_overrides.get(model_key), + ) + if self.branch_has_hessian[model_key]: + merged_branch_model.enable_hessian() + if not finetune_rule.get_resuming(): + log.info( + "Model branch %s will be fine-tuned. This may take a long time...", + model_key, + ) + merged_branch_model = model_change_out_bias( + merged_branch_model, + _pack_data_for_bias_adjust( + train_data[model_key], self.data_bias_nsample[model_key] + ), + bias_adjust_mode=( + "set-by-statistic" + if finetune_rule.get_random_fitting() + else "change-by-statistic" + ), + ) + else: + log.info("Model branch %s will resume training.", model_key) + merged_branch_models[model_key] = merged_branch_model.serialize() + self.model = ModelWrapper.deserialize( + {"model_dict": merged_branch_models}, + shared_links=self.shared_links, + case_embd_index=self.case_embd_index, + ) + self._apply_hessian_flags(self.model) + def train( - self, train_data: DeepmdDataSystem, valid_data: DeepmdDataSystem | None = None + self, + train_data: DeepmdDataSystem | dict[str, DeepmdDataSystem], + valid_data: DeepmdDataSystem | dict[str, DeepmdDataSystem | None] | None = None, + ) -> None: + if self.multi_task: + assert isinstance(train_data, dict) + valid_data = valid_data if isinstance(valid_data, dict) else {} + self._train_multi(train_data, valid_data) + else: + assert isinstance(train_data, DeepmdDataSystem) + valid_data = ( + valid_data if isinstance(valid_data, DeepmdDataSystem) else None + ) + self._train_single(train_data, valid_data) + + def _train_single( + self, + train_data: DeepmdDataSystem, + valid_data: DeepmdDataSystem | None = None, ) -> None: model = self.model + if isinstance(model, ModelWrapper): + raise TypeError("single-task JAX training expected a single-task model.") + _clear_jax_mesh_for_host_ops() tx = optax.adam( learning_rate=lambda step: self.lr.value(self.start_step + step), ) - optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param) - # data stat - if self.init_model is None and self.restart is None: - data_stat_nbatch = 10 # TODO - all_stat = make_stat_input(train_data, data_stat_nbatch, merge_sys=False) - all_stat["atype"] = all_stat.pop("type") - - # swap dict key and list idx - all_stat_sys = [ - { - kk: jnp.asarray(np.concatenate(vv[ii], axis=0)) - for kk, vv in all_stat.items() - if not kk.startswith("find_") - } - for ii in range(train_data.get_nsystems()) - ] - for ii, single_data in enumerate(all_stat_sys): - if not train_data.data_systems[ii].pbc: - single_data["box"] = None - model.atomic_model.descriptor.compute_input_stats(all_stat_sys) - model.atomic_model.fitting.compute_output_stats( - all_stat, mixed_type=train_data.mixed_type - ) + finetune_rule = self.finetune_links.get("Default") + finetune_has_new_type = ( + self.finetune_model is not None + and finetune_rule is not None + and finetune_rule.get_has_new_type() + ) + if ( + self.init_model is None + and self.restart is None + and (self.finetune_model is None or finetune_has_new_type) + ): + _compute_single_data_stat(model, train_data) + + if self.finetune_model is not None: + self._finetune_single(train_data) + model = self.model + if isinstance(model, ModelWrapper): + raise TypeError( + "single-task JAX finetune produced a multitask model unexpectedly." + ) - # parallel auto_mesh = jax.make_mesh( - ( - jax.process_count(), - jax.local_device_count(), - ), + (jax.process_count(), jax.local_device_count()), ("data", "natoms"), ) - nnx.use_eager_sharding(True) - jax.set_mesh(auto_mesh) - if int(os.environ.get("DP_JAX_MULTI_NPROC", "0")) > 1: - sharding = NamedSharding(auto_mesh, P("data")) - else: - sharding = None - # a hack to apply the sharding to all parameters + _set_nnx_eager_sharding(True) + _set_jax_mesh(auto_mesh) + sharding = ( + NamedSharding(auto_mesh, P("data")) + if int(os.environ.get("DP_JAX_MULTI_NPROC", "0")) > 1 + else None + ) model = BaseModel.deserialize(model.serialize()) + self._apply_hessian_flags(model) + optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param) def loss_fn( model: BaseModel, @@ -231,7 +1117,7 @@ def loss_fn( mapping, do_atomic_virial=False, ) - loss, more_loss = self.loss( + loss, _ = self.loss( learning_rate=lr, natoms=label_dict["coord"].shape[1], model_dict=model_dict, @@ -265,7 +1151,7 @@ def loss_fn_more_loss( mapping, do_atomic_virial=False, ) - loss, more_loss = self.loss( + _, more_loss = self.loss( learning_rate=lr, natoms=label_dict["coord"].shape[1], model_dict=model_dict, @@ -306,8 +1192,11 @@ def train_step( disp_file_fp = open(self.disp_file, "w") for step in range(self.start_step, self.num_steps): batch_data = train_data.get_batch() - # numpy to jax - jax_data = convert_numpy_data_to_jax_data(batch_data, sharding) + jax_data = convert_numpy_data_to_jax_data( + batch_data, + sharding, + natoms_axis_size=auto_mesh.shape.get("natoms", 1), + ) extended_coord, extended_atype, nlist, mapping, fp, ap = prepare_input( rcut=model.get_rcut(), sel=model.get_sel(), @@ -333,12 +1222,7 @@ def train_step( step == 0 or (step + 1) % self.disp_freq == 0 ): wall_time = time.time() - start_time - log.info( - format_training_message( - batch=step + 1, - wall_time=wall_time, - ) - ) + log.info(format_training_message(batch=step + 1, wall_time=wall_time)) more_loss = loss_fn_more_loss( model, self.lr.value(step), @@ -353,7 +1237,9 @@ def train_step( if valid_data is not None: valid_batch_data = valid_data.get_batch() jax_valid_data = convert_numpy_data_to_jax_data( - valid_batch_data, sharding + valid_batch_data, + sharding, + natoms_axis_size=auto_mesh.shape.get("natoms", 1), ) extended_coord, extended_atype, nlist, mapping, fp, ap = ( prepare_input( @@ -382,44 +1268,379 @@ def train_step( else: valid_more_loss = None if step == 0: - self.print_header( - disp_file_fp, - train_results=more_loss, - valid_results=valid_more_loss, - ) + self.print_header(disp_file_fp, more_loss, valid_more_loss) self.print_on_training( disp_file_fp, - train_results=more_loss, - valid_results=valid_more_loss, + more_loss, + valid_more_loss, cur_batch=step + 1, cur_lr=self.lr.value(step), ) start_time = time.time() if (step + 1) % self.save_freq == 0: - # save model - _, state = nnx.split(model) - ckpt_path = Path(f"{self.save_ckpt}-{step + 1}.jax") - if ckpt_path.is_dir(): - # remove old checkpoint if it exists - shutil.rmtree(ckpt_path) - model_def_script_cpy = self.model_def_script.copy() - model_def_script_cpy["current_step"] = step + 1 - with ocp.Checkpointer( - ocp.CompositeCheckpointHandler("state", "model_def_script") - ) as checkpointer: - checkpointer.save( - ckpt_path.absolute(), - ocp.args.Composite( - state=ocp.args.StandardSave(state.to_pure_dict()), - model_def_script=ocp.args.JsonSave(model_def_script_cpy), - ), + self._save_checkpoint(model, step + 1) + log.info( + f"Trained model has been saved to: {Path(f'{self.save_ckpt}-{step + 1}.jax')!s}" + ) + disp_file_fp.close() + self.model = model + + def _train_multi( + self, + train_data: dict[str, DeepmdDataSystem], + valid_data: dict[str, DeepmdDataSystem | None], + ) -> None: + model = self.model + assert isinstance(model, ModelWrapper) + _clear_jax_mesh_for_host_ops() + self.model_prob, self.num_steps = _resolve_model_prob_multi( + self.model_keys, + self.training_param, + train_data, + ) + finetune_has_new_type = self.finetune_model is not None and any( + rule.get_has_new_type() for rule in self.finetune_links.values() + ) + if self.init_model is None and self.restart is None: + if self.finetune_model is None or finetune_has_new_type: + data_stat_protect_map = { + model_key: float( + self.model_def_script["model_dict"][model_key].get( + "data_stat_protect", 1e-2 + ) ) - log.info(f"Trained model has been saved to: {ckpt_path!s}") - symlink_prefix_files(f"{self.save_ckpt}-{step + 1}", self.save_ckpt) - with open("checkpoint", "w") as fp: - fp.write(f"{self.save_ckpt}.jax") + for model_key in self.model_keys + } + _compute_multitask_data_stat( + model, + train_data, + dict(zip(self.model_keys, self.model_prob, strict=True)), + data_stat_protect_map, + ) + if self.finetune_model is not None: + self._finetune_multi(train_data) + model = self.model + assert isinstance(model, ModelWrapper) + self.lr = LearningRateExp( + **self.learning_rate_param, + num_steps=self.num_steps, + ) + tx = optax.adam( + learning_rate=lambda step: self.lr.value(self.start_step + step), + ) + auto_mesh = jax.make_mesh( + (jax.process_count(), jax.local_device_count()), + ("data", "natoms"), + ) + _set_nnx_eager_sharding(True) + _set_jax_mesh(auto_mesh) + sharding = ( + NamedSharding(auto_mesh, P("data")) + if int(os.environ.get("DP_JAX_MULTI_NPROC", "0")) > 1 + else None + ) + model = ModelWrapper.deserialize( + model.serialize(), + shared_links=self.shared_links, + case_embd_index=self.case_embd_index, + ) + self._apply_hessian_flags(model) + optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param) + + loss_fns = {} + more_loss_fns = {} + train_step_fns = {} + for model_key in self.model_keys: + branch_loss = self.loss[model_key] + + def make_loss_fn( + task_key: str, task_loss: DefLoss + ) -> Callable[..., jnp.ndarray]: + def loss_fn( + wrapper: ModelWrapper, + lr: float, + label_dict: dict[str, jnp.ndarray], + extended_coord: jnp.ndarray, + extended_atype: jnp.ndarray, + nlist: jnp.ndarray, + mapping: jnp.ndarray | None, + fp: jnp.ndarray | None, + ap: jnp.ndarray | None, + ) -> jnp.ndarray: + branch_model = wrapper[task_key] + model_dict_lower = branch_model.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fp, + ap, + ) + model_dict = communicate_extended_output( + model_dict_lower, + branch_model.model_output_def(), + mapping, + do_atomic_virial=False, + ) + loss, _ = task_loss( + learning_rate=lr, + natoms=label_dict["coord"].shape[1], + model_dict=model_dict, + label_dict=label_dict, + ) + return loss + + return loss_fn + + def make_more_loss_fn( + task_key: str, task_loss: DefLoss + ) -> Callable[..., dict[str, jnp.ndarray]]: + @nnx.jit + def more_loss_fn( + wrapper: ModelWrapper, + lr: float, + label_dict: dict[str, jnp.ndarray], + extended_coord: jnp.ndarray, + extended_atype: jnp.ndarray, + nlist: jnp.ndarray, + mapping: jnp.ndarray | None, + fp: jnp.ndarray | None, + ap: jnp.ndarray | None, + ) -> dict[str, jnp.ndarray]: + branch_model = wrapper[task_key] + model_dict_lower = branch_model.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fp, + ap, + ) + model_dict = communicate_extended_output( + model_dict_lower, + branch_model.model_output_def(), + mapping, + do_atomic_virial=False, + ) + _, more_loss = task_loss( + learning_rate=lr, + natoms=label_dict["coord"].shape[1], + model_dict=model_dict, + label_dict=label_dict, + ) + return more_loss + + return more_loss_fn + + def make_train_step( + task_loss_fn: Callable[..., jnp.ndarray], + ) -> Callable[..., None]: + @nnx.jit + def train_step( + wrapper: ModelWrapper, + optimizer: nnx.Optimizer, + lr: float, + label_dict: dict[str, jnp.ndarray], + extended_coord: jnp.ndarray, + extended_atype: jnp.ndarray, + nlist: jnp.ndarray, + mapping: jnp.ndarray | None, + fp: jnp.ndarray | None, + ap: jnp.ndarray | None, + ) -> None: + grads = nnx.grad(task_loss_fn)( + wrapper, + lr, + label_dict, + extended_coord, + extended_atype, + nlist, + mapping, + fp, + ap, + ) + if Version(flax_version) >= Version("0.11.0"): + optimizer.update(wrapper, grads) + else: + optimizer.update(grads) + + return train_step + + loss_fns[model_key] = make_loss_fn(model_key, branch_loss) + more_loss_fns[model_key] = make_more_loss_fn(model_key, branch_loss) + train_step_fns[model_key] = make_train_step(loss_fns[model_key]) + + start_time = time.time() + disp_file_fp = open(self.disp_file, "w") + for step in range(self.start_step, self.num_steps): + model_index = dp_random.choice(len(self.model_keys), p=self.model_prob) + task_key = self.model_keys[model_index] + model.set_case_embd(task_key) + batch_data = train_data[task_key].get_batch() + jax_data = convert_numpy_data_to_jax_data( + batch_data, + sharding, + natoms_axis_size=auto_mesh.shape.get("natoms", 1), + ) + branch_model = model[task_key] + extended_coord, extended_atype, nlist, mapping, fp, ap = prepare_input( + rcut=branch_model.get_rcut(), + sel=branch_model.get_sel(), + coord=jax_data["coord"], + atype=jax_data["type"], + box=jax_data["box"] if jax_data["default_mesh"].size > 1 else None, + fparam=jax_data.get("fparam", None), + aparam=jax_data.get("aparam", None), + ) + train_step_fns[task_key]( + model, + optimizer, + self.lr.value(step), + jax_data, + extended_coord, + extended_atype, + nlist, + mapping, + fp, + ap, + ) + if self.display_in_training and ( + step == 0 or (step + 1) % self.disp_freq == 0 + ): + wall_time = time.time() - start_time + log.info(format_training_message(batch=step + 1, wall_time=wall_time)) + train_results = {_key: {} for _key in self.model_keys} + valid_results = {_key: {} for _key in self.model_keys} + model.set_case_embd(task_key) + train_results[task_key] = more_loss_fns[task_key]( + model, + self.lr.value(step), + jax_data, + extended_coord, + extended_atype, + nlist, + mapping, + fp, + ap, + ) + for _key in self.model_keys: + if _key != task_key: + train_batch_data = train_data[_key].get_batch() + jax_train_data = convert_numpy_data_to_jax_data( + train_batch_data, + sharding, + natoms_axis_size=auto_mesh.shape.get("natoms", 1), + ) + branch_model = model[_key] + ( + train_extended_coord, + train_extended_atype, + train_nlist, + train_mapping, + train_fp, + train_ap, + ) = prepare_input( + rcut=branch_model.get_rcut(), + sel=branch_model.get_sel(), + coord=jax_train_data["coord"], + atype=jax_train_data["type"], + box=jax_train_data["box"] + if jax_train_data["default_mesh"].size > 1 + else None, + fparam=jax_train_data.get("fparam", None), + aparam=jax_train_data.get("aparam", None), + ) + model.set_case_embd(_key) + train_results[_key] = more_loss_fns[_key]( + model, + self.lr.value(step), + jax_train_data, + train_extended_coord, + train_extended_atype, + train_nlist, + train_mapping, + train_fp, + train_ap, + ) + if valid_data.get(_key) is not None: + valid_batch_data = valid_data[_key].get_batch() + jax_valid_data = convert_numpy_data_to_jax_data( + valid_batch_data, + sharding, + natoms_axis_size=auto_mesh.shape.get("natoms", 1), + ) + branch_model = model[_key] + ( + valid_extended_coord, + valid_extended_atype, + valid_nlist, + valid_mapping, + valid_fp, + valid_ap, + ) = prepare_input( + rcut=branch_model.get_rcut(), + sel=branch_model.get_sel(), + coord=jax_valid_data["coord"], + atype=jax_valid_data["type"], + box=jax_valid_data["box"] + if jax_valid_data["find_box"] + else None, + fparam=jax_valid_data.get("fparam", None), + aparam=jax_valid_data.get("aparam", None), + ) + model.set_case_embd(_key) + valid_results[_key] = more_loss_fns[_key]( + model, + self.lr.value(step), + jax_valid_data, + valid_extended_coord, + valid_extended_atype, + valid_nlist, + valid_mapping, + valid_fp, + valid_ap, + ) + if step == 0: + self.print_header_multitask( + disp_file_fp, train_results, valid_results + ) + self.print_on_training_multitask( + disp_file_fp, + train_results, + valid_results, + cur_batch=step + 1, + cur_lr=self.lr.value(step), + ) + start_time = time.time() + if (step + 1) % self.save_freq == 0: + self._save_checkpoint(model, step + 1) + log.info( + f"Trained model has been saved to: {Path(f'{self.save_ckpt}-{step + 1}.jax')!s}" + ) disp_file_fp.close() + self.model = model + + def _save_checkpoint(self, model: DefModel, step: int) -> None: + _, state = nnx.split(model) + ckpt_path = Path(f"{self.save_ckpt}-{step}.jax") + if ckpt_path.is_dir(): + shutil.rmtree(ckpt_path) + model_def_script_cpy = deepcopy(self.model_def_script) + model_def_script_cpy["current_step"] = step + with ocp.Checkpointer( + ocp.CompositeCheckpointHandler("state", "model_def_script") + ) as checkpointer: + checkpointer.save( + ckpt_path.absolute(), + ocp.args.Composite( + state=ocp.args.StandardSave(state.to_pure_dict()), + model_def_script=ocp.args.JsonSave(model_def_script_cpy), + ), + ) + symlink_prefix_files(f"{self.save_ckpt}-{step}", self.save_ckpt) + with open("checkpoint", "w") as fp: + fp.write(f"{self.save_ckpt}.jax") @staticmethod def print_on_training( @@ -429,17 +1650,15 @@ def print_on_training( cur_batch: int, cur_lr: float, ) -> None: - print_str = "" - print_str += f"{cur_batch:7d}" + print_str = f"{cur_batch:7d}" if valid_results is not None: prop_fmt = " %11.2e %11.2e" - for k in valid_results.keys(): - # assert k in train_results.keys() - print_str += prop_fmt % (valid_results[k], train_results[k]) + for key in valid_results.keys(): + print_str += prop_fmt % (valid_results[key], train_results[key]) else: prop_fmt = " %11.2e" - for k in train_results.keys(): - print_str += prop_fmt % (train_results[k]) + for key in train_results.keys(): + print_str += prop_fmt % (train_results[key]) print_str += f" {cur_lr:8.1e}\n" log.info( format_training_message_per_task( @@ -461,22 +1680,88 @@ def print_on_training( fp.write(print_str) fp.flush() + @staticmethod + def print_on_training_multitask( + fp: TextIO, + train_results: dict[str, dict[str, float]], + valid_results: dict[str, dict[str, float]], + cur_batch: int, + cur_lr: float, + ) -> None: + print_str = f"{cur_batch:7d}" + for model_key in train_results.keys(): + if valid_results.get(model_key): + prop_fmt = " %11.2e %11.2e" + for key in sorted(train_results[model_key].keys()): + print_str += prop_fmt % ( + valid_results[model_key][key], + train_results[model_key][key], + ) + else: + prop_fmt = " %11.2e" + for key in sorted(train_results[model_key].keys()): + print_str += prop_fmt % (train_results[model_key][key]) + print_str += f" {cur_lr:8.1e}\n" + for model_key in train_results.keys(): + log.info( + format_training_message_per_task( + batch=cur_batch, + task_name=model_key + "_trn", + rmse=train_results[model_key], + learning_rate=cur_lr, + ) + ) + if valid_results.get(model_key): + log.info( + format_training_message_per_task( + batch=cur_batch, + task_name=model_key + "_val", + rmse=valid_results[model_key], + learning_rate=None, + ) + ) + fp.write(print_str) + fp.flush() + @staticmethod def print_header( fp: TextIO, train_results: dict[str, float], valid_results: dict[str, float] | None, ) -> None: - print_str = "" - print_str += "# {:5s}".format("step") + print_str = "# {:5s}".format("step") if valid_results is not None: prop_fmt = " %11s %11s" - for k in train_results.keys(): - print_str += prop_fmt % (k + "_val", k + "_trn") + for key in train_results.keys(): + print_str += prop_fmt % (key + "_val", key + "_trn") else: prop_fmt = " %11s" - for k in train_results.keys(): - print_str += prop_fmt % (k + "_trn") + for key in train_results.keys(): + print_str += prop_fmt % (key + "_trn") + print_str += " {:8s}\n".format("lr") + print_str += "# If there is no available reference data, rmse_*_{val,trn} will print nan\n" + fp.write(print_str) + fp.flush() + + @staticmethod + def print_header_multitask( + fp: TextIO, + train_results: dict[str, dict[str, float]], + valid_results: dict[str, dict[str, float]], + ) -> None: + print_str = "# {:5s}".format("step") + for model_key in train_results.keys(): + if valid_results.get(model_key): + prop_fmt = " %11s %11s" + for key in sorted(train_results[model_key].keys()): + print_str += prop_fmt % ( + key + f"_val_{model_key}", + key + f"_trn_{model_key}", + ) + else: + prop_fmt = " %11s" + for key in sorted(train_results[model_key].keys()): + print_str += prop_fmt % (key + f"_trn_{model_key}") print_str += " {:8s}\n".format("lr") print_str += "# If there is no available reference data, rmse_*_{val,trn} will print nan\n" fp.write(print_str) @@ -484,21 +1769,21 @@ def print_header( def prepare_input( - *, # enforce keyword-only arguments + *, rcut: float, sel: list[int], coord: np.ndarray, atype: np.ndarray, - box: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + box: np.ndarray | None = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, ) -> tuple[ np.ndarray, np.ndarray, np.ndarray, np.ndarray, - Optional[np.ndarray], - Optional[np.ndarray], + np.ndarray | None, + np.ndarray | None, ]: nframes, nloc = atype.shape[:2] cc, bb, fp, ap = coord, box, fparam, aparam @@ -519,8 +1804,6 @@ def prepare_input( nloc, rcut, sel, - # types will be distinguished in the lower interface, - # so it doesn't need to be distinguished here distinguish_types=False, ) extended_coord = extended_coord.reshape(nframes, -1, 3) @@ -530,22 +1813,8 @@ def prepare_input( def convert_numpy_data_to_jax_data( numpy_data: dict[str, np.ndarray | np.floating], sharding: Any | None = None, + natoms_axis_size: int = 1, ) -> dict[str, jnp.ndarray | bool]: - """Convert NumPy data to JAX data. - - Parameters - ---------- - numpy_data : dict[str, np.ndarray | np.floating] - NumPy data - sharding : Any | None - sharding - - Returns - ------- - jax_data - JAX data - """ - # numpy to jax jax_data = { kk: jnp.asarray(vv) if not kk.startswith("find_") else bool(vv.item()) for kk, vv in numpy_data.items() @@ -555,31 +1824,34 @@ def convert_numpy_data_to_jax_data( kk: jax.make_array_from_process_local_data(sharding, vv) if not kk.startswith("find_") and vv is not None - and kk - not in { - "natoms_vec", - "default_mesh", - } + and kk not in {"natoms_vec", "default_mesh"} else vv for kk, vv in jax_data.items() } + def _label_sharding(key: str, value: jnp.ndarray) -> Any: + if key in {"energy", "box", "numb_copy", "virial", "real_natoms_vec"}: + spec = P("data") + elif natoms_axis_size <= 1 or value.ndim < 2: + spec = P("data") + elif value.shape[1] % natoms_axis_size != 0: + spec = P("data") + else: + spec = P("data", "natoms") + if sharding is not None and hasattr(sharding, "mesh"): + return NamedSharding(sharding.mesh, spec) + if getattr(jax, "set_mesh", None) is None: + return None + return spec + jax_data = { kk: jax.device_put( vv, - P("data", "natoms") - if kk not in {"energy", "box", "numb_copy", "virial", "real_natoms_vec"} - else P( - "data", - ), + _label_sharding(kk, vv), ) if not kk.startswith("find_") and vv is not None - and kk - not in { - "natoms_vec", - "default_mesh", - } + and kk not in {"natoms_vec", "default_mesh"} else vv for kk, vv in jax_data.items() } diff --git a/deepmd/jax/utils/finetune.py b/deepmd/jax/utils/finetune.py new file mode 100644 index 0000000000..61a5a82642 --- /dev/null +++ b/deepmd/jax/utils/finetune.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from copy import ( + deepcopy, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.jax.utils.serialization import ( + serialize_from_file, +) +from deepmd.utils.finetune import ( + FinetuneRuleItem, +) +from deepmd.utils.model_branch_dict import ( + get_model_dict, +) + +log = logging.getLogger(__name__) + + +def _validate_finetune_source(finetune_model: str) -> dict[str, Any]: + if finetune_model.endswith(".savedmodel"): + raise ValueError( + "JAX fine-tuning does not support loading from .savedmodel. " + "Please use a .jax checkpoint or a .hlo frozen model." + ) + if not (finetune_model.endswith(".jax") or finetune_model.endswith(".hlo")): + raise ValueError( + "JAX fine-tuning only supports .jax checkpoints or .hlo frozen models." + ) + return serialize_from_file(finetune_model) + + +def get_finetune_rule_single( + target_model_config: dict[str, Any], + pretrained_model_config: dict[str, Any], + *, + from_multitask: bool = False, + model_branch: str = "Default", + model_branch_from: str = "", + change_model_params: bool = False, +) -> tuple[dict[str, Any], FinetuneRuleItem]: + single_config = deepcopy(target_model_config) + new_fitting = False + model_branch_chosen = "Default" + + if not from_multitask: + single_config_chosen = deepcopy(pretrained_model_config) + if model_branch_from == "RANDOM": + new_fitting = True + else: + model_dict_params = pretrained_model_config["model_dict"] + if model_branch_from in ("", "RANDOM"): + model_branch_chosen = next(iter(model_dict_params.keys())) + new_fitting = True + log.warning( + "The fitting net will be re-initialized instead of using that in the " + "pretrained multitask model because no explicit source branch was " + "selected." + ) + else: + model_branch_chosen = model_branch_from + model_alias_dict, _ = get_model_dict(model_dict_params) + if model_branch_from not in model_alias_dict: + if model_branch_from in ("", "RANDOM"): + model_branch_from = model_branch_chosen + else: + raise ValueError( + f"No model branch or alias named '{model_branch_from}'. " + f"Available ones are {list(model_dict_params.keys())}." + ) + if model_branch_from not in ("", "RANDOM"): + model_branch_chosen = model_branch_from + if model_branch_chosen not in model_alias_dict: + raise ValueError( + f"No model branch or alias named '{model_branch_chosen}'. " + f"Available ones are {list(model_dict_params.keys())}." + ) + model_branch_chosen = model_alias_dict[model_branch_chosen] + single_config_chosen = deepcopy(model_dict_params[model_branch_chosen]) + + finetune_rule = FinetuneRuleItem( + p_type_map=single_config_chosen["type_map"], + type_map=single_config["type_map"], + model_branch=model_branch_chosen, + random_fitting=new_fitting, + ) + if change_model_params: + trainable_param = { + "descriptor": single_config.get("descriptor", {}).get("trainable", True), + "fitting_net": single_config.get("fitting_net", {}).get("trainable", True), + } + single_config["descriptor"] = deepcopy(single_config_chosen["descriptor"]) + if not new_fitting: + single_config["fitting_net"] = deepcopy(single_config_chosen["fitting_net"]) + log.info( + "Change the '%s' model configurations according to the model branch '%s' in the pretrained one...", + model_branch, + model_branch_chosen, + ) + for net_type, trainable in trainable_param.items(): + if net_type in single_config: + single_config[net_type]["trainable"] = trainable + else: + single_config[net_type] = {"trainable": trainable} + return single_config, finetune_rule + + +def get_finetune_rules( + finetune_model: str, + model_config: dict[str, Any], + *, + model_branch: str = "", + change_model_params: bool = True, +) -> tuple[dict[str, Any], dict[str, FinetuneRuleItem], dict[str, Any]]: + finetune_data = _validate_finetune_source(finetune_model) + pretrained_model_config = finetune_data["model_def_script"] + finetune_from_multi_task = "model_dict" in pretrained_model_config + multi_task = "model_dict" in model_config + finetune_links: dict[str, FinetuneRuleItem] = {} + + if not multi_task: + if model_branch == "" and "finetune_head" in model_config: + model_branch = model_config["finetune_head"] + + updated_model_config, finetune_rule = get_finetune_rule_single( + model_config, + pretrained_model_config, + from_multitask=finetune_from_multi_task, + model_branch="Default", + model_branch_from=model_branch, + change_model_params=change_model_params, + ) + finetune_links["Default"] = finetune_rule + return updated_model_config, finetune_links, finetune_data + + if model_branch != "": + raise AssertionError( + "Multi-task fine-tuning does not support command-line branches chosen!" + "Please define the 'finetune_head' in each model params!" + ) + + target_keys = model_config["model_dict"].keys() + pretrained_keys = ( + pretrained_model_config["model_dict"].keys() + if finetune_from_multi_task + else ["Default"] + ) + + for model_key in target_keys: + branch_config = model_config["model_dict"][model_key] + model_branch_from = "RANDOM" + resuming = False + if ( + "finetune_head" in branch_config + and branch_config["finetune_head"] != "RANDOM" + ): + pretrained_key = branch_config["finetune_head"] + if pretrained_key not in pretrained_keys: + raise AssertionError( + f"'{pretrained_key}' head chosen to finetune not exist in the pretrained model!" + f"Available heads are: {list(pretrained_keys)}" + ) + model_branch_from = pretrained_key + elif "finetune_head" not in branch_config and model_key in pretrained_keys: + model_branch_from = model_key + resuming = True + + model_config["model_dict"][model_key], finetune_rule = get_finetune_rule_single( + branch_config, + pretrained_model_config, + from_multitask=finetune_from_multi_task, + model_branch=model_key, + model_branch_from=model_branch_from, + change_model_params=change_model_params, + ) + finetune_rule.resuming = resuming + finetune_links[model_key] = finetune_rule + return model_config, finetune_links, finetune_data + + +def _merge_array_leaves( + target_node: Any, + source_node: Any, + *, + path: tuple[Any, ...] = (), + keep_target_on_shape_mismatch: bool = False, +) -> Any: + if isinstance(target_node, dict): + if not isinstance(source_node, dict): + raise TypeError( + f"Expected dict at {path}, got {type(source_node).__name__}" + ) + merged = deepcopy(target_node) + for key, value in target_node.items(): + if key not in source_node: + raise KeyError(f"Missing key {'.'.join(map(str, (*path, key)))}") + merged[key] = _merge_array_leaves( + value, + source_node[key], + path=(*path, key), + keep_target_on_shape_mismatch=keep_target_on_shape_mismatch, + ) + return merged + if isinstance(target_node, list): + if not isinstance(source_node, list) or len(target_node) != len(source_node): + raise ValueError(f"List mismatch at {path}") + return [ + _merge_array_leaves( + tv, + sv, + path=(*path, idx), + keep_target_on_shape_mismatch=keep_target_on_shape_mismatch, + ) + for idx, (tv, sv) in enumerate(zip(target_node, source_node, strict=True)) + ] + if isinstance(target_node, tuple): + if not isinstance(source_node, tuple) or len(target_node) != len(source_node): + raise ValueError(f"Tuple mismatch at {path}") + return tuple( + _merge_array_leaves( + tv, + sv, + path=(*path, idx), + keep_target_on_shape_mismatch=keep_target_on_shape_mismatch, + ) + for idx, (tv, sv) in enumerate(zip(target_node, source_node, strict=True)) + ) + if isinstance(target_node, np.ndarray): + if not isinstance(source_node, np.ndarray): + raise TypeError( + f"Expected ndarray at {path}, got {type(source_node).__name__}" + ) + if target_node.shape == source_node.shape: + return np.array(source_node, copy=True) + if target_node.size == source_node.size: + return np.array(source_node, copy=True).reshape(target_node.shape) + if keep_target_on_shape_mismatch: + log.info( + "Keeping target-initialized leaf at %s due to shape mismatch: target %s, source %s", + ".".join(map(str, path)), + target_node.shape, + source_node.shape, + ) + return np.array(target_node, copy=True) + raise ValueError( + f"Shape mismatch at {path}: target {target_node.shape}, source {source_node.shape}" + ) + return deepcopy(target_node) + + +def _get_by_path(node: Any, path: tuple[Any, ...]) -> Any: + current = node + for key in path: + current = current[key] + return current + + +def _set_by_path(node: Any, path: tuple[Any, ...], value: Any) -> None: + current = node + for key in path[:-1]: + current = current[key] + current[path[-1]] = value + + +def merge_finetune_model_data( + target_model_data: dict[str, Any], + pretrained_model_data: dict[str, Any], + finetune_rule: FinetuneRuleItem, + *, + source_overrides: dict[tuple[Any, ...], dict[str, Any]] | None = None, +) -> dict[str, Any]: + target_model_data = deepcopy(target_model_data) + target_template = deepcopy(target_model_data) + pretrained_model_data = deepcopy(pretrained_model_data) + source_overrides = source_overrides or {} + if finetune_rule.get_random_fitting(): + if ( + "descriptor" not in target_model_data + or "descriptor" not in pretrained_model_data + ): + raise NotImplementedError( + "JAX random-fitting fine-tuning currently requires standard descriptor/fitting models." + ) + target_model_data["descriptor"] = _merge_array_leaves( + target_model_data["descriptor"], + pretrained_model_data["descriptor"], + path=("descriptor",), + ) + merged_model_data = target_model_data + else: + target_case_embd_dim = target_model_data.get("fitting", {}).get( + "dim_case_embd", 0 + ) + source_case_embd_dim = pretrained_model_data.get("fitting", {}).get( + "dim_case_embd", 0 + ) + if target_case_embd_dim != source_case_embd_dim: + if ( + "descriptor" in target_model_data + and "descriptor" in pretrained_model_data + ): + target_model_data["descriptor"] = _merge_array_leaves( + target_model_data["descriptor"], + pretrained_model_data["descriptor"], + path=("descriptor",), + ) + if ( + "fitting" not in target_model_data + or "fitting" not in pretrained_model_data + ): + raise NotImplementedError( + "JAX case embedding fine-tuning currently requires standard descriptor/fitting models." + ) + target_model_data["fitting"] = _merge_array_leaves( + target_model_data["fitting"], + pretrained_model_data["fitting"], + path=("fitting",), + keep_target_on_shape_mismatch=True, + ) + if ( + "@variables" in target_model_data + and "@variables" in pretrained_model_data + ): + target_model_data["@variables"] = _merge_array_leaves( + target_model_data["@variables"], + pretrained_model_data["@variables"], + path=("@variables",), + ) + merged_model_data = target_model_data + else: + merged_model_data = _merge_array_leaves( + target_model_data, pretrained_model_data + ) + for path, override_source_model_data in source_overrides.items(): + source_subtree = _get_by_path(override_source_model_data, path) + target_subtree = _get_by_path(target_template, path) + _set_by_path( + merged_model_data, + path, + _merge_array_leaves( + target_subtree, + source_subtree, + path=path, + ), + ) + return merged_model_data diff --git a/deepmd/jax/utils/multi_task.py b/deepmd/jax/utils/multi_task.py new file mode 100644 index 0000000000..55681db49f --- /dev/null +++ b/deepmd/jax/utils/multi_task.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from copy import ( + deepcopy, +) +from typing import ( + Any, +) + +from deepmd.jax.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.jax.fitting.base_fitting import ( + BaseFitting, +) + + +def get_class_name(item_key: str, item_params: dict[str, Any]) -> str: + if item_key == "descriptor": + return BaseDescriptor.get_class_by_type( + item_params.get("type", "se_e2_a") + ).__name__ + if item_key == "fitting_net": + return BaseFitting.get_class_by_type(item_params.get("type", "ener")).__name__ + raise RuntimeError(f"Unknown class_name type {item_key}") + + +def preprocess_shared_params( + model_config: dict[str, Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + 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: dict[str, Any] = {} + type_map_keys: list[str] = [] + + def replace_one_item( + params_dict: dict[str, Any], + key_type: str, + key_in_dict: str, + model_key: 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 list(model_params_item.keys()): + 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, + model_key, + ) + 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"], + item_key, + hybrid_item, + model_key, + 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"], + ) + assert len(type_map_keys) == 1, "Multitask model must have only one type_map!" + return model_config, shared_links + + +def get_case_embd_config(model_params: dict[str, Any]) -> tuple[bool, dict[str, int]]: + assert "model_dict" in model_params, ( + "Only support setting case embedding for multi-task model!" + ) + model_keys = list(model_params["model_dict"]) + sorted_model_keys = sorted(model_keys) + numb_case_embd_list = [ + model_params["model_dict"][model_key] + .get("fitting_net", {}) + .get("dim_case_embd", 0) + for model_key in sorted_model_keys + ] + if not all(item == numb_case_embd_list[0] for item in numb_case_embd_list): + raise ValueError( + "All models must have the same dimension of case embedding, " + f"while the settings are: {numb_case_embd_list}" + ) + if numb_case_embd_list[0] == 0: + return False, {} + case_embd_index = { + model_key: idx for idx, model_key in enumerate(sorted_model_keys) + } + return True, case_embd_index diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 86c31e9e78..650eb7b822 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from copy import ( + deepcopy, +) from pathlib import ( Path, ) @@ -16,30 +19,80 @@ jnp, nnx, ) -from deepmd.jax.model.model import ( +from deepmd.jax.model.base_model import ( BaseModel, - get_model, ) +from deepmd.jax.model.model import ( + get_model_for_wrapper, +) +from deepmd.jax.model.multitask import ( + ModelWrapper, +) +from deepmd.jax.utils.multi_task import ( + get_case_embd_config, +) +from deepmd.utils.model_branch_dict import ( + get_model_dict, +) + + +def _is_topology_mismatch_error(exc: Exception) -> bool: + message = str(exc) + return ( + "Topology mismatch detected" in message + or "available devices are different from the devices used to save the checkpoint" + in message + ) + + +def select_model_branch( + data: dict, + model_branch: str | None, +) -> dict: + model_def_script = data["model_def_script"] + if "model_dict" not in model_def_script: + return data + if not model_branch: + raise ValueError( + "Freezing a multitask JAX checkpoint to a single model requires " + "selecting a branch with --head/--model-branch." + ) + model_alias_dict, _ = get_model_dict(model_def_script["model_dict"]) + if model_branch not in model_alias_dict: + raise ValueError( + f"No model branch or alias named '{model_branch}'. " + f"Available branches are: {list(model_def_script['model_dict'].keys())}" + ) + model_branch = model_alias_dict[model_branch] + return { + **data, + "model_def_script": deepcopy(model_def_script["model_dict"][model_branch]), + "model": deepcopy(data["model"]["model_dict"][model_branch]), + } def deserialize_to_file(model_file: str, data: dict, hessian: bool = False) -> None: - """Deserialize the dictionary to a model file. - - Parameters - ---------- - model_file : str - The model file to be saved. - data : dict - The dictionary to be deserialized. - hessian : bool - Add the Hessian to the model output. - """ + """Deserialize the dictionary to a model file.""" if model_file.endswith(".jax"): - model = BaseModel.deserialize(data["model"]) - model_def_script = data["model_def_script"] - if hessian: - model.enable_hessian() - model_def_script["hessian_mode"] = True + model_def_script = data["model_def_script"].copy() + shared_links = model_def_script.get("shared_links") + if "model_dict" in model_def_script: + _, case_embd_index = get_case_embd_config(model_def_script) + model = ModelWrapper.deserialize( + data["model"], + shared_links=shared_links, + case_embd_index=case_embd_index, + ) + if hessian: + raise ValueError( + "Freezing Hessian into a multitask .jax checkpoint is not supported. " + "Please select a single branch first." + ) + else: + model = BaseModel.deserialize(data["model"]) + if hessian: + model.enable_hessian() + model_def_script["hessian_mode"] = True _, state = nnx.split(model) with ocp.Checkpointer( ocp.CompositeCheckpointHandler("state", "model_def_script") @@ -52,6 +105,10 @@ def deserialize_to_file(model_file: str, data: dict, hessian: bool = False) -> N ), ) elif model_file.endswith(".hlo"): + if "model_dict" in data["model_def_script"]: + raise ValueError( + "Freezing a multitask JAX checkpoint to .hlo requires selecting a single branch with --head/--model-branch." + ) model = BaseModel.deserialize(data["model"]) model_def_script = data["model_def_script"] if hessian: @@ -88,18 +145,16 @@ def call_lower_with_fixed_do_atomic_virial( nghost_ = 0 return jax_export.export(jax.jit(call_lower_with_fixed_do_atomic_virial))( - jax.ShapeDtypeStruct( - (nf, nloc + nghost_, 3), jnp.float64 - ), # extended_coord - jax.ShapeDtypeStruct((nf, nloc + nghost_), jnp.int32), # extended_atype - jax.ShapeDtypeStruct((nf, nloc, model.get_nnei()), jnp.int64), # nlist - jax.ShapeDtypeStruct((nf, nloc + nghost_), jnp.int64), # mapping + jax.ShapeDtypeStruct((nf, nloc + nghost_, 3), jnp.float64), + jax.ShapeDtypeStruct((nf, nloc + nghost_), jnp.int32), + jax.ShapeDtypeStruct((nf, nloc, model.get_nnei()), jnp.int64), + jax.ShapeDtypeStruct((nf, nloc + nghost_), jnp.int64), jax.ShapeDtypeStruct((nf, model.get_dim_fparam()), jnp.float64) if model.get_dim_fparam() - else None, # fparam + else None, jax.ShapeDtypeStruct((nf, nloc, model.get_dim_aparam()), jnp.float64) if model.get_dim_aparam() - else None, # aparam + else None, ) exported = exported_whether_do_atomic_virial( @@ -156,32 +211,52 @@ def call_lower_with_fixed_do_atomic_virial( def serialize_from_file(model_file: str) -> dict: - """Serialize the model file to a dictionary. - - Parameters - ---------- - model_file : str - The model file to be serialized. - - Returns - ------- - dict - The serialized model data. - """ + """Serialize the model file to a dictionary.""" if model_file.endswith(".jax"): with ocp.Checkpointer( ocp.CompositeCheckpointHandler("state", "model_def_script") ) as checkpointer: - data = checkpointer.restore( - Path(model_file).absolute(), - ocp.args.Composite( - state=ocp.args.StandardRestore(), - model_def_script=ocp.args.JsonRestore(), - ), - ) + try: + data = checkpointer.restore( + Path(model_file).absolute(), + ocp.args.Composite( + state=ocp.args.StandardRestore(), + model_def_script=ocp.args.JsonRestore(), + ), + ) + except ValueError as exc: + if not _is_topology_mismatch_error(exc): + raise + model_def_script = checkpointer.restore( + Path(model_file).absolute(), + ocp.args.Composite(model_def_script=ocp.args.JsonRestore()), + ).model_def_script + shared_links = model_def_script.get("shared_links") + abstract_model = get_model_for_wrapper( + model_def_script, + shared_links=shared_links, + ) + if "model_dict" in model_def_script: + for model_key in model_def_script["model_dict"]: + if model_def_script["model_dict"][model_key].get( + "hessian_mode", False + ): + abstract_model[model_key].enable_hessian() + elif model_def_script.get("hessian_mode", False): + abstract_model.enable_hessian() + _, abstract_state = nnx.split(abstract_model) + data = checkpointer.restore( + Path(model_file).absolute(), + ocp.args.Composite( + state=ocp.args.StandardRestore( + item=abstract_state.to_pure_dict(), + strict=False, + ), + model_def_script=ocp.args.JsonRestore(), + ), + ) state = data.state - # convert str "1" to int 1 key def convert_str_to_int_key(item: dict) -> None: for key, value in item.copy().items(): if isinstance(value, dict): @@ -193,21 +268,29 @@ def convert_str_to_int_key(item: dict) -> None: model_def_script = data.model_def_script current_step = model_def_script.pop("current_step", 0) - abstract_model = get_model(model_def_script) + shared_links = model_def_script.get("shared_links") + abstract_model = get_model_for_wrapper( + model_def_script, + shared_links=shared_links, + ) + if "model_dict" in model_def_script: + for model_key in model_def_script["model_dict"]: + if model_def_script["model_dict"][model_key].get("hessian_mode", False): + abstract_model[model_key].enable_hessian() + elif model_def_script.get("hessian_mode", False): + abstract_model.enable_hessian() graphdef, abstract_state = nnx.split(abstract_model) abstract_state.replace_by_pure_dict(state) model = nnx.merge(graphdef, abstract_state) - model_dict = model.serialize() - data = { + return { "backend": "JAX", "jax_version": jax.__version__, - "model": model_dict, + "model": model.serialize(), "model_def_script": model_def_script, "@variables": { "current_step": current_step, }, } - return data elif model_file.endswith(".hlo"): data = load_dp_model(model_file) data.pop("constants") diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index dd60d9a7e0..c3ec538539 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -552,12 +552,28 @@ def _merge_batch_data(self, batch_data: list[dict]) -> dict: if not vv["atomic"]: b_data[kk] = np.concatenate([bb[kk] for bb in batch_data], axis=0) else: - b_data[kk] = np.zeros( - (len(batch_data), max_natoms * vv["ndof"] * vv["repeat"]), - dtype=batch_data[0][kk].dtype, - ) - for ii, bb in enumerate(batch_data): - b_data[kk][ii, : bb[kk].shape[1]] = bb[kk][0] + if kk == "hessian": + max_hessian_size = (3 * max_natoms) ** 2 + b_data[kk] = np.zeros( + (len(batch_data), max_hessian_size), + dtype=batch_data[0][kk].dtype, + ) + for ii, bb in enumerate(batch_data): + natoms = bb["natoms_vec"][0] + hessian = bb[kk][0].reshape(3 * natoms, 3 * natoms) + padded = np.zeros( + (3 * max_natoms, 3 * max_natoms), + dtype=bb[kk].dtype, + ) + padded[: 3 * natoms, : 3 * natoms] = hessian + b_data[kk][ii] = padded.reshape(-1) + else: + b_data[kk] = np.zeros( + (len(batch_data), max_natoms * vv["ndof"] * vv["repeat"]), + dtype=batch_data[0][kk].dtype, + ) + for ii, bb in enumerate(batch_data): + b_data[kk][ii, : bb[kk].shape[1]] = bb[kk][0] return b_data # ! altered by Marián Rynik diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py index 8061c7aa9c..87760fa08b 100644 --- a/deepmd/utils/model_stat.py +++ b/deepmd/utils/model_stat.py @@ -9,11 +9,29 @@ import numpy as np +def _get_batch_by_system(data: Any, sys_idx: int) -> dict[str, Any]: + """Get one statistics batch from a concrete underlying system.""" + if not getattr(data, "mixed_systems", False): + return data.get_batch(sys_idx=sys_idx) + + # DeepmdDataSystem.get_batch(sys_idx=...) ignores sys_idx for mixed + # systems and returns a padded mixed batch. Statistics need to stay + # grouped by original system so arrays with different natoms are not + # concatenated before model/stat code sees them. + stat_data = data.data_systems[sys_idx].get_batch(int(data.batch_size[sys_idx])) + stat_data["natoms_vec"] = data.natoms_vec[sys_idx] + stat_data["real_natoms_vec"] = np.tile( + data.natoms_vec[sys_idx], (stat_data["type"].shape[0], 1) + ) + stat_data["default_mesh"] = data.default_mesh[sys_idx] + return stat_data + + def _make_all_stat_ref(data: Any, nbatches: int) -> dict[str, list[Any]]: all_stat = defaultdict(list) for ii in range(data.get_nsystems()): for jj in range(nbatches): - stat_data = data.get_batch(sys_idx=ii) + stat_data = _get_batch_by_system(data, ii) for dd in stat_data: if dd == "natoms_vec": stat_data[dd] = stat_data[dd].astype(np.int32) @@ -48,7 +66,7 @@ def make_stat_input( for ii in range(data.get_nsystems()): sys_stat = defaultdict(list) for jj in range(nbatches): - stat_data = data.get_batch(sys_idx=ii) + stat_data = _get_batch_by_system(data, ii) for dd in stat_data: if dd == "natoms_vec": stat_data[dd] = stat_data[dd].astype(np.int32) diff --git a/pyproject.toml b/pyproject.toml index 2b6810a88c..a7152a4886 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -153,12 +153,18 @@ dev = [ pin_tensorflow_cpu = [ # https://github.com/tensorflow/tensorflow/issues/75279 # macos x86 has been deprecated - "tensorflow-cpu>=2.18,<2.21; platform_machine=='x86_64' and platform_system == 'Linux'", + "tensorflow-cpu==2.20.0; platform_machine=='x86_64' and platform_system == 'Linux'", "tensorflow~=2.18.0; (platform_machine!='x86_64' or platform_system != 'Linux') and (platform_machine!='x86_64' or platform_system != 'Darwin')", "tensorflow; platform_machine=='x86_64' and platform_system == 'Darwin'", + # TODO: unpin protobuf after TF is upgraded to 2.21 + # See: https://github.com/tensorflow/tensorflow/pull/103382 + "protobuf<7.34.0", ] pin_tensorflow_gpu = [ - "tensorflow~=2.18.0", + "tensorflow==2.18.0", + # TODO: unpin protobuf after TF is upgraded to 2.21 + # See: https://github.com/tensorflow/tensorflow/pull/103382 + "protobuf<7.34.0", ] pin_pytorch_cpu = [ # https://github.com/pytorch/pytorch/issues/114602 diff --git a/source/tests/common/test_model_stat.py b/source/tests/common/test_model_stat.py new file mode 100644 index 0000000000..2e685d4bbc --- /dev/null +++ b/source/tests/common/test_model_stat.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + + +class TestModelStatMixedSystems(unittest.TestCase): + class _FakeSystem: + def __init__(self, natoms: int, atype: list[int]) -> None: + self.natoms = natoms + self.atype = np.asarray(atype, dtype=np.int64) + self.calls = 0 + + def get_batch(self, batch_size: int) -> dict: + self.calls += 1 + coord = np.full( + (batch_size, self.natoms * 3), + float(self.natoms), + dtype=np.float64, + ) + return { + "type": np.tile(self.atype.reshape(1, -1), (batch_size, 1)), + "coord": coord, + "energy": np.full((batch_size, 1), float(self.natoms)), + "force": -coord, + "find_energy": np.float32(1.0), + "find_force": np.float32(1.0), + } + + class _FakeMixedData: + mixed_systems = True + + def __init__(self) -> None: + self.data_systems = [ + TestModelStatMixedSystems._FakeSystem(4, [0, 1, 1, 1]), + TestModelStatMixedSystems._FakeSystem(7, [0, 0, 1, 1, 1, 1, 1]), + ] + self.batch_size = np.asarray([2, 2], dtype=np.int64) + self.natoms_vec = [ + np.asarray([4, 4, 1, 3], dtype=np.int32), + np.asarray([7, 7, 2, 5], dtype=np.int32), + ] + self.default_mesh = [ + np.zeros(6, dtype=np.int32), + np.zeros(6, dtype=np.int32), + ] + self.fallback_get_batch_calls = 0 + + def get_nsystems(self) -> int: + return len(self.data_systems) + + def get_batch(self, sys_idx=None) -> dict: + self.fallback_get_batch_calls += 1 + raise AssertionError("mixed-system stat collection must not use get_batch") + + def test_make_stat_input_keeps_mixed_systems_separate(self) -> None: + from deepmd.utils.model_stat import ( + make_stat_input, + ) + + data = self._FakeMixedData() + all_stat = make_stat_input(data, nbatches=2, merge_sys=False) + + self.assertEqual(data.fallback_get_batch_calls, 0) + self.assertEqual([sys.calls for sys in data.data_systems], [2, 2]) + self.assertEqual(len(all_stat["coord"]), 2) + self.assertEqual([batch.shape for batch in all_stat["coord"][0]], [(2, 12)] * 2) + self.assertEqual([batch.shape for batch in all_stat["coord"][1]], [(2, 21)] * 2) + self.assertEqual([batch.shape for batch in all_stat["force"][0]], [(2, 12)] * 2) + self.assertEqual([batch.shape for batch in all_stat["force"][1]], [(2, 21)] * 2) + self.assertEqual( + [batch.shape for batch in all_stat["real_natoms_vec"][0]], + [(2, 4)] * 2, + ) + self.assertEqual( + [batch.shape for batch in all_stat["real_natoms_vec"][1]], + [(2, 4)] * 2, + ) + np.testing.assert_array_equal( + all_stat["real_natoms_vec"][0][0], + np.tile(np.asarray([4, 4, 1, 3], dtype=np.int32), (2, 1)), + ) + np.testing.assert_array_equal( + all_stat["real_natoms_vec"][1][0], + np.tile(np.asarray([7, 7, 2, 5], dtype=np.int32), (2, 1)), + ) + self.assertFalse(np.any(all_stat["type"][0][0] < 0)) + self.assertFalse(np.any(all_stat["type"][1][0] < 0)) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/jax/test_finetune.py b/source/tests/jax/test_finetune.py new file mode 100644 index 0000000000..30ad3329c2 --- /dev/null +++ b/source/tests/jax/test_finetune.py @@ -0,0 +1,1617 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import sys +import unittest +from copy import ( + deepcopy, +) +from unittest.mock import ( + Mock, + patch, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.jax.common import ( + to_jax_array, +) +from deepmd.jax.descriptor.dpa3 import ( + DescrptDPA3, +) +from deepmd.jax.entrypoints import train as jax_train_entrypoint +from deepmd.jax.fitting.fitting import ( + EnergyFittingNet, +) +from deepmd.jax.model import ( + EnergyModel, +) +from deepmd.jax.train.trainer import ( + DPTrainer, + _merge_init_frz_model_data, +) +from deepmd.jax.utils.finetune import ( + get_finetune_rule_single, + get_finetune_rules, + merge_finetune_model_data, +) +from deepmd.utils.finetune import ( + FinetuneRuleItem, +) + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestJAXFinetuneRules(unittest.TestCase): + def setUp(self) -> None: + self.target_model_config = { + "type": "standard", + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa3", + "trainable": False, + "trainable_ln": False, + "rcut": 6.0, + }, + "fitting_net": { + "type": "ener", + "trainable": [False, True], + "neuron": [16], + }, + } + self.pretrained_model_config = { + "type": "standard", + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa3", + "trainable": True, + "trainable_ln": True, + "rcut": 5.0, + }, + "fitting_net": { + "type": "ener", + "trainable": [True, True], + "neuron": [32], + }, + } + + def test_single_task_preserves_trainable_flags(self) -> None: + updated_config, finetune_rule = get_finetune_rule_single( + self.target_model_config, + self.pretrained_model_config, + change_model_params=True, + ) + self.assertFalse(updated_config["descriptor"]["trainable"]) + self.assertTrue(updated_config["descriptor"]["trainable_ln"]) + self.assertEqual(updated_config["descriptor"]["rcut"], 5.0) + self.assertEqual(updated_config["fitting_net"]["trainable"], [False, True]) + self.assertEqual(updated_config["fitting_net"]["neuron"], [32]) + self.assertFalse(finetune_rule.get_random_fitting()) + + def test_single_task_random_fitting_keeps_target_fitting(self) -> None: + updated_config, finetune_rule = get_finetune_rule_single( + self.target_model_config, + self.pretrained_model_config, + model_branch_from="RANDOM", + change_model_params=True, + ) + self.assertTrue(finetune_rule.get_random_fitting()) + self.assertEqual(updated_config["descriptor"]["rcut"], 5.0) + self.assertEqual(updated_config["fitting_net"]["neuron"], [16]) + + def test_get_finetune_rules_accepts_jax(self) -> None: + finetune_data = { + "model_def_script": self.pretrained_model_config, + "model": {"unused": True}, + } + with patch( + "deepmd.jax.utils.finetune.serialize_from_file", + return_value=finetune_data, + ): + updated_config, finetune_links, returned_data = get_finetune_rules( + "pretrained.jax", + deepcopy(self.target_model_config), + change_model_params=True, + ) + self.assertEqual(updated_config["descriptor"]["rcut"], 5.0) + self.assertEqual(finetune_links["Default"].get_model_branch(), "Default") + self.assertIs(returned_data, finetune_data) + + def test_get_finetune_rules_accepts_hlo(self) -> None: + finetune_data = { + "model_def_script": self.pretrained_model_config, + "model": {"unused": True}, + } + with patch( + "deepmd.jax.utils.finetune.serialize_from_file", + return_value=finetune_data, + ): + _, finetune_links, _ = get_finetune_rules( + "pretrained.hlo", + deepcopy(self.target_model_config), + change_model_params=False, + ) + self.assertEqual(finetune_links["Default"].get_model_branch(), "Default") + + def test_get_finetune_rules_rejects_savedmodel(self) -> None: + with self.assertRaisesRegex(ValueError, ".savedmodel"): + get_finetune_rules( + "pretrained.savedmodel", + deepcopy(self.target_model_config), + ) + + def test_multitask_source_random_fitting_uses_first_branch(self) -> None: + updated_config, finetune_rule = get_finetune_rule_single( + deepcopy(self.target_model_config), + {"model_dict": {"branch_a": deepcopy(self.pretrained_model_config)}}, + from_multitask=True, + change_model_params=True, + ) + self.assertTrue(finetune_rule.get_random_fitting()) + self.assertEqual(finetune_rule.get_model_branch(), "branch_a") + self.assertEqual(updated_config["descriptor"]["rcut"], 5.0) + self.assertEqual(updated_config["fitting_net"]["neuron"], [16]) + + def test_multitask_source_explicit_random_keyword_uses_first_branch(self) -> None: + _, finetune_rule = get_finetune_rule_single( + deepcopy(self.target_model_config), + {"model_dict": {"branch_a": deepcopy(self.pretrained_model_config)}}, + from_multitask=True, + model_branch_from="RANDOM", + change_model_params=True, + ) + self.assertTrue(finetune_rule.get_random_fitting()) + self.assertEqual(finetune_rule.get_model_branch(), "branch_a") + + def test_get_finetune_rules_from_multitask_source_branch(self) -> None: + finetune_data = { + "model_def_script": { + "model_dict": { + "branch_a": deepcopy(self.pretrained_model_config), + "branch_b": { + **deepcopy(self.pretrained_model_config), + "descriptor": { + **deepcopy(self.pretrained_model_config["descriptor"]), + "rcut": 7.0, + }, + }, + } + }, + "model": {"unused": True}, + } + with patch( + "deepmd.jax.utils.finetune._validate_finetune_source", + return_value=finetune_data, + ): + updated_config, finetune_links, returned_data = get_finetune_rules( + "pretrained.hlo", + { + **deepcopy(self.target_model_config), + "finetune_head": "branch_b", + }, + change_model_params=True, + ) + self.assertEqual(updated_config["descriptor"]["rcut"], 7.0) + self.assertEqual(finetune_links["Default"].get_model_branch(), "branch_b") + self.assertIs(returned_data, finetune_data) + + def test_multitask_target_from_single_source(self) -> None: + finetune_data = { + "model_def_script": deepcopy(self.pretrained_model_config), + "model": {"unused": True}, + } + target = { + "model_dict": { + "Default": deepcopy(self.target_model_config), + "new_head": { + **deepcopy(self.target_model_config), + "finetune_head": "RANDOM", + }, + } + } + with patch( + "deepmd.jax.utils.finetune._validate_finetune_source", + return_value=finetune_data, + ): + updated_config, finetune_links, _ = get_finetune_rules( + "pretrained.jax", + deepcopy(target), + change_model_params=True, + ) + self.assertTrue(finetune_links["Default"].get_resuming()) + self.assertEqual(finetune_links["Default"].get_model_branch(), "Default") + self.assertFalse(finetune_links["Default"].get_random_fitting()) + self.assertTrue(finetune_links["new_head"].get_random_fitting()) + self.assertFalse(finetune_links["new_head"].get_resuming()) + self.assertEqual( + updated_config["model_dict"]["Default"]["descriptor"]["rcut"], 5.0 + ) + self.assertEqual( + updated_config["model_dict"]["new_head"]["fitting_net"]["neuron"], [16] + ) + + def test_multitask_target_from_multitask_source(self) -> None: + finetune_data = { + "model_def_script": { + "model_dict": { + "head_a": deepcopy(self.pretrained_model_config), + "head_b": { + **deepcopy(self.pretrained_model_config), + "descriptor": { + **deepcopy(self.pretrained_model_config["descriptor"]), + "rcut": 7.0, + }, + }, + } + }, + "model": {"unused": True}, + } + target = { + "model_dict": { + "head_a": deepcopy(self.target_model_config), + "new_head": { + **deepcopy(self.target_model_config), + "finetune_head": "head_b", + }, + "fresh_head": deepcopy(self.target_model_config), + } + } + with patch( + "deepmd.jax.utils.finetune._validate_finetune_source", + return_value=finetune_data, + ): + updated_config, finetune_links, _ = get_finetune_rules( + "pretrained.jax", + deepcopy(target), + change_model_params=True, + ) + self.assertTrue(finetune_links["head_a"].get_resuming()) + self.assertEqual(finetune_links["head_a"].get_model_branch(), "head_a") + self.assertFalse(finetune_links["new_head"].get_resuming()) + self.assertEqual(finetune_links["new_head"].get_model_branch(), "head_b") + self.assertTrue(finetune_links["fresh_head"].get_random_fitting()) + self.assertEqual( + updated_config["model_dict"]["new_head"]["descriptor"]["rcut"], 7.0 + ) + + def test_multitask_target_rejects_command_line_model_branch(self) -> None: + finetune_data = { + "model_def_script": deepcopy(self.pretrained_model_config), + "model": {"unused": True}, + } + with ( + patch( + "deepmd.jax.utils.finetune._validate_finetune_source", + return_value=finetune_data, + ), + self.assertRaisesRegex( + AssertionError, "Multi-task fine-tuning does not support" + ), + ): + get_finetune_rules( + "pretrained.jax", + {"model_dict": {"Default": deepcopy(self.target_model_config)}}, + model_branch="Default", + ) + + def test_multitask_target_invalid_branch_raises(self) -> None: + finetune_data = { + "model_def_script": { + "model_dict": {"head_a": deepcopy(self.pretrained_model_config)} + }, + "model": {"unused": True}, + } + with ( + patch( + "deepmd.jax.utils.finetune._validate_finetune_source", + return_value=finetune_data, + ), + self.assertRaisesRegex(AssertionError, "chosen to finetune not exist"), + ): + get_finetune_rules( + "pretrained.jax", + { + "model_dict": { + "head_b": { + **deepcopy(self.target_model_config), + "finetune_head": "missing", + } + } + }, + ) + + def test_change_model_params_uses_source_case_embd_dimension(self) -> None: + target = deepcopy(self.target_model_config) + target["fitting_net"]["dim_case_embd"] = 3 + pretrained = deepcopy(self.pretrained_model_config) + pretrained["fitting_net"]["dim_case_embd"] = 2 + updated_config, _ = get_finetune_rule_single( + target, + pretrained, + change_model_params=True, + ) + self.assertEqual(updated_config["fitting_net"]["dim_case_embd"], 2) + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestJAXFinetuneMerge(unittest.TestCase): + def test_merge_replaces_array_leaves_only(self) -> None: + target = { + "descriptor": { + "weights": np.zeros((2, 2)), + "trainable": False, + }, + "fitting": { + "weights": np.ones((2, 2)), + "trainable": True, + }, + } + source = { + "descriptor": { + "weights": np.full((2, 2), 2.0), + "trainable": True, + }, + "fitting": { + "weights": np.full((2, 2), 3.0), + "trainable": False, + }, + } + merged = merge_finetune_model_data( + target, + source, + FinetuneRuleItem(["O", "H"], ["O", "H"]), + ) + np.testing.assert_array_equal( + merged["descriptor"]["weights"], source["descriptor"]["weights"] + ) + np.testing.assert_array_equal( + merged["fitting"]["weights"], source["fitting"]["weights"] + ) + self.assertFalse(merged["descriptor"]["trainable"]) + self.assertTrue(merged["fitting"]["trainable"]) + + def test_merge_reshapes_same_size_leaves(self) -> None: + target = {"descriptor": {"weights": np.zeros((2, 3))}} + source = {"descriptor": {"weights": np.arange(6).reshape(3, 2)}} + merged = merge_finetune_model_data( + target, + source, + FinetuneRuleItem(["O", "H"], ["O", "H"]), + ) + np.testing.assert_array_equal( + merged["descriptor"]["weights"], + np.arange(6).reshape(2, 3), + ) + + def test_random_fitting_only_inherits_descriptor(self) -> None: + target = { + "descriptor": {"weights": np.zeros((2, 2))}, + "fitting": {"weights": np.ones((2, 2))}, + } + source = { + "descriptor": {"weights": np.full((2, 2), 2.0)}, + "fitting": {"weights": np.full((2, 2), 3.0)}, + } + merged = merge_finetune_model_data( + target, + source, + FinetuneRuleItem(["O", "H"], ["O", "H"], random_fitting=True), + ) + np.testing.assert_array_equal( + merged["descriptor"]["weights"], source["descriptor"]["weights"] + ) + np.testing.assert_array_equal( + merged["fitting"]["weights"], target["fitting"]["weights"] + ) + + def test_case_embedding_migration_keeps_target_shape_mismatch_leaves(self) -> None: + target = { + "descriptor": {"weights": np.zeros((2, 2))}, + "fitting": { + "dim_case_embd": 3, + "weights": np.ones((2, 5)), + }, + "@variables": {"bias": np.zeros((2,))}, + } + source = { + "descriptor": {"weights": np.full((2, 2), 2.0)}, + "fitting": { + "dim_case_embd": 2, + "weights": np.full((2, 4), 3.0), + }, + "@variables": {"bias": np.full((2,), 4.0)}, + } + merged = merge_finetune_model_data( + target, + source, + FinetuneRuleItem(["O", "H"], ["O", "H"]), + ) + np.testing.assert_array_equal( + merged["descriptor"]["weights"], source["descriptor"]["weights"] + ) + np.testing.assert_array_equal( + merged["fitting"]["weights"], target["fitting"]["weights"] + ) + np.testing.assert_array_equal( + merged["@variables"]["bias"], source["@variables"]["bias"] + ) + + def test_source_override_replaces_only_requested_subtree(self) -> None: + target = { + "descriptor": {"weights": np.zeros((2, 2))}, + "fitting": {"nets": np.ones((2, 2)), "keep": np.ones((1,))}, + } + source = { + "descriptor": {"weights": np.full((2, 2), 2.0)}, + "fitting": {"nets": np.full((2, 2), 3.0), "keep": np.full((1,), 4.0)}, + } + override_source = { + "fitting": {"nets": np.full((2, 2), 5.0), "keep": np.full((1,), 6.0)}, + } + merged = merge_finetune_model_data( + target, + source, + FinetuneRuleItem(["O", "H"], ["O", "H"]), + source_overrides={("fitting", "nets"): override_source}, + ) + np.testing.assert_array_equal( + merged["descriptor"]["weights"], source["descriptor"]["weights"] + ) + np.testing.assert_array_equal( + merged["fitting"]["nets"], override_source["fitting"]["nets"] + ) + np.testing.assert_array_equal( + merged["fitting"]["keep"], source["fitting"]["keep"] + ) + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestJAXFinetuneSharedOverrides(unittest.TestCase): + def test_shared_type_to_serialized_paths_descriptor_partial(self) -> None: + paths = DPTrainer._shared_type_to_serialized_paths( + "descriptor", + 1, + {"descriptor": {"type": "dpa3", "type_embedding": {"w": np.array([1.0])}}}, + ) + self.assertEqual(paths, [("descriptor", "type_embedding")]) + + def test_shared_type_to_serialized_paths_hybrid_partial(self) -> None: + paths = DPTrainer._shared_type_to_serialized_paths( + "descriptor_hybrid_0", + 1, + { + "descriptor": { + "list": [ + {"type": "dpa3", "type_embedding": {"w": np.array([1.0])}}, + {"type": "se_e2_a"}, + ] + } + }, + ) + self.assertEqual(paths, [("descriptor", "list", 0, "type_embedding")]) + + def test_collect_shared_source_overrides_partial_descriptor_only(self) -> None: + trainer = DPTrainer.__new__(DPTrainer) + trainer.model_keys = ["head_a", "head_b"] + trainer.shared_links = { + "shared_desc": { + "links": [ + { + "model_key": "head_a", + "shared_type": "descriptor", + "shared_level": 0, + }, + { + "model_key": "head_b", + "shared_type": "descriptor", + "shared_level": 1, + }, + ] + } + } + trainer.finetune_links = { + "head_a": FinetuneRuleItem(["H"], ["H"], model_branch="src_a"), + "head_b": FinetuneRuleItem(["H"], ["H"], model_branch="src_b"), + } + trainer.finetune_model_data = { + "model": { + "model_dict": { + "src_a": { + "descriptor": { + "type": "dpa3", + "type_embedding": {"w": np.array([1.0])}, + "repflows": {"w": np.array([2.0])}, + } + }, + "src_b": { + "descriptor": { + "type": "dpa3", + "type_embedding": {"w": np.array([3.0])}, + "repflows": {"w": np.array([4.0])}, + } + }, + } + } + } + overrides = trainer._collect_shared_source_overrides(source_multi=True) + self.assertEqual( + list(overrides["head_b"].keys()), + [("descriptor", "type_embedding")], + ) + self.assertIs( + overrides["head_b"][("descriptor", "type_embedding")], + trainer.finetune_model_data["model"]["model_dict"]["src_a"], + ) + + def test_collect_shared_source_overrides_hybrid_partial_descriptor_only( + self, + ) -> None: + trainer = DPTrainer.__new__(DPTrainer) + trainer.model_keys = ["head_a", "head_b"] + trainer.shared_links = { + "shared_desc": { + "links": [ + { + "model_key": "head_a", + "shared_type": "descriptor_hybrid_0", + "shared_level": 0, + }, + { + "model_key": "head_b", + "shared_type": "descriptor_hybrid_0", + "shared_level": 1, + }, + ] + } + } + trainer.finetune_links = { + "head_a": FinetuneRuleItem(["H"], ["H"], model_branch="src_a"), + "head_b": FinetuneRuleItem(["H"], ["H"], model_branch="src_b"), + } + trainer.finetune_model_data = { + "model": { + "model_dict": { + "src_a": { + "descriptor": { + "list": [ + { + "type": "dpa3", + "type_embedding": {"w": np.array([1.0])}, + "repflows": {"w": np.array([2.0])}, + }, + {"type": "se_e2_a"}, + ] + } + }, + "src_b": { + "descriptor": { + "list": [ + { + "type": "dpa3", + "type_embedding": {"w": np.array([3.0])}, + "repflows": {"w": np.array([4.0])}, + }, + {"type": "se_e2_a"}, + ] + } + }, + } + } + } + overrides = trainer._collect_shared_source_overrides(source_multi=True) + self.assertEqual( + list(overrides["head_b"].keys()), + [("descriptor", "list", 0, "type_embedding")], + ) + self.assertIs( + overrides["head_b"][("descriptor", "list", 0, "type_embedding")], + trainer.finetune_model_data["model"]["model_dict"]["src_a"], + ) + + def test_validate_shared_finetune_rules_rejects_partial_fitting_share(self) -> None: + trainer = DPTrainer.__new__(DPTrainer) + trainer.multi_task = True + trainer.shared_links = { + "shared_fit": { + "links": [ + { + "model_key": "head_a", + "shared_type": "fitting_net", + "shared_level": 1, + } + ] + } + } + trainer.finetune_links = {"head_a": FinetuneRuleItem(["H"], ["H"])} + with self.assertRaisesRegex( + NotImplementedError, "fitting_net sharing only supports shared_level=0" + ): + trainer._validate_shared_finetune_rules() + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestJAXFinetuneTypeMapConsistency(unittest.TestCase): + def test_change_type_map_consistency(self) -> None: + descriptor_kwargs = { + "repflow": { + "n_dim": 8, + "e_dim": 4, + "a_dim": 4, + "nlayers": 2, + "e_rcut": 4.0, + "e_rcut_smth": 2.0, + "e_sel": 8, + "a_rcut": 3.0, + "a_rcut_smth": 1.5, + "a_sel": 4, + "axis_neuron": 2, + "a_compress_rate": 1, + "a_compress_e_rate": 1, + "a_compress_use_split": True, + "update_angle": True, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const", + "smooth_edge_update": True, + }, + "activation_function": "tanh", + "use_tebd_bias": False, + "precision": "float64", + "concat_output_tebd": False, + } + fitting_kwargs = { + "neuron": [8, 8], + "resnet_dt": True, + "precision": "float64", + "activation_function": "tanh", + "seed": 1, + } + pretrained_ds = DescrptDPA3( + ntypes=3, + type_map=["H", "O", "B"], + **deepcopy(descriptor_kwargs), + ) + pretrained_ft = EnergyFittingNet( + 3, + pretrained_ds.get_dim_out(), + mixed_types=pretrained_ds.mixed_types(), + type_map=["H", "O", "B"], + **deepcopy(fitting_kwargs), + ) + pretrained_model = EnergyModel( + pretrained_ds, + pretrained_ft, + type_map=["H", "O", "B"], + ) + target_ds = DescrptDPA3( + ntypes=3, + type_map=["O", "H", "B"], + **deepcopy(descriptor_kwargs), + ) + target_ft = EnergyFittingNet( + 3, + target_ds.get_dim_out(), + mixed_types=target_ds.mixed_types(), + type_map=["O", "H", "B"], + **deepcopy(fitting_kwargs), + ) + target_model = EnergyModel( + target_ds, + target_ft, + type_map=["O", "H", "B"], + ) + finetune_rule = FinetuneRuleItem(["H", "O", "B"], ["O", "H", "B"]) + + pretrained_model.change_type_map( + target_model.get_type_map(), + model_with_new_type_stat=target_model.atomic_model, + ) + changed_pretrained_model = EnergyModel.deserialize(pretrained_model.serialize()) + merged = merge_finetune_model_data( + target_model.serialize(), + pretrained_model.serialize(), + finetune_rule, + ) + finetuned_model = EnergyModel.deserialize(merged) + + coord = np.array( + [[[0.0, 0.0, 0.0], [0.0, 1.1, 0.0], [0.9, 0.0, 0.0]]], + dtype=np.float64, + ).reshape(1, 9) + box = (5.0 * np.eye(3)).reshape(1, 9) + atype_new = np.array([[1, 0, 1]], dtype=np.int64) + + old_ret = changed_pretrained_model.call( + to_jax_array(coord), + to_jax_array(atype_new), + box=to_jax_array(box), + ) + new_ret = finetuned_model.call( + to_jax_array(coord), + to_jax_array(atype_new), + box=to_jax_array(box), + ) + np.testing.assert_allclose( + to_numpy_array(old_ret["energy"]), + to_numpy_array(new_ret["energy"]), + atol=1e-10, + ) + np.testing.assert_allclose( + to_numpy_array(old_ret["force"]), + to_numpy_array(new_ret["force"]), + atol=1e-10, + ) + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestJAXFinetuneWiring(unittest.TestCase): + def _trainer_jdata(self) -> dict: + return { + "model": {"type": "standard"}, + "learning_rate": {"start_lr": 1e-3, "decay_steps": 1, "stop_lr": 1e-8}, + "training": {"numb_steps": 1}, + } + + @staticmethod + def _dummy_wrapper_cls(): + class DummyWrapper(dict): + @classmethod + def deserialize(cls, *args, **kwargs): + return Mock() + + return DummyWrapper + + def test_trainer_finetune_does_not_restore_step(self) -> None: + dummy_model = Mock() + dummy_model.get_dim_fparam.return_value = 0 + with ( + patch( + "deepmd.jax.train.trainer.get_model_for_wrapper", + return_value=dummy_model, + ), + patch( + "deepmd.jax.train.trainer.EnergyLoss.get_loss", + return_value=Mock(label_requirement=[]), + ), + patch("deepmd.jax.train.trainer.serialize_from_file") as mock_serialize, + ): + trainer = DPTrainer( + self._trainer_jdata(), + finetune_model="pretrained.jax", + finetune_links={"Default": FinetuneRuleItem(["H"], ["H"])}, + finetune_model_data={"model": {"dummy": True}}, + ) + self.assertEqual(trainer.start_step, 0) + mock_serialize.assert_not_called() + + def test_trainer_finetune_selects_multitask_source_branch(self) -> None: + dummy_model = Mock() + dummy_model.get_dim_fparam.return_value = 0 + with ( + patch( + "deepmd.jax.train.trainer.get_model_for_wrapper", + return_value=dummy_model, + ), + patch( + "deepmd.jax.train.trainer.EnergyLoss.get_loss", + return_value=Mock(label_requirement=[]), + ), + patch.object( + DPTrainer, + "_apply_single_finetune", + return_value=dummy_model, + ) as mock_apply, + patch( + "deepmd.jax.train.trainer.select_model_branch", + return_value={ + "model_def_script": {"type": "standard"}, + "model": {"branch_value": "selected"}, + }, + ) as mock_select, + patch( + "deepmd.jax.train.trainer._pack_data_for_bias_adjust", + return_value={"coord": np.zeros((1, 1, 3))}, + ), + patch( + "deepmd.jax.train.trainer.model_change_out_bias", + return_value=dummy_model, + ), + ): + trainer = DPTrainer( + self._trainer_jdata(), + finetune_model="pretrained.jax", + finetune_links={ + "Default": FinetuneRuleItem(["H"], ["H"], model_branch="branch_b") + }, + finetune_model_data={ + "model_def_script": { + "model_dict": { + "branch_a": {"type": "standard"}, + "branch_b": {"type": "standard"}, + } + }, + "model": { + "model_dict": { + "branch_a": {"branch_value": "a"}, + "branch_b": {"branch_value": "b"}, + } + }, + }, + ) + train_data = Mock() + train_data.get_nsystems.return_value = 1 + trainer._finetune_single(train_data) + mock_select.assert_called_once() + self.assertEqual(mock_select.call_args.args[1], "branch_b") + self.assertEqual(mock_apply.call_args.args[1], {"branch_value": "selected"}) + + def test_trainer_finetune_with_new_type_computes_stats(self) -> None: + dummy_model = Mock() + dummy_model.get_dim_fparam.return_value = 0 + dummy_model.atomic_model = Mock() + dummy_model.serialize.return_value = {"descriptor": np.zeros((1,))} + loss = Mock(label_requirement=[]) + + with ( + patch( + "deepmd.jax.train.trainer.get_model_for_wrapper", + return_value=dummy_model, + ), + patch( + "deepmd.jax.train.trainer.EnergyLoss.get_loss", + return_value=loss, + ), + patch( + "deepmd.jax.train.trainer.make_stat_input", + return_value={ + "type": [[np.array([[0]], dtype=np.int32)]], + "coord": [[np.zeros((1, 3))]], + }, + ), + patch( + "deepmd.jax.train.trainer.jnp.asarray", + side_effect=lambda x: x, + ), + patch.object( + DPTrainer, + "_apply_single_finetune", + return_value=dummy_model, + ), + patch( + "deepmd.jax.train.trainer._pack_data_for_bias_adjust", + return_value={"coord": np.zeros((1, 1, 3))}, + ), + patch( + "deepmd.jax.train.trainer.model_change_out_bias", + return_value=dummy_model, + ), + patch( + "deepmd.jax.train.trainer.jax.make_mesh", + side_effect=RuntimeError("stop_after_stats"), + ), + ): + trainer = DPTrainer( + self._trainer_jdata(), + finetune_model="pretrained.jax", + finetune_links={"Default": FinetuneRuleItem(["H"], ["H", "He"])}, + finetune_model_data={"model": {"dummy": True}}, + ) + train_data = Mock() + train_data.get_nsystems.return_value = 1 + train_data.mixed_type = False + train_data.data_systems = [Mock(pbc=False)] + with self.assertRaisesRegex(RuntimeError, "stop_after_stats"): + trainer._train_single(train_data) + dummy_model.atomic_model.descriptor.compute_input_stats.assert_called_once() + dummy_model.atomic_model.fitting.compute_output_stats.assert_called_once() + + def test_train_entrypoint_wires_finetune_rule(self) -> None: + fake_jdata = { + "model": {"type": "standard", "type_map": ["H"]}, + "learning_rate": {"start_lr": 1e-3, "decay_steps": 1, "stop_lr": 1e-8}, + "training": { + "numb_steps": 1, + "training_data": {"systems": []}, + }, + } + train_data = Mock() + train_data.type_map = ["H"] + trainer_instance = Mock() + trainer_instance.model.get_rcut.return_value = 6.0 + trainer_instance.model.get_type_map.return_value = ["H"] + trainer_instance.data_requirements = [] + + with ( + patch.object( + jax_train_entrypoint, "j_loader", return_value=deepcopy(fake_jdata) + ), + patch.object( + jax_train_entrypoint, + "get_finetune_rules", + return_value=( + deepcopy(fake_jdata["model"]), + {"Default": FinetuneRuleItem(["H"], ["H"])}, + { + "model": {"dummy": True}, + "model_def_script": deepcopy(fake_jdata["model"]), + }, + ), + ) as mock_rules, + patch.object( + jax_train_entrypoint, + "update_deepmd_input", + side_effect=lambda x, **kwargs: x, + ), + patch.object( + jax_train_entrypoint, "normalize", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, "update_sel", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, + "SummaryPrinter", + return_value=Mock(__call__=Mock()), + ), + patch.object( + jax_train_entrypoint, "DPTrainer", return_value=trainer_instance + ) as mock_trainer, + patch.object(jax_train_entrypoint, "get_data", return_value=train_data), + patch.object(jax_train_entrypoint.dp_random, "seed"), + patch( + "builtins.open", + unittest.mock.mock_open(), + ), + patch.object(jax_train_entrypoint.json, "dump"), + ): + jax_train_entrypoint.train( + INPUT="input.json", + init_model=None, + restart=None, + output="out.json", + init_frz_model="", + mpi_log="master", + log_level=2, + log_path=None, + finetune="pretrained.jax", + use_pretrain_script=False, + ) + mock_rules.assert_called_once() + _, rule_kwargs = mock_rules.call_args + self.assertFalse(rule_kwargs["change_model_params"]) + _, kwargs = mock_trainer.call_args + self.assertEqual(kwargs["finetune_model"], "pretrained.jax") + self.assertEqual( + kwargs["finetune_links"]["Default"].get_model_branch(), "Default" + ) + self.assertIn("model", kwargs["finetune_model_data"]) + + def test_train_entrypoint_init_model_without_pretrain_script_keeps_target_model( + self, + ) -> None: + fake_jdata = { + "model": {"type": "standard", "type_map": ["H"]}, + "learning_rate": {"start_lr": 1e-3, "decay_steps": 1, "stop_lr": 1e-8}, + "training": { + "numb_steps": 1, + "training_data": {"systems": []}, + }, + } + train_data = Mock() + train_data.type_map = ["H"] + trainer_instance = Mock() + trainer_instance.model.get_rcut.return_value = 6.0 + trainer_instance.model.get_type_map.return_value = ["H"] + trainer_instance.data_requirements = [] + + with ( + patch.object( + jax_train_entrypoint, "j_loader", return_value=deepcopy(fake_jdata) + ), + patch.object( + jax_train_entrypoint, + "update_deepmd_input", + side_effect=lambda x, **kwargs: x, + ), + patch.object( + jax_train_entrypoint, "normalize", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, "update_sel", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, + "SummaryPrinter", + return_value=Mock(__call__=Mock()), + ), + patch.object( + jax_train_entrypoint, "DPTrainer", return_value=trainer_instance + ) as mock_trainer, + patch.object(jax_train_entrypoint, "get_data", return_value=train_data), + patch.object(jax_train_entrypoint.dp_random, "seed"), + patch.object(jax_train_entrypoint, "serialize_from_file") as mock_serialize, + patch( + "builtins.open", + unittest.mock.mock_open(), + ), + patch.object(jax_train_entrypoint.json, "dump"), + ): + jax_train_entrypoint.train( + INPUT="input.json", + init_model="init_model.jax", + restart=None, + output="out.json", + init_frz_model="", + mpi_log="master", + log_level=2, + log_path=None, + use_pretrain_script=False, + ) + mock_serialize.assert_not_called() + self.assertEqual(mock_trainer.call_args.args[0]["model"], fake_jdata["model"]) + + def test_train_entrypoint_init_model_with_pretrain_script_uses_source_model( + self, + ) -> None: + fake_jdata = { + "model": {"type": "standard", "type_map": ["H"]}, + "learning_rate": {"start_lr": 1e-3, "decay_steps": 1, "stop_lr": 1e-8}, + "training": { + "numb_steps": 1, + "training_data": {"systems": []}, + }, + } + source_model = { + "type": "standard", + "type_map": ["O"], + "descriptor": {"type": "dpa3"}, + } + train_data = Mock() + train_data.type_map = ["O"] + trainer_instance = Mock() + trainer_instance.model.get_rcut.return_value = 6.0 + trainer_instance.model.get_type_map.return_value = ["O"] + trainer_instance.data_requirements = [] + + with ( + patch.object( + jax_train_entrypoint, "j_loader", return_value=deepcopy(fake_jdata) + ), + patch.object( + jax_train_entrypoint, + "update_deepmd_input", + side_effect=lambda x, **kwargs: x, + ), + patch.object( + jax_train_entrypoint, "normalize", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, "update_sel", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, + "SummaryPrinter", + return_value=Mock(__call__=Mock()), + ), + patch.object( + jax_train_entrypoint, "DPTrainer", return_value=trainer_instance + ) as mock_trainer, + patch.object(jax_train_entrypoint, "get_data", return_value=train_data), + patch.object(jax_train_entrypoint.dp_random, "seed"), + patch.object( + jax_train_entrypoint, + "serialize_from_file", + return_value={ + "model_def_script": deepcopy(source_model), + "model": {"unused": True}, + }, + ) as mock_serialize, + patch( + "builtins.open", + unittest.mock.mock_open(), + ), + patch.object(jax_train_entrypoint.json, "dump"), + ): + jax_train_entrypoint.train( + INPUT="input.json", + init_model="init_model.jax", + restart=None, + output="out.json", + init_frz_model="", + mpi_log="master", + log_level=2, + log_path=None, + use_pretrain_script=True, + ) + mock_serialize.assert_called_once_with("init_model.jax") + self.assertEqual(mock_trainer.call_args.args[0]["model"], source_model) + + def test_train_entrypoint_init_frz_model_with_pretrain_script_uses_source_model( + self, + ) -> None: + fake_jdata = { + "model": {"type": "standard", "type_map": ["H"]}, + "learning_rate": {"start_lr": 1e-3, "decay_steps": 1, "stop_lr": 1e-8}, + "training": { + "numb_steps": 1, + "training_data": {"systems": []}, + }, + } + source_model = { + "type": "standard", + "type_map": ["O"], + "descriptor": {"type": "dpa3"}, + } + train_data = Mock() + train_data.type_map = ["O"] + trainer_instance = Mock() + trainer_instance.model.get_rcut.return_value = 6.0 + trainer_instance.model.get_type_map.return_value = ["O"] + trainer_instance.data_requirements = [] + + with ( + patch.object( + jax_train_entrypoint, "j_loader", return_value=deepcopy(fake_jdata) + ), + patch.object( + jax_train_entrypoint, + "update_deepmd_input", + side_effect=lambda x, **kwargs: x, + ), + patch.object( + jax_train_entrypoint, "normalize", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, "update_sel", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, + "SummaryPrinter", + return_value=Mock(__call__=Mock()), + ), + patch.object( + jax_train_entrypoint, "DPTrainer", return_value=trainer_instance + ) as mock_trainer, + patch.object(jax_train_entrypoint, "get_data", return_value=train_data), + patch.object(jax_train_entrypoint.dp_random, "seed"), + patch.object( + jax_train_entrypoint, + "serialize_from_file", + return_value={ + "model_def_script": deepcopy(source_model), + "model": {"unused": True}, + }, + ) as mock_serialize, + patch( + "builtins.open", + unittest.mock.mock_open(), + ), + patch.object(jax_train_entrypoint.json, "dump"), + ): + jax_train_entrypoint.train( + INPUT="input.json", + init_model=None, + restart=None, + output="out.json", + init_frz_model="frozen_model.hlo", + mpi_log="master", + log_level=2, + log_path=None, + use_pretrain_script=True, + ) + mock_serialize.assert_called_once_with("frozen_model.hlo") + self.assertEqual(mock_trainer.call_args.args[0]["model"], source_model) + self.assertEqual( + mock_trainer.call_args.kwargs["init_frz_model"], "frozen_model.hlo" + ) + self.assertFalse(mock_trainer.call_args.kwargs["force_load"]) + + def test_train_entrypoint_passes_force_load(self) -> None: + fake_jdata = { + "model": {"type": "standard", "type_map": ["H"]}, + "learning_rate": {"start_lr": 1e-3, "decay_steps": 1, "stop_lr": 1e-8}, + "training": { + "numb_steps": 1, + "training_data": {"systems": []}, + }, + } + train_data = Mock() + train_data.type_map = ["H"] + trainer_instance = Mock() + trainer_instance.model.get_rcut.return_value = 6.0 + trainer_instance.model.get_type_map.return_value = ["H"] + trainer_instance.data_requirements = [] + + with ( + patch.object( + jax_train_entrypoint, "j_loader", return_value=deepcopy(fake_jdata) + ), + patch.object( + jax_train_entrypoint, + "update_deepmd_input", + side_effect=lambda x, **kwargs: x, + ), + patch.object( + jax_train_entrypoint, "normalize", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, "update_sel", side_effect=lambda x, **kwargs: x + ), + patch.object( + jax_train_entrypoint, + "SummaryPrinter", + return_value=Mock(__call__=Mock()), + ), + patch.object( + jax_train_entrypoint, "DPTrainer", return_value=trainer_instance + ) as mock_trainer, + patch.object(jax_train_entrypoint, "get_data", return_value=train_data), + patch.object(jax_train_entrypoint.dp_random, "seed"), + patch( + "builtins.open", + unittest.mock.mock_open(), + ), + patch.object(jax_train_entrypoint.json, "dump"), + ): + jax_train_entrypoint.train( + INPUT="input.json", + init_model="init_model.jax", + restart=None, + output="out.json", + init_frz_model="", + mpi_log="master", + log_level=2, + log_path=None, + force_load=True, + ) + self.assertTrue(mock_trainer.call_args.kwargs["force_load"]) + + def test_trainer_init_model_preserves_input_model_def_script(self) -> None: + dummy_model = Mock() + dummy_model.get_dim_fparam.return_value = 0 + source_model_script = {"type": "standard", "type_map": ["O"]} + target_jdata = { + "model": {"type": "standard", "type_map": ["H"]}, + "learning_rate": {"start_lr": 1e-3, "decay_steps": 1, "stop_lr": 1e-8}, + "training": {"numb_steps": 1}, + } + with ( + patch( + "deepmd.jax.train.trainer.get_model_for_wrapper", + return_value=dummy_model, + ), + patch( + "deepmd.jax.train.trainer.EnergyLoss.get_loss", + return_value=Mock(label_requirement=[]), + ), + patch( + "deepmd.jax.train.trainer.serialize_from_file", + return_value={ + "model_def_script": deepcopy(source_model_script), + "model": {"dummy": True}, + }, + ), + patch.object(DPTrainer, "_load_model_data"), + ): + trainer = DPTrainer( + deepcopy(target_jdata), + init_model="init_model.jax", + ) + self.assertEqual(trainer.model_def_script, target_jdata["model"]) + + def test_trainer_init_frz_model_preserves_input_model_def_script(self) -> None: + dummy_model = Mock() + dummy_model.get_dim_fparam.return_value = 0 + source_model_script = {"type": "standard", "type_map": ["O"]} + target_jdata = { + "model": {"type": "standard", "type_map": ["H"]}, + "learning_rate": {"start_lr": 1e-3, "decay_steps": 1, "stop_lr": 1e-8}, + "training": {"numb_steps": 1}, + } + with ( + patch( + "deepmd.jax.train.trainer.get_model_for_wrapper", + return_value=dummy_model, + ), + patch( + "deepmd.jax.train.trainer.EnergyLoss.get_loss", + return_value=Mock(label_requirement=[]), + ), + patch( + "deepmd.jax.train.trainer.serialize_from_file", + return_value={ + "model_def_script": deepcopy(source_model_script), + "model": {"dummy": True}, + }, + ), + patch.object(DPTrainer, "_load_frozen_model_data"), + ): + trainer = DPTrainer( + deepcopy(target_jdata), + init_frz_model="frozen_model.hlo", + ) + self.assertEqual(trainer.model_def_script, target_jdata["model"]) + + def test_merge_init_frz_model_data_copies_overlapping_arrays_only(self) -> None: + target = { + "descriptor": { + "w": np.array([[1.0, 2.0]], dtype=np.float32), + "config": {"seed": 1}, + }, + "fitting": { + "b": np.array([3.0], dtype=np.float32), + }, + } + source = { + "descriptor": { + "w": np.array([[9.0, 8.0]], dtype=np.float32), + "extra": np.array([5.0], dtype=np.float32), + }, + } + missing: list[tuple[object, ...]] = [] + unexpected: list[tuple[object, ...]] = [] + merged = _merge_init_frz_model_data( + target, + source, + missing=missing, + unexpected=unexpected, + ) + np.testing.assert_allclose(merged["descriptor"]["w"], source["descriptor"]["w"]) + self.assertEqual(merged["descriptor"]["config"], target["descriptor"]["config"]) + np.testing.assert_allclose(merged["fitting"]["b"], target["fitting"]["b"]) + self.assertIn(("descriptor", "extra"), unexpected) + self.assertIn(("descriptor", "config"), missing) + self.assertIn(("fitting",), missing) + + def test_merge_init_frz_model_data_raises_on_shape_mismatch(self) -> None: + with self.assertRaisesRegex(ValueError, "Shape mismatch"): + _merge_init_frz_model_data( + {"w": np.zeros((2, 2), dtype=np.float32)}, + {"w": np.zeros((3, 2), dtype=np.float32)}, + ) + + def test_merge_init_frz_model_data_reshapes_same_size_arrays(self) -> None: + merged = _merge_init_frz_model_data( + {"w": np.zeros((2, 1), dtype=np.float32)}, + {"w": np.array([1.0, 2.0], dtype=np.float32)}, + ) + np.testing.assert_allclose( + merged["w"], + np.array([[1.0], [2.0]], dtype=np.float32), + ) + + def test_trainer_load_model_data_force_load_reinitializes_missing_keys( + self, + ) -> None: + trainer = DPTrainer.__new__(DPTrainer) + trainer.force_load = True + trainer.shared_links = {} + trainer.case_embd_index = {} + trainer.multi_task = False + trainer.branch_has_hessian = {"Default": False} + trainer.model = Mock() + trainer.model.serialize.return_value = { + "descriptor": {"w": np.array([1.0], dtype=np.float32)}, + "fitting": {"b": np.array([2.0], dtype=np.float32)}, + } + with ( + patch( + "deepmd.jax.train.trainer.BaseModel.deserialize", + return_value="loaded_model", + ) as mock_deserialize, + patch.object(trainer, "_apply_hessian_flags"), + ): + trainer._load_model_data( + { + "model": { + "descriptor": {"w": np.array([9.0], dtype=np.float32)}, + } + } + ) + merged = mock_deserialize.call_args.args[0] + np.testing.assert_allclose( + merged["descriptor"]["w"], np.array([9.0], dtype=np.float32) + ) + np.testing.assert_allclose( + merged["fitting"]["b"], np.array([2.0], dtype=np.float32) + ) + self.assertEqual(trainer.model, "loaded_model") + + def test_apply_single_finetune_force_load_reinitializes_missing_pretrained_keys( + self, + ) -> None: + trainer = DPTrainer.__new__(DPTrainer) + trainer.force_load = True + target_model = Mock() + target_model.serialize.return_value = { + "descriptor": {"w": np.array([1.0], dtype=np.float32)}, + "fitting": {"b": np.array([2.0], dtype=np.float32)}, + } + finetune_rule = FinetuneRuleItem(["H"], ["H"]) + with ( + patch("deepmd.jax.train.trainer.BaseModel.deserialize") as mock_deserialize, + patch( + "deepmd.jax.train.trainer.merge_finetune_model_data", + side_effect=lambda target, pretrained, rule, source_overrides=None: ( + pretrained + ), + ), + ): + pretrained_model = Mock() + pretrained_model.get_type_map.return_value = ["H"] + mock_deserialize.side_effect = [pretrained_model, "merged_model"] + result = trainer._apply_single_finetune( + target_model, + {"descriptor": {"w": np.array([9.0], dtype=np.float32)}}, + finetune_rule, + ) + first_pretrained = mock_deserialize.call_args_list[0].args[0] + np.testing.assert_allclose( + first_pretrained["descriptor"]["w"], + np.array([9.0], dtype=np.float32), + ) + np.testing.assert_allclose( + first_pretrained["fitting"]["b"], + np.array([2.0], dtype=np.float32), + ) + self.assertEqual(result, "merged_model") + + def test_train_multi_calls_multitask_finetune(self) -> None: + DummyWrapper = self._dummy_wrapper_cls() + trainer = DPTrainer.__new__(DPTrainer) + trainer.multi_task = True + trainer.model_keys = ["head_a"] + trainer.finetune_model = "pretrained.jax" + trainer.finetune_links = {"head_a": FinetuneRuleItem(["H"], ["H"])} + trainer.init_model = None + trainer.restart = None + trainer.training_param = {"numb_steps": 1, "model_prob": {"head_a": 1.0}} + trainer.learning_rate_param = { + "start_lr": 1e-3, + "decay_steps": 1, + "stop_lr": 1e-8, + } + trainer.model_def_script = {"model_dict": {"head_a": {}}} + trainer.shared_links = {} + trainer.case_embd_index = {} + trainer.branch_has_hessian = {"head_a": False} + trainer.loss = {"head_a": Mock(label_requirement=[])} + trainer.model = DummyWrapper(head_a=Mock()) + trainer.model_prob = None + with ( + patch("deepmd.jax.train.trainer.ModelWrapper", DummyWrapper), + patch("deepmd.jax.train.trainer._clear_jax_mesh_for_host_ops"), + patch( + "deepmd.jax.train.trainer._resolve_model_prob_multi", + return_value=(np.array([1.0]), 1), + ), + patch.object( + trainer, + "_finetune_multi", + ) as mock_finetune, + patch( + "deepmd.jax.train.trainer.jax.make_mesh", + side_effect=RuntimeError("stop_after_finetune"), + ), + ): + train_data = {"head_a": Mock()} + with self.assertRaisesRegex(RuntimeError, "stop_after_finetune"): + trainer._train_multi(train_data, {}) + mock_finetune.assert_called_once_with(train_data) + + def test_finetune_multi_resuming_skips_bias_adjust(self) -> None: + DummyWrapper = self._dummy_wrapper_cls() + trainer = DPTrainer.__new__(DPTrainer) + trainer.model_keys = ["head_a"] + trainer.shared_links = {} + trainer.finetune_links = { + "head_a": FinetuneRuleItem( + ["H"], ["H"], model_branch="head_a", resuming=True + ) + } + trainer.finetune_model = "pretrained.jax" + trainer.finetune_model_data = { + "model": {"model_dict": {"head_a": {"source": "a"}}}, + } + trainer.branch_has_hessian = {"head_a": False} + trainer.case_embd_index = {} + trainer.data_bias_nsample = {"head_a": 10} + branch_model = Mock() + trainer.model = DummyWrapper(head_a=branch_model) + trainer._apply_hessian_flags = Mock() + merged_model = Mock() + merged_model.serialize.return_value = {"merged": True} + with ( + patch( + "deepmd.jax.train.trainer.ModelWrapper", + DummyWrapper, + ), + patch.object( + trainer, + "_validate_shared_finetune_rules", + ), + patch.object( + trainer, + "_collect_shared_source_overrides", + return_value={"head_a": {}}, + ), + patch.object( + trainer, + "_apply_single_finetune", + return_value=merged_model, + ) as mock_apply, + patch("deepmd.jax.train.trainer.model_change_out_bias") as mock_bias, + ): + trainer._finetune_multi({"head_a": Mock()}) + mock_apply.assert_called_once() + mock_bias.assert_not_called() + + def test_finetune_multi_non_resuming_adjusts_bias(self) -> None: + DummyWrapper = self._dummy_wrapper_cls() + trainer = DPTrainer.__new__(DPTrainer) + trainer.model_keys = ["head_a"] + trainer.shared_links = {} + trainer.finetune_links = { + "head_a": FinetuneRuleItem( + ["H"], ["H"], model_branch="head_a", random_fitting=True, resuming=False + ) + } + trainer.finetune_model = "pretrained.jax" + trainer.finetune_model_data = { + "model": {"model_dict": {"head_a": {"source": "a"}}}, + } + trainer.branch_has_hessian = {"head_a": False} + trainer.case_embd_index = {} + trainer.data_bias_nsample = {"head_a": 10} + branch_model = Mock() + trainer.model = DummyWrapper(head_a=branch_model) + trainer._apply_hessian_flags = Mock() + merged_model = Mock() + merged_model.serialize.return_value = {"merged": True} + with ( + patch( + "deepmd.jax.train.trainer.ModelWrapper", + DummyWrapper, + ), + patch.object( + trainer, + "_validate_shared_finetune_rules", + ), + patch.object( + trainer, + "_collect_shared_source_overrides", + return_value={"head_a": {}}, + ), + patch.object( + trainer, + "_apply_single_finetune", + return_value=merged_model, + ), + patch( + "deepmd.jax.train.trainer._pack_data_for_bias_adjust", + return_value={"coord": np.zeros((1, 1, 3))}, + ), + patch( + "deepmd.jax.train.trainer.model_change_out_bias", + return_value=merged_model, + ) as mock_bias, + ): + trainer._finetune_multi({"head_a": Mock()}) + self.assertEqual( + mock_bias.call_args.kwargs["bias_adjust_mode"], "set-by-statistic" + ) diff --git a/source/tests/jax/test_multitask.py b/source/tests/jax/test_multitask.py new file mode 100644 index 0000000000..df88b149c3 --- /dev/null +++ b/source/tests/jax/test_multitask.py @@ -0,0 +1,833 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import shutil +import sys +import tempfile +import unittest +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) +from types import ( + SimpleNamespace, +) +from unittest.mock import ( + Mock, + patch, +) + +import numpy as np + +from deepmd.dpmodel.descriptor.repflows import ( + _maybe_apply_jax_placeholder_sharding, +) +from deepmd.jax.entrypoints import freeze as jax_freeze_entrypoint +from deepmd.jax.entrypoints import train as jax_train_entrypoint +from deepmd.jax.model.model import ( + get_model_for_wrapper, +) +from deepmd.jax.model.multitask import ( + ModelWrapper, +) +from deepmd.jax.train.trainer import ( + _compute_multitask_data_stat, +) +from deepmd.jax.utils.multi_task import ( + get_case_embd_config, + preprocess_shared_params, +) +from deepmd.jax.utils.serialization import ( + select_model_branch, + serialize_from_file, +) + +from ..pt.model.test_permutation import ( + model_dpa3, + model_se_e2_a, +) + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestJAXMultiTaskHelpers(unittest.TestCase): + def test_preprocess_shared_params_and_case_embd(self) -> None: + model_config = { + "shared_dict": { + "shared_tm": ["O", "H"], + "shared_desc": {"type": "se_e2_a", "sel": [4, 4], "rcut": 6.0}, + "shared_fit": {"type": "ener", "neuron": [16], "dim_case_embd": 2}, + }, + "model_dict": { + "task_b": { + "type_map": "shared_tm", + "descriptor": "shared_desc", + "fitting_net": "shared_fit", + }, + "task_a": { + "type_map": "shared_tm", + "descriptor": "shared_desc", + "fitting_net": "shared_fit", + }, + }, + } + updated_model_config, shared_links = preprocess_shared_params( + deepcopy(model_config) + ) + self.assertIn("shared_desc", shared_links) + self.assertIn("shared_fit", shared_links) + enabled, case_embd_index = get_case_embd_config(updated_model_config) + self.assertTrue(enabled) + self.assertEqual(case_embd_index, {"task_a": 0, "task_b": 1}) + + def test_partial_shared_level_shares_only_type_embedding(self) -> None: + model_config = { + "shared_dict": { + "shared_tm": ["O", "H"], + "shared_desc": deepcopy(model_dpa3["descriptor"]), + }, + "model_dict": { + "task_a": { + "type_map": "shared_tm", + "descriptor": "shared_desc", + "fitting_net": { + "type": "ener", + **deepcopy(model_dpa3["fitting_net"]), + }, + }, + "task_b": { + "type_map": "shared_tm", + "descriptor": "shared_desc:1", + "fitting_net": { + "type": "ener", + **deepcopy(model_dpa3["fitting_net"]), + }, + }, + }, + } + updated_model_config, shared_links = preprocess_shared_params( + deepcopy(model_config) + ) + model = get_model_for_wrapper(updated_model_config, shared_links=shared_links) + desc_a = model["task_a"].atomic_model.descriptor + desc_b = model["task_b"].atomic_model.descriptor + self.assertIs(desc_a.type_embedding, desc_b.type_embedding) + self.assertIsNot(desc_a, desc_b) + self.assertIsNot(desc_a.repflows, desc_b.repflows) + + def test_hybrid_subcomponent_full_sharing_is_supported(self) -> None: + hybrid_descriptor = { + "type": "hybrid", + "list": [ + "shared_desc", + deepcopy(model_se_e2_a["descriptor"]), + ], + } + model_config = { + "shared_dict": { + "shared_tm": ["O", "H", "B"], + "shared_desc": deepcopy(model_se_e2_a["descriptor"]), + }, + "model_dict": { + "task_a": { + "type_map": "shared_tm", + "descriptor": deepcopy(hybrid_descriptor), + "fitting_net": { + "type": "ener", + **deepcopy(model_se_e2_a["fitting_net"]), + }, + }, + "task_b": { + "type_map": "shared_tm", + "descriptor": deepcopy(hybrid_descriptor), + "fitting_net": { + "type": "ener", + **deepcopy(model_se_e2_a["fitting_net"]), + }, + }, + }, + } + updated_model_config, shared_links = preprocess_shared_params( + deepcopy(model_config) + ) + model = get_model_for_wrapper(updated_model_config, shared_links=shared_links) + desc_a = model["task_a"].atomic_model.descriptor.descrpt_list[0] + desc_b = model["task_b"].atomic_model.descriptor.descrpt_list[0] + self.assertIs(desc_a, desc_b) + + def test_multitask_data_stat_merges_shared_components_once(self) -> None: + class FakeWrapper: + def __init__(self, model_dict: dict[str, SimpleNamespace]) -> None: + self._model_dict = model_dict + + def keys(self) -> list[str]: + return list(self._model_dict) + + def __getitem__(self, key: str) -> SimpleNamespace: + return self._model_dict[key] + + shared_descriptor = Mock() + shared_nets = object() + fitting_a = SimpleNamespace( + numb_fparam=0, + numb_aparam=0, + nets=shared_nets, + compute_output_stats=Mock(), + ) + fitting_b = SimpleNamespace( + numb_fparam=0, + numb_aparam=0, + nets=shared_nets, + compute_output_stats=Mock(), + ) + model_dict = { + "task_a": SimpleNamespace( + atomic_model=SimpleNamespace( + descriptor=shared_descriptor, + fitting=fitting_a, + ) + ), + "task_b": SimpleNamespace( + atomic_model=SimpleNamespace( + descriptor=shared_descriptor, + fitting=fitting_b, + ) + ), + } + wrapper = FakeWrapper(model_dict) + + train_data = { + "task_a": SimpleNamespace(mixed_type=False), + "task_b": SimpleNamespace(mixed_type=False), + } + with patch( + "deepmd.jax.train.trainer._build_single_data_stat", + side_effect=[ + ( + [{"coord": "a"}], + { + "energy": [[[np.array([2.0])]]], + "natoms_vec": [[np.array([0.0, 0.0, 1.0, 0.0])]], + }, + ), + ( + [{"coord": "b"}], + { + "energy": [[[np.array([4.0])]]], + "natoms_vec": [[np.array([0.0, 0.0, 0.0, 1.0])]], + }, + ), + ], + ): + _compute_multitask_data_stat( + wrapper, + train_data, + {"task_a": 0.2, "task_b": 0.8}, + {"task_a": 1e-2, "task_b": 1e-2}, + ) + shared_descriptor.compute_input_stats.assert_called_once_with( + [{"coord": "a"}, {"coord": "b"}] + ) + fitting_a.compute_output_stats.assert_called_once() + fitting_b.compute_output_stats.assert_called_once() + fitting_stat_a = fitting_a.compute_output_stats.call_args.args[0] + fitting_stat_b = fitting_b.compute_output_stats.call_args.args[0] + np.testing.assert_allclose(fitting_stat_a["energy"][0][0][0], np.array([2.0])) + np.testing.assert_allclose(fitting_stat_b["energy"][0][0][0], np.array([4.0])) + np.testing.assert_allclose( + fitting_stat_a["natoms_vec"][0][0], np.array([0.0, 0.0, 1.0, 0.0]) + ) + np.testing.assert_allclose( + fitting_stat_b["natoms_vec"][0][0], np.array([0.0, 0.0, 0.0, 1.0]) + ) + self.assertFalse(fitting_a.compute_output_stats.call_args.kwargs["mixed_type"]) + self.assertFalse(fitting_b.compute_output_stats.call_args.kwargs["mixed_type"]) + + def test_multitask_shared_fitting_input_stats_follow_weights_and_protection( + self, + ) -> None: + class FakeWrapper: + def __init__(self, model_dict: dict[str, SimpleNamespace]) -> None: + self._model_dict = model_dict + + def keys(self) -> list[str]: + return list(self._model_dict) + + def __getitem__(self, key: str) -> SimpleNamespace: + return self._model_dict[key] + + shared_descriptor = Mock() + shared_fitting = SimpleNamespace( + numb_fparam=1, + numb_aparam=1, + nets=object(), + fparam_avg=np.zeros(1, dtype=np.float64), + fparam_inv_std=np.ones(1, dtype=np.float64), + aparam_avg=np.zeros(1, dtype=np.float64), + aparam_inv_std=np.ones(1, dtype=np.float64), + compute_output_stats=Mock(), + ) + model_dict = { + "task_a": SimpleNamespace( + atomic_model=SimpleNamespace( + descriptor=shared_descriptor, + fitting=shared_fitting, + ) + ), + "task_b": SimpleNamespace( + atomic_model=SimpleNamespace( + descriptor=shared_descriptor, + fitting=shared_fitting, + ) + ), + } + wrapper = FakeWrapper(model_dict) + train_data = { + "task_a": SimpleNamespace(mixed_type=False), + "task_b": SimpleNamespace(mixed_type=False), + } + with patch( + "deepmd.jax.train.trainer._build_single_data_stat", + side_effect=[ + ( + [ + { + "coord": "a", + "fparam": np.array([[1.0], [3.0]]), + "aparam": np.array([[[2.0]], [[4.0]]]), + } + ], + { + "energy": [[[np.array([2.0])]]], + "natoms_vec": [[np.array([0.0, 0.0, 1.0, 0.0])]], + }, + ), + ( + [ + { + "coord": "b", + "fparam": np.array([[10.0], [14.0]]), + "aparam": np.array([[[12.0]], [[16.0]]]), + } + ], + { + "energy": [[[np.array([4.0])]]], + "natoms_vec": [[np.array([0.0, 0.0, 0.0, 1.0])]], + }, + ), + ], + ): + _compute_multitask_data_stat( + wrapper, + train_data, + {"task_a": 0.25, "task_b": 0.75}, + {"task_a": 0.5, "task_b": 0.5}, + ) + np.testing.assert_allclose(shared_fitting.fparam_avg, np.array([9.5])) + np.testing.assert_allclose( + shared_fitting.fparam_inv_std, + np.array([1.0 / np.sqrt(22.0)]), + ) + np.testing.assert_allclose(shared_fitting.aparam_avg, np.array([11.25])) + np.testing.assert_allclose( + shared_fitting.aparam_inv_std, + np.array([1.0 / np.sqrt(25.9375)]), + ) + + def test_multitask_shared_fitting_requires_same_data_stat_protect(self) -> None: + class FakeWrapper: + def __init__(self, model_dict: dict[str, SimpleNamespace]) -> None: + self._model_dict = model_dict + + def keys(self) -> list[str]: + return list(self._model_dict) + + def __getitem__(self, key: str) -> SimpleNamespace: + return self._model_dict[key] + + shared_descriptor = Mock() + shared_fitting = SimpleNamespace( + numb_fparam=0, + numb_aparam=0, + nets=object(), + compute_output_stats=Mock(), + ) + model_dict = { + "task_a": SimpleNamespace( + atomic_model=SimpleNamespace( + descriptor=shared_descriptor, + fitting=shared_fitting, + ) + ), + "task_b": SimpleNamespace( + atomic_model=SimpleNamespace( + descriptor=shared_descriptor, + fitting=shared_fitting, + ) + ), + } + wrapper = FakeWrapper(model_dict) + train_data = { + "task_a": SimpleNamespace(mixed_type=False), + "task_b": SimpleNamespace(mixed_type=False), + } + with patch( + "deepmd.jax.train.trainer._build_single_data_stat", + side_effect=[ + ( + [{"coord": "a"}], + { + "energy": [[[np.array([1.0])]]], + "natoms_vec": [[np.array([0.0, 0.0, 1.0, 0.0])]], + }, + ), + ( + [{"coord": "b"}], + { + "energy": [[[np.array([1.0])]]], + "natoms_vec": [[np.array([0.0, 0.0, 1.0, 0.0])]], + }, + ), + ], + ): + with self.assertRaisesRegex(ValueError, "data_stat_protect"): + _compute_multitask_data_stat( + wrapper, + train_data, + {"task_a": 0.5, "task_b": 0.5}, + {"task_a": 1e-2, "task_b": 1e-1}, + ) + + def test_select_model_branch_resolves_alias(self) -> None: + data = { + "model_def_script": { + "model_dict": { + "branch_a": { + "type": "standard", + "model_branch_alias": ["A"], + }, + "branch_b": { + "type": "standard", + }, + } + }, + "model": { + "model_dict": { + "branch_a": {"value": "a"}, + "branch_b": {"value": "b"}, + } + }, + } + selected = select_model_branch(deepcopy(data), "A") + self.assertEqual(selected["model_def_script"]["type"], "standard") + self.assertEqual(selected["model"]["value"], "a") + + def test_select_model_branch_requires_explicit_head_for_multitask(self) -> None: + data = { + "model_def_script": { + "model_dict": { + "branch_a": {"type": "standard"}, + } + }, + "model": { + "model_dict": { + "branch_a": {"value": "a"}, + } + }, + } + with self.assertRaisesRegex(ValueError, "--head/--model-branch"): + select_model_branch(deepcopy(data), None) + + def test_placeholder_sharding_skips_missing_mesh_context(self) -> None: + placeholder = object() + with patch( + "deepmd.jax.env.jax.lax.with_sharding_constraint", + side_effect=RuntimeError("requires a non-empty mesh in context"), + ): + self.assertIs( + _maybe_apply_jax_placeholder_sharding(placeholder), + placeholder, + ) + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestJAXMultiTaskTraining(unittest.TestCase): + def setUp(self) -> None: + self.repo_root = Path(__file__).resolve().parents[3] + self.pt_water_dir = self.repo_root / "source/tests/pt/model/water" + self.data_dir = self.repo_root / "source/tests/pt/water/data/data_0" + self.tmpdir = Path(tempfile.mkdtemp(prefix="jax_mt_test_")) + self.prev_cwd = Path.cwd() + os.chdir(self.tmpdir) + + def tearDown(self) -> None: + os.chdir(self.prev_cwd) + shutil.rmtree(self.tmpdir) + + def _load_template(self, filename: str) -> dict: + with open(self.pt_water_dir / filename) as fp: + return json.load(fp) + + def _write_config(self, filename: str, config: dict) -> Path: + path = self.tmpdir / filename + with open(path, "w") as fp: + json.dump(config, fp) + return path + + def _base_multitask_config( + self, *, descriptor: dict, sharefit: bool = False + ) -> dict: + template_name = "multitask_sharefit.json" if sharefit else "multitask.json" + config = self._load_template(template_name) + config["model"]["shared_dict"]["my_descriptor"] = deepcopy(descriptor) + if sharefit: + config["model"]["shared_dict"]["my_fitting"]["seed"] = 1 + else: + for model_key in config["model"]["model_dict"]: + config["model"]["model_dict"][model_key]["fitting_net"]["seed"] = 1 + for model_key in config["model"]["model_dict"]: + config["model"]["model_dict"][model_key]["data_stat_nbatch"] = 1 + config["training"]["data_dict"][model_key]["training_data"]["systems"] = [ + str(self.data_dir) + ] + config["training"]["data_dict"][model_key]["validation_data"]["systems"] = [ + str(self.data_dir) + ] + config["training"]["data_dict"][model_key]["stat_file"] = str( + self.tmpdir / f"{model_key}.hdf5" + ) + config["training"]["numb_steps"] = 1 + config["training"]["save_freq"] = 1 + config["training"]["disp_freq"] = 1 + config["training"]["save_ckpt"] = "model.ckpt" + config["training"]["disp_file"] = "lcurve.out" + return config + + def _run_entrypoint(self, input_path: Path, *, restart: str | None = None) -> None: + jax_train_entrypoint.train( + INPUT=str(input_path), + init_model=None, + restart=restart, + output="out.json", + init_frz_model="", + mpi_log="master", + log_level=2, + log_path=None, + skip_neighbor_stat=True, + finetune=None, + use_pretrain_script=False, + ) + + def _load_wrapper(self, ckpt_path: Path) -> tuple[dict, ModelWrapper]: + serialized = serialize_from_file(str(ckpt_path)) + model_def_script = serialized["model_def_script"] + shared_links = model_def_script.get("shared_links", {}) + _, case_embd_index = get_case_embd_config(model_def_script) + wrapper = ModelWrapper.deserialize( + serialized["model"], + shared_links=shared_links, + case_embd_index=case_embd_index, + ) + return serialized, wrapper + + def test_entrypoint_multitask_train_and_restart_se_e2_a(self) -> None: + config = self._base_multitask_config( + descriptor=deepcopy(model_se_e2_a["descriptor"]) + ) + input_path = self._write_config("multitask_se_e2_a.json", config) + + self._run_entrypoint(input_path) + ckpt_path = self.tmpdir / "model.ckpt-1.jax" + self.assertTrue(ckpt_path.is_dir()) + + serialized, wrapper = self._load_wrapper(ckpt_path) + self.assertEqual( + set(serialized["model_def_script"]["model_dict"].keys()), + {"model_1", "model_2"}, + ) + self.assertIn("shared_links", serialized["model_def_script"]) + self.assertIs( + wrapper["model_1"].atomic_model.descriptor, + wrapper["model_2"].atomic_model.descriptor, + ) + + self._run_entrypoint(input_path, restart=str(ckpt_path)) + restarted = serialize_from_file(str(ckpt_path)) + self.assertEqual(restarted["@variables"]["current_step"], 1) + + def test_sharefit_case_embd_checkpoint_restores_shared_objects(self) -> None: + config = self._base_multitask_config( + descriptor=deepcopy(model_se_e2_a["descriptor"]), + sharefit=True, + ) + input_path = self._write_config("multitask_sharefit.json", config) + + self._run_entrypoint(input_path) + serialized, wrapper = self._load_wrapper(self.tmpdir / "model.ckpt-1.jax") + + self.assertIs( + wrapper["model_1"].atomic_model.descriptor, + wrapper["model_2"].atomic_model.descriptor, + ) + self.assertIsNot( + wrapper["model_1"].atomic_model.fitting, + wrapper["model_2"].atomic_model.fitting, + ) + self.assertIs( + wrapper["model_1"].atomic_model.fitting.nets, + wrapper["model_2"].atomic_model.fitting.nets, + ) + self.assertIs( + wrapper["model_1"].atomic_model.fitting.fparam_avg, + wrapper["model_2"].atomic_model.fitting.fparam_avg, + ) + self.assertIs( + wrapper["model_1"].atomic_model.fitting.fparam_inv_std, + wrapper["model_2"].atomic_model.fitting.fparam_inv_std, + ) + self.assertIsNot( + wrapper["model_1"].atomic_model.fitting.bias_atom_e, + wrapper["model_2"].atomic_model.fitting.bias_atom_e, + ) + self.assertEqual( + serialized["model_def_script"]["shared_links"]["my_fitting"]["links"][0][ + "shared_type" + ], + "fitting_net", + ) + + wrapper.set_case_embd("model_1") + case_embd = wrapper["model_1"].atomic_model.fitting.serialize()["@variables"][ + "case_embd" + ] + np.testing.assert_array_equal(case_embd, np.array([1.0, 0.0])) + wrapper.set_case_embd("model_2") + case_embd = wrapper["model_2"].atomic_model.fitting.serialize()["@variables"][ + "case_embd" + ] + np.testing.assert_array_equal(case_embd, np.array([0.0, 1.0])) + case_embd_model_1 = wrapper["model_1"].atomic_model.fitting.serialize()[ + "@variables" + ]["case_embd"] + np.testing.assert_array_equal(case_embd_model_1, np.array([1.0, 0.0])) + + def test_entrypoint_multitask_train_dpa3(self) -> None: + config = self._base_multitask_config( + descriptor=deepcopy(model_dpa3["descriptor"]) + ) + for model_key in config["model"]["model_dict"]: + config["model"]["model_dict"][model_key]["fitting_net"] = deepcopy( + model_dpa3["fitting_net"] + ) + input_path = self._write_config("multitask_dpa3.json", config) + + self._run_entrypoint(input_path) + serialized, wrapper = self._load_wrapper(self.tmpdir / "model.ckpt-1.jax") + self.assertEqual(serialized["@variables"]["current_step"], 1) + self.assertIs( + wrapper["model_1"].atomic_model.descriptor, + wrapper["model_2"].atomic_model.descriptor, + ) + + def test_multitask_checkpoint_rejects_single_task_restart(self) -> None: + input_path = self._write_config( + "multitask_reject.json", + self._base_multitask_config( + descriptor=deepcopy(model_se_e2_a["descriptor"]) + ), + ) + self._run_entrypoint(input_path) + ckpt_path = self.tmpdir / "model.ckpt-1.jax" + + single_task_config = { + "model": deepcopy(model_se_e2_a), + "learning_rate": { + "type": "exp", + "start_lr": 1e-3, + "decay_steps": 1, + "stop_lr": 1e-8, + }, + "loss": {"type": "ener"}, + "training": {"numb_steps": 1}, + } + single_task_config["model"]["fitting_net"].setdefault("type", "ener") + from deepmd.jax.train.trainer import ( + DPTrainer, + ) + + with self.assertRaisesRegex( + ValueError, "single-task JAX target does not accept a multitask checkpoint" + ): + DPTrainer(single_task_config, restart=str(ckpt_path)) + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestJAXFreezeMultiTask(unittest.TestCase): + def test_freeze_jax_selects_branch(self) -> None: + data = { + "model_def_script": { + "model_dict": { + "branch_a": { + "type": "standard", + "model_branch_alias": ["A"], + }, + "branch_b": { + "type": "standard", + }, + } + }, + "model": { + "model_dict": { + "branch_a": {"value": "a"}, + "branch_b": {"value": "b"}, + } + }, + } + with ( + patch.object( + jax_freeze_entrypoint, + "serialize_from_file", + return_value=deepcopy(data), + ), + patch.object( + jax_freeze_entrypoint, + "deserialize_to_file", + ) as mock_deserialize, + patch( + "deepmd.jax.entrypoints.freeze.Path.is_dir", + return_value=True, + ), + ): + jax_freeze_entrypoint.freeze( + checkpoint_folder="ckpt_dir", + output="frozen.jax", + head="A", + ) + selected_data = mock_deserialize.call_args.args[1] + self.assertEqual(selected_data["model"]["value"], "a") + self.assertNotIn("model_dict", selected_data["model_def_script"]) + + def test_freeze_jax_hessian_selects_branch(self) -> None: + data = { + "model_def_script": { + "model_dict": { + "branch_a": { + "type": "standard", + "model_branch_alias": ["A"], + }, + } + }, + "model": { + "model_dict": { + "branch_a": {"value": "a"}, + } + }, + } + with ( + patch.object( + jax_freeze_entrypoint, + "serialize_from_file", + return_value=deepcopy(data), + ), + patch.object( + jax_freeze_entrypoint, + "deserialize_to_file", + ) as mock_deserialize, + patch( + "deepmd.jax.entrypoints.freeze.Path.is_dir", + return_value=True, + ), + ): + jax_freeze_entrypoint.freeze( + checkpoint_folder="ckpt_dir", + output="frozen.jax", + head="A", + hessian=True, + ) + selected_data = mock_deserialize.call_args.args[1] + self.assertEqual(selected_data["model"]["value"], "a") + self.assertTrue(mock_deserialize.call_args.kwargs["hessian"]) + + def test_freeze_hlo_selects_branch(self) -> None: + data = { + "model_def_script": { + "model_dict": { + "branch_a": { + "type": "standard", + "model_branch_alias": ["A"], + }, + "branch_b": { + "type": "standard", + }, + } + }, + "model": { + "model_dict": { + "branch_a": {"value": "a"}, + "branch_b": {"value": "b"}, + } + }, + } + with ( + patch.object( + jax_freeze_entrypoint, + "serialize_from_file", + return_value=deepcopy(data), + ), + patch.object( + jax_freeze_entrypoint, + "deserialize_to_file", + ) as mock_deserialize, + patch( + "deepmd.jax.entrypoints.freeze.Path.is_dir", + return_value=True, + ), + ): + jax_freeze_entrypoint.freeze( + checkpoint_folder="ckpt_dir", + output="frozen.hlo", + head="A", + ) + selected_data = mock_deserialize.call_args.args[1] + self.assertEqual(selected_data["model"]["value"], "a") + self.assertNotIn("model_dict", selected_data["model_def_script"]) + + def test_freeze_hlo_rejects_multitask_without_head(self) -> None: + data = { + "model_def_script": { + "model_dict": { + "branch_a": {"type": "standard"}, + } + }, + "model": { + "model_dict": { + "branch_a": {"value": "a"}, + } + }, + } + with ( + patch.object( + jax_freeze_entrypoint, + "serialize_from_file", + return_value=deepcopy(data), + ), + patch( + "deepmd.jax.entrypoints.freeze.Path.is_dir", + return_value=True, + ), + ): + with self.assertRaisesRegex(ValueError, "--head/--model-branch"): + jax_freeze_entrypoint.freeze( + checkpoint_folder="ckpt_dir", + output="frozen.hlo", + )