From 91e4a09538f0a94c9b4cf9bcc0e3395b44994f7f Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 30 Jun 2026 02:16:21 +0800 Subject: [PATCH 1/3] feat(jax): freeze models with hessian output --- deepmd/jax/entrypoints/freeze.py | 5 +++- deepmd/jax/infer/deep_eval.py | 5 ++++ deepmd/jax/jax2tf/serialization.py | 10 +++++-- deepmd/jax/model/hlo.py | 21 ++++++++++++- deepmd/jax/utils/serialization.py | 36 +++++++++++++++++++--- deepmd/main.py | 6 ++++ source/tests/jax/test_training.py | 48 ++++++++++++++++++++++++++++-- 7 files changed, 120 insertions(+), 11 deletions(-) diff --git a/deepmd/jax/entrypoints/freeze.py b/deepmd/jax/entrypoints/freeze.py index fbc126ffc7..03536d4031 100644 --- a/deepmd/jax/entrypoints/freeze.py +++ b/deepmd/jax/entrypoints/freeze.py @@ -18,6 +18,7 @@ def freeze( *, checkpoint_folder: str, output: str, + hessian: bool = False, **kwargs: object, ) -> None: """Freeze a JAX checkpoint into a serialized model file. @@ -30,6 +31,8 @@ def freeze( output : str Output model filename or prefix. The JAX model suffix is added when the filename has no supported backend suffix. + hessian : bool, default=False + Whether to include the Hessian in the frozen model outputs. **kwargs Other CLI arguments accepted for backend entry-point compatibility. """ @@ -46,4 +49,4 @@ def freeze( strict_prefer=True, ) data = serialize_from_file(checkpoint_folder) - deserialize_to_file(output, data) + deserialize_to_file(output, data, hessian=hessian) diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 09b5783ba9..efd2aac574 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -306,6 +306,7 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: OutputVariableCategory.REDU, OutputVariableCategory.DERV_R, OutputVariableCategory.DERV_C_REDU, + OutputVariableCategory.DERV_R_DERV_R, ) ] @@ -458,6 +459,10 @@ def get_model_def_script(self) -> dict: """Get model definition script.""" return json.loads(self.dp.get_model_def_script()) + def get_has_hessian(self) -> bool: + """Check if the model has Hessian output.""" + return self.get_model_def_script().get("hessian_mode", False) + def get_model(self) -> Any: """Get the JAX model as BaseModel. diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index dd496dedf0..9dbc3bd390 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -29,16 +29,22 @@ BaseModel, ) from deepmd.jax.utils.serialization import ( + _prepare_hessian_model_def_script, _set_model_min_nbor_dist_from_data, ) -def deserialize_to_file(model_file: str, data: dict) -> None: +def deserialize_to_file(model_file: str, data: dict, hessian: bool = False) -> None: """Deserialize the dictionary to a JAX/jax2tf SavedModel.""" if model_file.endswith(".savedmodel"): model = BaseModel.deserialize(data["model"]) _set_model_min_nbor_dist_from_data(model, data) - model_def_script = data["model_def_script"] + model_def_script, hessian = _prepare_hessian_model_def_script( + data["model_def_script"], + hessian, + ) + if hessian: + model.enable_hessian() call_lower = model.call_common_lower tf_model = tf.Module() diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index 8c1e85c59c..b0d81c31ee 100644 --- a/deepmd/jax/model/hlo.py +++ b/deepmd/jax/model/hlo.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import json from typing import ( Any, ) @@ -30,6 +31,14 @@ r_differentiable=True, c_differentiable=True, ), + "energy_hessian": OutputVariableDef( + "energy", + shape=[1], + reducible=True, + r_differentiable=True, + c_differentiable=True, + r_hessian=True, + ), "mask": OutputVariableDef( "mask", shape=[1], @@ -171,7 +180,17 @@ def call( def model_output_def(self) -> ModelOutputDef: return ModelOutputDef( - FittingOutputDef([OUTPUT_DEFS[tt] for tt in self.model_output_type()]) + FittingOutputDef( + [ + OUTPUT_DEFS[ + f"{tt}_hessian" + if tt == "energy" + and json.loads(self.model_def_script).get("hessian_mode", False) + else tt + ] + for tt in self.model_output_type() + ] + ) ) def call_lower( diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 59240b41ab..bf334317a6 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -185,7 +185,19 @@ def _check_compressed_hlo_exportable(data: dict) -> None: ) -def deserialize_to_file(model_file: str, data: dict) -> None: +def _prepare_hessian_model_def_script( + model_def_script: dict, + hessian: bool, +) -> tuple[dict, bool]: + """Return a copied model definition and whether Hessian should be enabled.""" + model_def_script = model_def_script.copy() + hessian = hessian or model_def_script.get("hessian_mode", False) + if hessian: + model_def_script["hessian_mode"] = True + return model_def_script, hessian + + +def deserialize_to_file(model_file: str, data: dict, hessian: bool = False) -> None: """Deserialize the dictionary to a model file. Parameters @@ -194,9 +206,14 @@ def deserialize_to_file(model_file: str, data: dict) -> None: The model file to be saved. data : dict The dictionary to be deserialized. + hessian : bool, default=False + Whether to include the Hessian in the model outputs. """ if model_file.endswith(".jax"): - model_def_script = data["model_def_script"].copy() + model_def_script, hessian = _prepare_hessian_model_def_script( + data["model_def_script"], + hessian, + ) min_nbor_dist = _to_optional_float(data.get("min_nbor_dist")) if min_nbor_dist is None: min_nbor_dist = _to_optional_float( @@ -209,6 +226,9 @@ def deserialize_to_file(model_file: str, data: dict) -> None: model_key: BaseModel.deserialize(data["model"]["model_dict"][model_key]) for model_key in model_def_script["model_dict"] } + if hessian: + for model in models.values(): + model.enable_hessian() state = { "models": { model_key: nnx.split(model)[1].to_pure_dict() @@ -217,6 +237,8 @@ def deserialize_to_file(model_file: str, data: dict) -> None: } else: model = BaseModel.deserialize(data["model"]) + if hessian: + model.enable_hessian() _, state = nnx.split(model) state = state.to_pure_dict() with ocp.Checkpointer( @@ -233,7 +255,12 @@ def deserialize_to_file(model_file: str, data: dict) -> None: _check_compressed_hlo_exportable(data) model = BaseModel.deserialize(data["model"]) _set_model_min_nbor_dist_from_data(model, data) - model_def_script = data["model_def_script"] + model_def_script, hessian = _prepare_hessian_model_def_script( + data["model_def_script"], + hessian, + ) + if hessian: + model.enable_hessian() call_lower = model.call_common_lower nf, nloc, nghost = jax_export.symbolic_shape("nf, nloc, nghost") @@ -298,6 +325,7 @@ def call_lower_with_fixed_do_atomic_virial( serialized_atomic_virial_no_ghost = exported_atomic_virial_no_ghost.serialize() data = data.copy() + data["model_def_script"] = model_def_script data.setdefault("@variables", {}) data["@variables"]["stablehlo"] = np.void(serialized) data["@variables"]["stablehlo_atomic_virial"] = np.void( @@ -331,7 +359,7 @@ def call_lower_with_fixed_do_atomic_virial( deserialize_to_file as deserialize_to_savedmodel, ) - return deserialize_to_savedmodel(model_file, data) + return deserialize_to_savedmodel(model_file, data, hessian=hessian) else: raise ValueError("Unsupported file extension") diff --git a/deepmd/main.py b/deepmd/main.py index 42a03bc3f0..03d1a2ec91 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -350,6 +350,12 @@ def main_parser() -> argparse.ArgumentParser: type=str, help="(Supported backend: PyTorch) Task head (alias: model branch) to freeze if in multi-task mode.", ) + parser_frz.add_argument( + "--hessian", + action="store_true", + default=False, + help="(Supported backend: JAX) Add the Hessian to the frozen model output.", + ) parser_frz.add_argument( "--lower-kind", default="nlist", diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index 25d3ccdc49..63811b5b02 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -24,6 +24,9 @@ import numpy as np import optax +from deepmd.dpmodel.output_def import ( + OutputVariableCategory, +) from deepmd.dpmodel.train import ( TrainEntrypointOptions, ) @@ -40,6 +43,12 @@ from deepmd.jax.env import ( jnp, ) +from deepmd.jax.infer.deep_eval import ( + DeepEval, +) +from deepmd.jax.model.hlo import ( + HLO, +) from deepmd.jax.train.trainer import ( DPTrainer, _copy_matching_state_tree, @@ -527,17 +536,19 @@ def test_update_sel_supports_multitask(self, get_nbor_stat, get_data) -> None: def test_freeze_entrypoint_uses_checkpoint_pointer( self, serialize_from_file, deserialize_to_file ) -> None: - """Freeze resolves the stable checkpoint pointer without Hessian options.""" + """Freeze resolves the stable checkpoint pointer and forwards Hessian.""" checkpoint_dir = self.work_dir / "ckpt" checkpoint_dir.mkdir() (checkpoint_dir / "checkpoint").write_text("model-1.jax") serialize_from_file.return_value = {"model": {}, "model_def_script": {}} - freeze(checkpoint_folder=str(checkpoint_dir), output="frozen_model") + freeze( + checkpoint_folder=str(checkpoint_dir), output="frozen_model", hessian=True + ) serialize_from_file.assert_called_once_with(str(checkpoint_dir / "model-1.jax")) deserialize_to_file.assert_called_once_with( - "frozen_model.hlo", serialize_from_file.return_value + "frozen_model.hlo", serialize_from_file.return_value, hessian=True ) @patch("deepmd.jax.entrypoints.main.freeze") @@ -549,12 +560,43 @@ def test_main_dispatches_freeze(self, freeze_entrypoint) -> None: log_path=None, checkpoint_folder=".", output="frozen_model", + hessian=False, ) main(args) freeze_entrypoint.assert_called_once() + def test_hlo_hessian_mode_updates_output_def(self) -> None: + """HLO output definition should expose Hessian when requested.""" + hlo = object.__new__(HLO) + hlo._model_output_type = ["energy"] + hlo.model_def_script = json.dumps({"hessian_mode": True}) + + output_def = hlo.model_output_def() + + self.assertTrue(output_def["energy"].r_hessian) + self.assertIn("energy_derv_r_derv_r", output_def.keys()) + + def test_deep_eval_requests_hessian_for_hessian_model(self) -> None: + """Non-atomic JAX evaluation should request Hessian outputs.""" + hlo = object.__new__(HLO) + hlo._model_output_type = ["energy"] + hlo.model_def_script = json.dumps({"hessian_mode": True}) + deep_eval = object.__new__(DeepEval) + deep_eval.output_def = hlo.model_output_def() + deep_eval.dp = SimpleNamespace( + get_model_def_script=lambda: json.dumps({"hessian_mode": True}) + ) + + request_defs = deep_eval._get_request_defs(atomic=False) + + self.assertTrue(deep_eval.get_has_hessian()) + self.assertIn( + OutputVariableCategory.DERV_R_DERV_R, + {odef.category for odef in request_defs}, + ) + def test_jax_finetune_state_copy_preserves_random_fitting_target_leaves() -> None: """Random fitting should copy descriptor leaves only.""" From 7d3223b1297c8ca8185bd9ea29b8f58b587b7e0d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 7 Jul 2026 23:29:00 +0800 Subject: [PATCH 2/3] fix(jax): avoid mixed returns in serialization --- deepmd/jax/utils/serialization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index bf334317a6..52b8c210c7 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -359,7 +359,7 @@ def call_lower_with_fixed_do_atomic_virial( deserialize_to_file as deserialize_to_savedmodel, ) - return deserialize_to_savedmodel(model_file, data, hessian=hessian) + deserialize_to_savedmodel(model_file, data, hessian=hessian) else: raise ValueError("Unsupported file extension") From 30735c11fbe714aca785cc8c3689c186b4c5c3c0 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 8 Jul 2026 22:26:22 +0800 Subject: [PATCH 3/3] test(jax): assert hessian freeze dispatch --- source/tests/jax/test_training.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index 63811b5b02..791b203d69 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -566,6 +566,8 @@ def test_main_dispatches_freeze(self, freeze_entrypoint) -> None: main(args) freeze_entrypoint.assert_called_once() + self.assertIn("hessian", freeze_entrypoint.call_args.kwargs) + self.assertFalse(freeze_entrypoint.call_args.kwargs["hessian"]) def test_hlo_hessian_mode_updates_output_def(self) -> None: """HLO output definition should expose Hessian when requested."""