diff --git a/deepmd/dpmodel/fitting/dpa4_ener.py b/deepmd/dpmodel/fitting/dpa4_ener.py index f08097f032..6f7f970947 100644 --- a/deepmd/dpmodel/fitting/dpa4_ener.py +++ b/deepmd/dpmodel/fitting/dpa4_ener.py @@ -92,11 +92,18 @@ def __init__( ) if neuron is None: neuron = [] - if isinstance(trainable, list): - trainable = all(trainable) self.in_dim = int(in_dim) self.out_dim = int(out_dim) self.neuron = [int(nn_dim) for nn_dim in neuron] + if isinstance(trainable, bool): + self.trainable = [trainable] * (len(self.neuron) + 1) + else: + self.trainable = [bool(flag) for flag in trainable] + if len(self.trainable) != len(self.neuron) + 1: + raise ValueError( + "trainable must contain one flag per hidden layer plus " + "one flag for the output layer" + ) self.activation_function = activation_function self.resnet_dt = bool(resnet_dt) self.precision = precision @@ -123,7 +130,7 @@ def __init__( resnet=False, precision=self.precision, seed=child_seed(seed, layer_idx), - trainable=trainable, + trainable=self.trainable[layer_idx], ) ) dim_in = hidden_dim @@ -139,7 +146,7 @@ def __init__( resnet=False, precision=self.precision, seed=child_seed(seed, len(self.neuron) + int(self.case_film_embd)), - trainable=trainable, + trainable=self.trainable[-1], ) def call_until_last(self, xx: Array) -> Array: @@ -181,6 +188,9 @@ def serialize(self) -> dict[str, Any]: "descriptor_dim": self.descriptor_dim, "dim_case_embd": self.dim_case_embd, "case_film_embd": self.case_film_embd, + # Preserve the effective per-layer freeze policy when backend + # wrappers rebuild this network from its serialized form. + "trainable": self.trainable.copy(), "@variables": variables, } diff --git a/deepmd/jax/descriptor/__init__.py b/deepmd/jax/descriptor/__init__.py index cda2faf24d..f88a8f6589 100644 --- a/deepmd/jax/descriptor/__init__.py +++ b/deepmd/jax/descriptor/__init__.py @@ -8,6 +8,9 @@ from deepmd.jax.descriptor.dpa3 import ( DescrptDPA3, ) +from deepmd.jax.descriptor.dpa4 import ( + DescrptDPA4, +) from deepmd.jax.descriptor.hybrid import ( DescrptHybrid, ) @@ -31,6 +34,7 @@ "DescrptDPA1", "DescrptDPA2", "DescrptDPA3", + "DescrptDPA4", "DescrptHybrid", "DescrptSeA", "DescrptSeAttenV2", diff --git a/deepmd/jax/descriptor/dpa4.py b/deepmd/jax/descriptor/dpa4.py new file mode 100644 index 0000000000..cda6507599 --- /dev/null +++ b/deepmd/jax/descriptor/dpa4.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Mapping, + Sequence, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4DP +from deepmd.dpmodel.descriptor.dpa4_nn.activation import SwiGLU as SwiGLUDP +from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import GridProduct as GridProductDP +from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + C3CutoffEnvelope as C3CutoffEnvelopeDP, +) +from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialMLP as RadialMLPDP +from deepmd.dpmodel.descriptor.dpa4_nn.so2 import SO2Linear as SO2LinearDP +from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + WignerDCalculator as WignerDCalculatorDP, +) +from deepmd.jax.common import ( + ArrayAPIVariable, + flax_module, + register_dpmodel_mapping, + to_jax_array, + try_convert_module, +) +from deepmd.jax.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.jax.env import ( + jnp, + nnx, +) +from deepmd.jax.utils.network import ( + ArrayAPIParam, +) + + +@flax_module +class SwiGLU(SwiGLUDP): + pass + + +register_dpmodel_mapping(SwiGLUDP, lambda v: SwiGLU()) + + +@flax_module +class C3CutoffEnvelope(C3CutoffEnvelopeDP): + pass + + +register_dpmodel_mapping( + C3CutoffEnvelopeDP, + lambda v: C3CutoffEnvelope(v.rcut, v.p, precision=v.precision), +) + + +@flax_module +class RadialMLP(RadialMLPDP): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + converted = [self._convert_layer(layer) for layer in self.net] + self.net = nnx.List(converted) if hasattr(nnx, "List") else converted + + @staticmethod + def _convert_layer(layer: Any) -> Any: + if isinstance(layer, nnx.Module): + return layer + if isinstance(layer, NativeOP): + converted = try_convert_module(layer) + if converted is not None: + return converted + return layer + + +register_dpmodel_mapping( + RadialMLPDP, + lambda v: RadialMLP.deserialize(v.serialize()), +) + + +@flax_module +class GridProduct(GridProductDP): + pass + + +register_dpmodel_mapping(GridProductDP, lambda v: GridProduct()) + + +@flax_module +class WignerDCalculator(WignerDCalculatorDP): + pass + + +register_dpmodel_mapping( + WignerDCalculatorDP, + lambda v: WignerDCalculator(v.lmax, eps=v.eps, precision=v.precision), +) + + +_TRAINABLE_ATTRS: dict[str, tuple[str, ...]] = { + "RMSNorm": ("adam_scale",), + "EquivariantRMSNorm": ("adam_scale", "bias"), + "ReducedEquivariantRMSNorm": ("adam_scale", "bias0"), + "ScalarRMSNorm": ("adam_scale",), + "RadialBasis": ("adam_freqs",), + "SO3Linear": ("weight", "bias"), + "FocusLinear": ("weight", "bias"), + "ChannelLinear": ("weight", "bias"), + "FrameContract": ("weight",), + "FrameExpand": ("weight",), + "SO2Linear": ("weight_m0", "bias0"), + "DynamicRadialDegreeMixer": ("weight", "channel_basis"), + "SO2Convolution": ( + "adamw_attn_logit_w", + "adamw_attn_z_bias_raw", + "adamw_attn_gate_w", + "adamw_focus_compete_w", + "focus_compete_bias", + ), + "SeZMTypeEmbedding": ("adam_type_embedding",), + "SpinEmbedding": ("adam_spin_vec_weight", "adam_spin_nbr_weight"), + "EnvironmentInitialEmbedding": ("spin_scale",), + "DepthAttnRes": ("adamw_pseudo_query",), + "S2GridNet": ("residual_scale",), + "SO3GridNet": ("residual_scale",), + "DescrptDPA4": ("film_scale_strength_log", "film_shift_strength_log"), +} + +_TRAINABLE_LIST_ATTRS: dict[str, tuple[str, ...]] = { + "SeZMInteractionBlock": ("adam_ffn_layer_scales",), + "SO2Linear": ("weight_m",), + "SO2Convolution": ("adam_so2_layer_scales",), +} + + +def _is_array_like(value: Any) -> bool: + return hasattr(value, "shape") and hasattr(value, "dtype") + + +def _array_value(value: Any) -> Any: + if isinstance(value, nnx.Variable): + return value.value + return value + + +def _is_floating_array(value: Any) -> bool: + value = _array_value(value) + if value is None or not _is_array_like(value): + return False + return bool(jnp.issubdtype(value.dtype, jnp.floating)) + + +def _as_parameter_variable(value: Any, *, trainable: bool) -> Any: + """Track a floating parameter with the requested optimizer visibility.""" + variable_type = ArrayAPIParam if trainable else ArrayAPIVariable + if type(value) is variable_type: + return value + if not _is_floating_array(value): + return value + if isinstance(value, nnx.Variable): + value = value.value + if isinstance(value, np.ndarray): + value = to_jax_array(value) + return variable_type(value) + + +def _as_parameter_variable_list(value: Any, *, trainable: bool) -> Any: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + return value + promoted = [] + changed = False + for item in value: + new_item = _as_parameter_variable(item, trainable=trainable) + promoted.append(new_item) + changed = changed or new_item is not item + if not changed: + return value + return nnx.List(promoted) if hasattr(nnx, "List") else promoted + + +def _iter_object_tree(root: Any) -> Any: + seen: set[int] = set() + + def visit(value: Any) -> Any: + if value is None or isinstance(value, (str, bytes, int, float, bool)): + return + if _is_array_like(value): + return + value_id = id(value) + if value_id in seen: + return + seen.add(value_id) + + if isinstance(value, Mapping): + for item in value.values(): + yield from visit(item) + return + if isinstance(value, Sequence): + for item in value: + yield from visit(item) + return + try: + value_dict = object.__getattribute__(value, "__dict__") + except AttributeError: + return + + yield value + for item in value_dict.values(): + yield from visit(item) + + yield from visit(root) + + +def _promote_parameters( + module: Any, names: tuple[str, ...], *, trainable: bool +) -> None: + for name in names: + if not hasattr(module, name): + continue + value = getattr(module, name) + new_value = _as_parameter_variable(value, trainable=trainable) + if new_value is not value: + setattr(module, name, new_value) + + +def _promote_parameter_lists( + module: Any, names: tuple[str, ...], *, trainable: bool +) -> None: + for name in names: + if not hasattr(module, name): + continue + value = getattr(module, name) + new_value = _as_parameter_variable_list(value, trainable=trainable) + if new_value is not value: + setattr(module, name, new_value) + + +def _promote_trainable_tree(module: Any) -> Any: + root_trainable = bool(getattr(module, "trainable", True)) + for submodule in _iter_object_tree(module): + # A frozen descriptor freezes every descendant, including helper + # modules such as RadialBasis that do not carry a local flag. + trainable = root_trainable and bool(getattr(submodule, "trainable", True)) + names = _TRAINABLE_ATTRS.get(type(submodule).__name__) + if names is not None: + _promote_parameters(submodule, names, trainable=trainable) + list_names = _TRAINABLE_LIST_ATTRS.get(type(submodule).__name__) + if list_names is not None: + _promote_parameter_lists(submodule, list_names, trainable=trainable) + return module + + +@flax_module +class SO2Linear(SO2LinearDP): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.weight_m = _as_parameter_variable_list( + self.weight_m, trainable=bool(self.trainable) + ) + + @classmethod + def deserialize(cls, data: dict) -> "SO2Linear": + obj = super().deserialize(data) + obj.weight_m = _as_parameter_variable_list( + obj.weight_m, trainable=bool(obj.trainable) + ) + return obj + + +register_dpmodel_mapping( + SO2LinearDP, + lambda v: SO2Linear.deserialize(v.serialize()), +) + + +@BaseDescriptor.register("SeZM") +@BaseDescriptor.register("sezm") +@BaseDescriptor.register("DPA4") +@BaseDescriptor.register("dpa4") +@flax_module +class DescrptDPA4(DescrptDPA4DP): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + _promote_trainable_tree(self) + + @classmethod + def deserialize(cls, data: dict) -> "DescrptDPA4": + obj = super().deserialize(data) + return _promote_trainable_tree(obj) diff --git a/deepmd/jax/fitting/__init__.py b/deepmd/jax/fitting/__init__.py index 77133e2bac..82cea44e1b 100644 --- a/deepmd/jax/fitting/__init__.py +++ b/deepmd/jax/fitting/__init__.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.jax.fitting.dpa4_ener import ( + SeZMEnergyFittingNet, +) from deepmd.jax.fitting.fitting import ( DipoleFittingNet, DOSFittingNet, @@ -11,4 +14,5 @@ "DipoleFittingNet", "EnergyFittingNet", "PolarFittingNet", + "SeZMEnergyFittingNet", ] diff --git a/deepmd/jax/fitting/dpa4_ener.py b/deepmd/jax/fitting/dpa4_ener.py new file mode 100644 index 0000000000..425268cd78 --- /dev/null +++ b/deepmd/jax/fitting/dpa4_ener.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + ClassVar, +) + +from deepmd.dpmodel.fitting.dpa4_ener import GLUFittingNet as GLUFittingNetDP +from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMEnergyFittingNet as SeZMEnergyFittingNetDP, +) +from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMNetworkCollection as SeZMNetworkCollectionDP, +) +from deepmd.jax.common import ( + flax_module, + register_dpmodel_mapping, +) +from deepmd.jax.fitting.base_fitting import ( + BaseFitting, +) + + +@flax_module +class GLUFittingNet(GLUFittingNetDP): + pass + + +register_dpmodel_mapping( + GLUFittingNetDP, + lambda v: GLUFittingNet.deserialize(v.serialize()), +) + + +@flax_module +class SeZMNetworkCollection(SeZMNetworkCollectionDP): + _jax_data_list_attrs: ClassVar[set[str]] = {"_networks", "networks"} + NETWORK_TYPE_MAP: ClassVar[dict[str, type]] = { + "sezm_fitting_network": GLUFittingNet, + } + + +register_dpmodel_mapping( + SeZMNetworkCollectionDP, + lambda v: SeZMNetworkCollection.deserialize(v.serialize()), +) + + +@BaseFitting.register("dpa4_ener") +@BaseFitting.register("sezm_ener") +@flax_module +class SeZMEnergyFittingNet(SeZMEnergyFittingNetDP): + pass diff --git a/deepmd/jax/model/base_model.py b/deepmd/jax/model/base_model.py index 8a5a55e8d7..1a341da902 100644 --- a/deepmd/jax/model/base_model.py +++ b/deepmd/jax/model/base_model.py @@ -1,5 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + from deepmd.dpmodel.model.base_model import ( make_base_model, ) @@ -12,8 +16,67 @@ jax, jnp, ) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +class BaseModel(make_base_model()): + """JAX model registry with adapters for regular PT SeZM checkpoints.""" + + _SEZM_MODEL_TYPES = frozenset({"sezm", "dpa4"}) + _SEZM_ATOMIC_TYPES = frozenset({"sezm_atomic"}) -BaseModel = make_base_model() + @classmethod + def deserialize(cls, data: dict[str, Any]) -> "BaseModel": + model_type = str(data.get("type", "standard")).lower() + if model_type in cls._SEZM_MODEL_TYPES: + return cls.deserialize(cls._unwrap_pt_sezm_model(data)) + if model_type in cls._SEZM_ATOMIC_TYPES: + return cls.deserialize(cls._normalize_pt_sezm_atomic(data)) + return super().deserialize(data) + + @staticmethod + def _unwrap_pt_sezm_model(data: dict[str, Any]) -> dict[str, Any]: + """Unwrap PT's model-level SeZM schema after validating its extras.""" + check_version_compatibility(int(data.get("@version", 1)), 1, 1) + if str(data.get("bridging_method", "none")).lower() not in ("none", ""): + raise NotImplementedError( + "PT SeZM/DPA4 checkpoints with bridging are not supported in JAX." + ) + if data.get("lora") is not None: + raise NotImplementedError( + "PT SeZM/DPA4 checkpoints with LoRA are not supported in JAX." + ) + atomic_model = data.get("atomic_model") + if atomic_model is None: + raise ValueError("SeZM/DPA4 model data is missing 'atomic_model'.") + return atomic_model + + @staticmethod + def _normalize_pt_sezm_atomic(data: dict[str, Any]) -> dict[str, Any]: + """Convert PT's energy-only ``sezm_atomic`` schema to ``standard``.""" + data = data.copy() + check_version_compatibility(int(data.get("@version", 2)), 3, 2) + if data.pop("dens_fitting", None) is not None: + raise NotImplementedError( + "PT SeZM/DPA4 checkpoints with a dens head are not supported in JAX." + ) + active_mode = data.pop("active_mode", None) + if active_mode not in (None, "ener"): + raise NotImplementedError( + f"PT SeZM/DPA4 active_mode {active_mode!r} is not supported in JAX." + ) + variables = data.get("@variables") + if isinstance(variables, dict): + data["@variables"] = { + key: value + for key, value in variables.items() + if key in ("out_bias", "out_std") + } + data["@version"] = 2 + data["type"] = "standard" + return data def forward_common_atomic( diff --git a/deepmd/jax/model/ener_model.py b/deepmd/jax/model/ener_model.py index 1d3e8a1d80..626997d18e 100644 --- a/deepmd/jax/model/ener_model.py +++ b/deepmd/jax/model/ener_model.py @@ -11,6 +11,8 @@ ) +@BaseModel.register("sezm_ener") +@BaseModel.register("dpa4_ener") @BaseModel.register("ener") class EnergyModel(make_jax_dp_model_from_dpmodel(EnergyModelDP, DPAtomicModelEnergy)): pass diff --git a/deepmd/jax/model/model.py b/deepmd/jax/model/model.py index a3d067c636..31ddc12073 100644 --- a/deepmd/jax/model/model.py +++ b/deepmd/jax/model/model.py @@ -109,6 +109,57 @@ def get_zbl_model(data: dict) -> DPZBLModel: ) +def get_sezm_model(data: dict) -> BaseModel: + """Build a DPA4/SeZM energy model from the pt-style model config.""" + data = deepcopy(data) + if "spin" in data: + raise NotImplementedError("Spin DPA4/SeZM models are not supported in JAX.") + if str(data.get("bridging_method", "none")).lower() != "none": + raise NotImplementedError("DPA4/SeZM bridging is not supported in JAX.") + if data.get("lora") is not None: + raise NotImplementedError("DPA4/SeZM LoRA is not supported in JAX.") + if data.get("use_compile"): + raise NotImplementedError("model.use_compile is not supported in JAX.") + if data.get("preset_out_bias"): + raise NotImplementedError("DPA4/SeZM preset_out_bias is not supported in JAX.") + + data.pop("type", None) + data["descriptor"] = data.get("descriptor") or {} + data["fitting_net"] = data.get("fitting_net") or {} + data["descriptor"].setdefault("type", "dpa4") + data["fitting_net"].setdefault("type", "dpa4_ener") + if data["descriptor"].get("add_chg_spin_ebd"): + raise NotImplementedError( + "DPA4/SeZM charge/spin conditioning is not supported in JAX." + ) + if data["descriptor"]["type"] not in ("dpa4", "DPA4", "sezm", "SeZM"): + raise ValueError( + "Model type 'dpa4' requires a DPA4/SeZM descriptor, but got " + f"descriptor type '{data['descriptor']['type']}'." + ) + if data["fitting_net"]["type"] not in ("dpa4_ener", "sezm_ener"): + raise ValueError( + "Model type 'dpa4' requires the DPA4/SeZM energy fitting net, but got " + f"fitting_net type '{data['fitting_net']['type']}'." + ) + + descriptor_exclude_types = [ + list(pair) for pair in (data["descriptor"].get("exclude_types") or []) + ] + pair_exclude_types = [list(pair) for pair in (data.get("pair_exclude_types") or [])] + if pair_exclude_types: + if descriptor_exclude_types and descriptor_exclude_types != pair_exclude_types: + raise ValueError( + "DPA4/SeZM pair_exclude_types and descriptor.exclude_types must " + "match when both are provided." + ) + else: + pair_exclude_types = descriptor_exclude_types + data["pair_exclude_types"] = pair_exclude_types + data["descriptor"]["exclude_types"] = deepcopy(pair_exclude_types) + return get_standard_model(data) + + def get_model(data: dict) -> BaseModel: """Get a model from a dictionary. @@ -125,5 +176,7 @@ def get_model(data: dict) -> BaseModel: return get_zbl_model(data) else: return get_standard_model(data) + elif model_type in ("SeZM", "sezm", "DPA4", "dpa4"): + return get_sezm_model(data) else: return BaseModel.get_class_by_type(model_type).get_model(data) diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index c19267250e..450e2795e1 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -86,6 +86,7 @@ preprocess_shared_params, ) from deepmd.jax.utils.serialization import ( + _drop_zero_size_array_leaves, serialize_from_file, ) from deepmd.utils.argcheck import ( @@ -657,6 +658,7 @@ def loss_fn( model_dict = _evaluate_model_dict( model, extended_coord, extended_atype, nlist, mapping, fp, ap ) + model_dict = _match_label_shapes(model_dict, label_dict) loss, _ = loss_obj( learning_rate=lr, natoms=label_dict["type"].shape[1], @@ -680,6 +682,7 @@ def loss_fn_more_loss( model_dict = _evaluate_model_dict( model, extended_coord, extended_atype, nlist, mapping, fp, ap ) + model_dict = _match_label_shapes(model_dict, label_dict) _, more_loss = loss_obj( learning_rate=lr, natoms=label_dict["type"].shape[1], @@ -930,6 +933,7 @@ def _write_checkpoint(self, ckpt_path: Path, *, step: int) -> None: else: _, single_state = nnx.split(self.models[DEFAULT_TASK_KEY]) state = single_state.to_pure_dict() + state = _drop_zero_size_array_leaves(state) if ckpt_path.is_dir(): shutil.rmtree(ckpt_path) model_def_script_cpy = deepcopy(self.model_def_script) @@ -987,11 +991,32 @@ def _evaluate_model_dict( ) model_dict["atom_energy"] = model_dict["energy"] model_dict["energy"] = model_dict["energy_redu"] - model_dict["force"] = model_dict["energy_derv_r"].squeeze(-2) + force = model_dict["energy_derv_r"].squeeze(-2) + if force.ndim == 2 or (force.ndim == 3 and force.shape[-1] != 3): + force = jnp.reshape(force, (force.shape[0], -1, 3)) + model_dict["force"] = force model_dict["virial"] = model_dict["energy_derv_c_redu"].squeeze(-2) return model_dict +def _match_label_shapes( + model_dict: dict[str, jnp.ndarray], + label_dict: dict[str, jnp.ndarray], +) -> dict[str, jnp.ndarray]: + """Match equivalent flattened model outputs to label tensor shapes.""" + force_hat = model_dict.get("force") + force = label_dict.get("force") + if ( + force_hat is not None + and force is not None + and force_hat.shape != force.shape + and force_hat.size == force.size + ): + model_dict = dict(model_dict) + model_dict["force"] = jnp.reshape(force_hat, force.shape) + return model_dict + + def share_jax_model_params( models: dict[str, BaseModel], shared_links: dict[str, Any], diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 39354cc1fe..fdbd50a3ce 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -24,6 +24,22 @@ get_model, ) +_DROP_ZERO_SIZE_LEAF = object() + + +def _drop_zero_size_array_leaves(value: Any) -> Any: + """Remove array leaves Orbax cannot save while preserving tree containers.""" + if isinstance(value, dict): + filtered = {} + for key, item in value.items(): + new_item = _drop_zero_size_array_leaves(item) + if new_item is not _DROP_ZERO_SIZE_LEAF: + filtered[key] = new_item + return filtered + if getattr(value, "size", None) == 0: + return _DROP_ZERO_SIZE_LEAF + return value + def _convert_str_to_int_key(item: dict) -> None: """Convert Orbax-restored numeric index keys from strings back to ints.""" @@ -50,6 +66,40 @@ def _normalize_restored_state_keys( _convert_str_to_int_key(state) +_NO_ZERO_SIZE_LEAF = object() + + +def _zero_size_subtree(value: Any) -> Any: + if isinstance(value, dict): + restored = {} + for key, item in value.items(): + subtree = _zero_size_subtree(item) + if subtree is not _NO_ZERO_SIZE_LEAF: + restored[key] = subtree + return restored if restored else _NO_ZERO_SIZE_LEAF + if getattr(value, "size", None) == 0: + return value + return _NO_ZERO_SIZE_LEAF + + +def _restore_missing_zero_size_leaves(template: Any, restored: Any) -> Any: + """Reinsert zero-size leaves dropped before Orbax checkpoint saving.""" + if not isinstance(template, dict) or not isinstance(restored, dict): + return restored + restored = dict(restored) + for key, template_value in template.items(): + if key in restored: + restored[key] = _restore_missing_zero_size_leaves( + template_value, + restored[key], + ) + continue + subtree = _zero_size_subtree(template_value) + if subtree is not _NO_ZERO_SIZE_LEAF: + restored[key] = subtree + return restored + + def _state_sequence_to_numpy_list(state_value: Any) -> list[np.ndarray]: """Convert an Orbax-restored list/dict sequence to NumPy arrays.""" if isinstance(state_value, dict): @@ -219,6 +269,7 @@ def deserialize_to_file(model_file: str, data: dict) -> None: model = BaseModel.deserialize(data["model"]) _, state = nnx.split(model) state = state.to_pure_dict() + state = _drop_zero_size_array_leaves(state) with ocp.Checkpointer( ocp.CompositeCheckpointHandler("state", "model_def_script") ) as checkpointer: @@ -382,6 +433,10 @@ def restore_model(model_params: dict, model_state: dict) -> BaseModel: abstract_model = get_model(model_params) _restore_compression_slots_from_state(abstract_model, model_state) graphdef, abstract_state = nnx.split(abstract_model) + model_state = _restore_missing_zero_size_leaves( + abstract_state.to_pure_dict(), + model_state, + ) abstract_state.replace_by_pure_dict(model_state) return nnx.merge(graphdef, abstract_state) diff --git a/deepmd/pt/model/task/sezm_ener.py b/deepmd/pt/model/task/sezm_ener.py index 0932ec7086..f83fd339f3 100644 --- a/deepmd/pt/model/task/sezm_ener.py +++ b/deepmd/pt/model/task/sezm_ener.py @@ -243,11 +243,18 @@ def __init__( super().__init__() if neuron is None: neuron = [] - if isinstance(trainable, list): - trainable = all(trainable) self.in_dim = int(in_dim) self.out_dim = int(out_dim) self.neuron = [int(nn_dim) for nn_dim in neuron] + if isinstance(trainable, bool): + self.trainable = [trainable] * (len(self.neuron) + 1) + else: + self.trainable = [bool(flag) for flag in trainable] + if len(self.trainable) != len(self.neuron) + 1: + raise ValueError( + "trainable must contain one flag per hidden layer plus " + "one flag for the output layer" + ) self.activation_function = activation_function self.resnet_dt = bool(resnet_dt) self.precision = precision @@ -270,7 +277,7 @@ def __init__( activation_function=self.activation_function, precision=self.precision, seed=child_seed(seed, layer_idx), - trainable=trainable, + trainable=self.trainable[layer_idx], ) ) dim_in = hidden_dim @@ -285,7 +292,7 @@ def __init__( activation_function=self.activation_function, precision=self.precision, seed=child_seed(seed, len(self.neuron)), - trainable=trainable, + trainable=all(self.trainable), ) else: self.case_film = None @@ -300,12 +307,9 @@ def __init__( resnet=False, precision=self.precision, seed=child_seed(seed, len(self.neuron) + int(self.case_film_embd)), - trainable=trainable, + trainable=self.trainable[-1], ) - for param in self.parameters(): - param.requires_grad = trainable - def _apply_input_film( self, xx: torch.Tensor, @@ -395,6 +399,8 @@ def serialize(self) -> dict[str, Any]: "descriptor_dim": self.descriptor_dim, "dim_case_embd": self.dim_case_embd, "case_film_embd": self.case_film_embd, + # Keep the per-layer freeze policy stable across backend round trips. + "trainable": self.trainable.copy(), "@variables": {key: to_numpy_array(value) for key, value in state.items()}, } diff --git a/source/tests/consistent/descriptor/test_dpa4.py b/source/tests/consistent/descriptor/test_dpa4.py index 1cfbf2f9a1..58f7edadd3 100644 --- a/source/tests/consistent/descriptor/test_dpa4.py +++ b/source/tests/consistent/descriptor/test_dpa4.py @@ -20,6 +20,7 @@ from ..common import ( INSTALLED_ARRAY_API_STRICT, + INSTALLED_JAX, INSTALLED_PT, INSTALLED_PT_EXPT, CommonTest, @@ -37,6 +38,10 @@ from deepmd.pt_expt.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4PTExpt else: DescrptDPA4PTExpt = None +if INSTALLED_JAX: + from deepmd.jax.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4JAX +else: + DescrptDPA4JAX = None if INSTALLED_ARRAY_API_STRICT: from ...array_api_strict.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4Strict else: @@ -155,7 +160,7 @@ def skip_pt(self) -> bool: skip_dp = False skip_tf = True - skip_jax = True + skip_jax = not INSTALLED_JAX or DescrptDPA4JAX is None skip_pd = True skip_pt_expt = not INSTALLED_PT_EXPT skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT @@ -164,7 +169,7 @@ def skip_pt(self) -> bool: dp_class = DescrptDPA4DP pt_class = DescrptDPA4PT pt_expt_class = DescrptDPA4PTExpt - jax_class = None + jax_class = DescrptDPA4JAX pd_class = None array_api_strict_class = DescrptDPA4Strict args: ClassVar[list] = [ @@ -239,6 +244,16 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: mixed_types=True, ) + def eval_jax(self, jax_obj: Any) -> Any: + return self.eval_jax_descriptor( + jax_obj, + self.natoms, + self.coords, + self.atype, + self.box, + mixed_types=True, + ) + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: return self.eval_array_api_strict_descriptor( array_api_strict_obj, diff --git a/source/tests/consistent/fitting/test_dpa4_ener.py b/source/tests/consistent/fitting/test_dpa4_ener.py index 3d007a9959..3982656870 100644 --- a/source/tests/consistent/fitting/test_dpa4_ener.py +++ b/source/tests/consistent/fitting/test_dpa4_ener.py @@ -19,6 +19,7 @@ from ..common import ( INSTALLED_ARRAY_API_STRICT, + INSTALLED_JAX, INSTALLED_PT, INSTALLED_PT_EXPT, CommonTest, @@ -44,6 +45,13 @@ from deepmd.pt_expt.utils.env import DEVICE as PT_EXPT_DEVICE else: SeZMEnerFittingPTExpt = None +if INSTALLED_JAX: + from deepmd.jax.env import ( + jnp, + ) + from deepmd.jax.fitting.dpa4_ener import SeZMEnergyFittingNet as SeZMEnerFittingJAX +else: + SeZMEnerFittingJAX = None if INSTALLED_ARRAY_API_STRICT: import array_api_strict @@ -86,7 +94,7 @@ def skip_pt(self) -> bool: skip_dp = False skip_tf = True - skip_jax = True + skip_jax = not INSTALLED_JAX or SeZMEnerFittingJAX is None skip_pd = True skip_pt_expt = not INSTALLED_PT_EXPT skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT @@ -95,7 +103,7 @@ def skip_pt(self) -> bool: dp_class = SeZMEnerFittingDP pt_class = SeZMEnerFittingPT pt_expt_class = SeZMEnerFittingPTExpt - jax_class = None + jax_class = SeZMEnerFittingJAX pd_class = None array_api_strict_class = SeZMEnerFittingStrict args = fitting_sezm_ener() @@ -150,6 +158,14 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: .numpy() ) + def eval_jax(self, jax_obj: Any) -> Any: + return np.asarray( + jax_obj( + jnp.asarray(self.inputs), + jnp.asarray(self.atype.reshape(1, -1)), + )["energy"] + ) + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: return to_numpy_array( array_api_strict_obj( diff --git a/source/tests/jax/test_dpa4.py b/source/tests/jax/test_dpa4.py new file mode 100644 index 0000000000..d25ec00f1c --- /dev/null +++ b/source/tests/jax/test_dpa4.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Focused tests for JAX DPA4 trainable-state conversion.""" + +from deepmd.jax.descriptor.dpa4 import ( + DescrptDPA4, + _iter_object_tree, +) +from deepmd.jax.env import ( + nnx, +) +from deepmd.jax.fitting.dpa4_ener import ( + SeZMEnergyFittingNet, +) +from deepmd.jax.utils.network import ( + ArrayAPIParam, +) + + +def _make_trainable_descriptor() -> DescrptDPA4: + """Build a small descriptor that enables the optional trainable leaves.""" + return DescrptDPA4( + ntypes=2, + sel=4, + rcut=4.0, + channels=4, + n_radial=4, + lmax=1, + mmax=1, + n_blocks=1, + grid_branch=0, + layer_scale=True, + message_node_so3=True, + random_gamma=False, + precision="float64", + trainable=True, + seed=20260711, + ) + + +def test_optional_dpa4_weights_are_jax_parameters() -> None: + """Optional cross-grid and FFN LayerScale weights must receive gradients.""" + descriptor = _make_trainable_descriptor() + modules = list(_iter_object_tree(descriptor)) + + frame_modules = [ + module + for module in modules + if type(module).__name__ in {"FrameContract", "FrameExpand"} + ] + assert {type(module).__name__ for module in frame_modules} == { + "FrameContract", + "FrameExpand", + } + assert all(isinstance(module.weight, ArrayAPIParam) for module in frame_modules) + + interaction_blocks = [ + module for module in modules if type(module).__name__ == "SeZMInteractionBlock" + ] + assert interaction_blocks + assert all( + isinstance(scale, ArrayAPIParam) + for block in interaction_blocks + for scale in block.adam_ffn_layer_scales + ) + + +def test_frozen_descriptor_has_no_optimizer_visible_parameters() -> None: + """The root freeze flag must demote every descendant parameter.""" + descriptor = DescrptDPA4( + ntypes=2, + sel=4, + rcut=4.0, + channels=4, + n_radial=4, + lmax=1, + mmax=1, + n_blocks=1, + grid_branch=0, + random_gamma=False, + precision="float64", + trainable=False, + seed=20260712, + ) + + assert len(nnx.to_flat_state(nnx.state(descriptor, nnx.Param))) == 0 + + +def test_frozen_fitting_stays_frozen_after_conversion_round_trip() -> None: + """Serialized GLU layers retain the all-false optimizer policy.""" + fitting = SeZMEnergyFittingNet( + ntypes=2, + dim_descrpt=4, + neuron=[4], + trainable=False, + precision="float64", + mixed_types=True, + seed=20260712, + ) + restored = SeZMEnergyFittingNet.deserialize(fitting.serialize()) + + assert restored.nets[0].trainable == [False, False] + assert len(nnx.to_flat_state(nnx.state(restored, nnx.Param))) == 0 diff --git a/source/tests/jax/test_dpa4_conversion.py b/source/tests/jax/test_dpa4_conversion.py new file mode 100644 index 0000000000..703677f1be --- /dev/null +++ b/source/tests/jax/test_dpa4_conversion.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""End-to-end PT checkpoint conversion coverage for JAX DPA4.""" + +from copy import ( + deepcopy, +) + +import torch + +from deepmd.jax.utils.serialization import ( + deserialize_to_file as deserialize_to_jax_file, +) +from deepmd.jax.utils.serialization import ( + serialize_from_file as serialize_from_jax_file, +) +from deepmd.pt.model.model import get_model as get_pt_model +from deepmd.pt.train.wrapper import ( + ModelWrapper, +) +from deepmd.pt.utils.serialization import serialize_from_file as serialize_from_pt_file +from deepmd.utils.argcheck import ( + model_args, +) + + +def _small_dpa4_config() -> dict: + """Return a small real PT DPA4 config with zero-size state leaves.""" + return model_args().normalize_value( + { + "type": "dpa4", + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4", + "sel": 4, + "rcut": 4.0, + "channels": 4, + "n_radial": 4, + "lmax": 1, + "mmax": 1, + "n_blocks": 1, + "random_gamma": False, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [4], + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_.*", + ) + + +def test_pt_dpa4_checkpoint_converts_to_real_jax_checkpoint(tmp_path) -> None: + """The public PT schema saves and restores through Orbax without loss.""" + model_params = _small_dpa4_config() + pt_model = get_pt_model(deepcopy(model_params)).to(torch.float64) + wrapper = ModelWrapper(pt_model, model_params=deepcopy(model_params)) + pt_path = tmp_path / "dpa4.pt" + jax_path = tmp_path / "dpa4.jax" + torch.save({"model": wrapper.state_dict()}, pt_path) + + data = serialize_from_pt_file(str(pt_path)) + assert data["model"]["type"] == "SeZM" + + deserialize_to_jax_file(str(jax_path), data) + restored = serialize_from_jax_file(str(jax_path)) + + assert restored["model"]["type"] == "standard" + assert restored["model"]["descriptor"]["type"] == "SeZM" diff --git a/source/tests/jax/test_model_factory.py b/source/tests/jax/test_model_factory.py index 75ffc519a1..12a137f9ba 100644 --- a/source/tests/jax/test_model_factory.py +++ b/source/tests/jax/test_model_factory.py @@ -10,6 +10,9 @@ """ import unittest +from unittest.mock import ( + patch, +) from deepmd.jax.model.ener_model import ( EnergyModel, @@ -17,6 +20,9 @@ from deepmd.jax.model.model import ( get_model, ) +from deepmd.utils.argcheck import ( + model_args, +) def _base_config() -> dict: @@ -40,6 +46,16 @@ def _base_config() -> dict: } +def _base_sezm_config() -> dict: + """Return the smallest config needed to exercise DPA4 factory routing.""" + return { + "type": "dpa4", + "type_map": ["O", "H"], + "descriptor": {"type": "dpa4"}, + "fitting_net": {"type": "dpa4_ener"}, + } + + class TestJAXModelFactoryFittingDefault(unittest.TestCase): def test_fitting_net_without_type_defaults_to_ener(self) -> None: # fitting_net present but no "type": must default to energy. @@ -62,5 +78,89 @@ def test_explicit_fitting_type_preserved(self) -> None: self.assertIsInstance(model, EnergyModel) +class TestJAXSeZMModelFactory(unittest.TestCase): + @patch("deepmd.jax.model.model.get_standard_model", side_effect=lambda data: data) + def test_null_blocks_receive_dpa4_defaults(self, _get_standard_model) -> None: + data = _base_sezm_config() + data["descriptor"] = None + data["fitting_net"] = None + + normalized = get_model(data) + + self.assertEqual(normalized["descriptor"]["type"], "dpa4") + self.assertEqual(normalized["fitting_net"]["type"], "dpa4_ener") + + def test_rejects_unsupported_features(self) -> None: + cases = ( + ("spin", {}), + ("bridging_method", "linear"), + ("lora", {}), + ("use_compile", True), + ("preset_out_bias", [0.0]), + ) + for key, value in cases: + with self.subTest(key=key): + data = _base_sezm_config() + data[key] = value + with self.assertRaises(NotImplementedError): + get_model(data) + + data = _base_sezm_config() + data["descriptor"]["add_chg_spin_ebd"] = True + with self.assertRaises(NotImplementedError): + get_model(data) + + def test_rejects_incompatible_descriptor_and_fitting_types(self) -> None: + data = _base_sezm_config() + data["descriptor"]["type"] = "se_e2_a" + with self.assertRaises(ValueError): + get_model(data) + + data = _base_sezm_config() + data["fitting_net"]["type"] = "ener" + with self.assertRaises(ValueError): + get_model(data) + + def test_rejects_mismatched_exclude_types(self) -> None: + data = _base_sezm_config() + data["descriptor"]["exclude_types"] = [[0, 1]] + data["pair_exclude_types"] = [[1, 1]] + + with self.assertRaises(ValueError): + get_model(data) + + @patch("deepmd.jax.model.model.get_standard_model", side_effect=lambda data: data) + def test_descriptor_exclude_types_feed_standard_model( + self, + _get_standard_model, + ) -> None: + data = _base_sezm_config() + data["descriptor"] = { + "type": "SeZM", + "exclude_types": [[0, 1]], + } + data["fitting_net"]["type"] = "sezm_ener" + + normalized = get_model(data) + + self.assertEqual(normalized["pair_exclude_types"], [[0, 1]]) + self.assertEqual(normalized["descriptor"]["exclude_types"], [[0, 1]]) + + @patch("deepmd.jax.model.model.get_standard_model", side_effect=lambda data: data) + def test_normalized_descriptor_exclusions_override_empty_default( + self, + _get_standard_model, + ) -> None: + """Argcheck's empty model-level default is not an explicit mismatch.""" + data = _base_sezm_config() + data["descriptor"]["exclude_types"] = [[0, 1]] + data = model_args().normalize_value(data, trim_pattern="_.*") + + normalized = get_model(data) + + self.assertEqual(normalized["pair_exclude_types"], [[0, 1]]) + self.assertEqual(normalized["descriptor"]["exclude_types"], [[0, 1]]) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index 72e8a47ec7..74e81d07f5 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -49,6 +49,9 @@ from deepmd.jax.train.trainer import ( DPTrainer, _copy_matching_state_tree, + _drop_zero_size_array_leaves, + _evaluate_model_dict, + _match_label_shapes, _merge_descriptor_stats, _merge_fitting_param_stats, _scale_by_global_learning_rate, @@ -58,6 +61,7 @@ ) from deepmd.jax.utils.serialization import ( _normalize_restored_state_keys, + _restore_missing_zero_size_leaves, ) from deepmd.utils.compat import ( convert_optimizer_v31_to_v32, @@ -929,3 +933,96 @@ def test_jax_multitask_state_key_normalization_preserves_numeric_task_names() -> assert 1 not in state["models"] assert 0 in state["models"]["1"]["layers"] assert 0 in state["models"]["task"]["layers"] + + +def test_jax_zero_size_checkpoint_leaves_round_trip() -> None: + """Checkpoint filtering and restore must preserve every zero-size path.""" + template = { + "model": { + "empty": jnp.zeros((0, 3)), + "nested": { + "empty": jnp.zeros((2, 0)), + "weight": jnp.ones((2,)), + }, + } + } + + filtered = _drop_zero_size_array_leaves(template) + restored = _restore_missing_zero_size_leaves(template, filtered) + + assert "empty" not in filtered["model"] + assert "empty" not in filtered["model"]["nested"] + np.testing.assert_array_equal( + restored["model"]["empty"], template["model"]["empty"] + ) + np.testing.assert_array_equal( + restored["model"]["nested"]["empty"], + template["model"]["nested"]["empty"], + ) + np.testing.assert_array_equal( + restored["model"]["nested"]["weight"], + template["model"]["nested"]["weight"], + ) + + +def test_jax_match_label_shapes_reshapes_only_equivalent_force_layouts() -> None: + """Flattened force tensors reshape, while an existing layout is untouched.""" + force = jnp.arange(6).reshape(1, 2, 3) + model_dict = {"force": force} + + reshaped = _match_label_shapes(model_dict, {"force": jnp.zeros((1, 6))}) + unchanged = _match_label_shapes(model_dict, {"force": jnp.zeros((1, 2, 3))}) + + assert reshaped is not model_dict + assert reshaped["force"].shape == (1, 6) + assert unchanged is model_dict + + +def test_jax_evaluate_model_dict_normalizes_flattened_force_only() -> None: + """Model evaluation preserves canonical force tensors and expands flat ones.""" + + class FakeModel: + def __init__(self, force_derivative: jnp.ndarray) -> None: + self.force_derivative = force_derivative + + def call_common_lower(self, *args, **kwargs): + del args, kwargs + return { + "energy": jnp.zeros((1, 2, 1)), + "energy_redu": jnp.zeros((1, 1)), + "energy_derv_r": self.force_derivative, + "energy_derv_c_redu": jnp.zeros((1, 1, 9)), + } + + def model_output_def(self): + return {} + + def passthrough(model_dict, *args, **kwargs): + del args, kwargs + return dict(model_dict) + + with patch( + "deepmd.jax.train.trainer.communicate_extended_output", + side_effect=passthrough, + ): + flattened = _evaluate_model_dict( + FakeModel(jnp.arange(6).reshape(1, 1, 6)), + jnp.zeros((1, 6)), + jnp.zeros((1, 2), dtype=jnp.int32), + jnp.zeros((1, 2, 1), dtype=jnp.int32), + None, + None, + None, + ) + canonical = _evaluate_model_dict( + FakeModel(jnp.arange(6).reshape(1, 2, 1, 3)), + jnp.zeros((1, 6)), + jnp.zeros((1, 2), dtype=jnp.int32), + jnp.zeros((1, 2, 1), dtype=jnp.int32), + None, + None, + None, + ) + + assert flattened["force"].shape == (1, 2, 3) + assert canonical["force"].shape == (1, 2, 3)