From 40585f656a8887e5dddeff43c39cee3f4500e57e Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 29 Jul 2026 23:12:40 +0800 Subject: [PATCH] docs(jax): document training and hessian freeze Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- doc/backend.md | 7 +++++-- doc/freeze/freeze.md | 23 +++++++++++++++++++++++ doc/model/overall.md | 4 +++- doc/model/train-energy-hessian.md | 10 +++++++--- doc/train/training.md | 11 +++++++++++ 5 files changed, 49 insertions(+), 6 deletions(-) diff --git a/doc/backend.md b/doc/backend.md index 3e5d634ac7..450dee5d4f 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -25,15 +25,18 @@ While `.pth` and `.pt` are the same in the PyTorch package, they have different ### JAX {{ jax_icon }} -- Model filename extension: `.xlo`, `.savedmodel` +- Model filename extension: `.hlo`, `.savedmodel` - Checkpoint filename extension: `.jax` [JAX](https://jax.readthedocs.io/) 0.4.33 or above is required. -Both `.xlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions. +Both `.hlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions. `.savedmodel` is the TensorFlow [SavedModel format](https://www.tensorflow.org/guide/saved_model) generated by [JAX2TF](https://www.tensorflow.org/guide/jax2tf), which needs the installation of TensorFlow. Only the `.savedmodel` format supports C++ inference, which needs the TensorFlow C++ interface. The model is device-specific, so that the model generated on the GPU device cannot be run on the CPUs. +JAX supports training with `dp --jax train`; training checkpoints use the +`.jax` extension and can be frozen as `.hlo`, `.jax`, or `.savedmodel` models. + ### Paddle {{ paddle_icon }} - Model filename extensions: `.json` and `.pdiparams` diff --git a/doc/freeze/freeze.md b/doc/freeze/freeze.md index 3aef7e22a6..f0e4945a06 100644 --- a/doc/freeze/freeze.md +++ b/doc/freeze/freeze.md @@ -51,4 +51,27 @@ $ dp --pd freeze -o model_branch1 --head CHOSEN_BRANCH The output model is called `model_branch1.json`, which is the specifically frozen model with the `CHOSEN_BRANCH` head. ::: +:::{tab-item} JAX {{ jax_icon }} + +```bash +$ dp --jax freeze -c model.ckpt.jax -o model.hlo +``` + +The JAX backend can write a StableHLO `.hlo` model, a lossless `.jax` model, or +a JAX2TF `.savedmodel` model. The `.savedmodel` format requires TensorFlow and +is the JAX format that supports the C++ inference interface. +::: + :::: + +## Freeze a JAX model with Hessian output {{ jax_icon }} + +Use `--hessian` to add coordinate-Hessian output to a frozen JAX energy model: + +```bash +$ dp --jax freeze -c model.ckpt.jax -o model-hessian.hlo --hessian +``` + +The option applies to JAX `.hlo`, `.jax`, and `.savedmodel` outputs. A model +whose serialized definition already enables Hessian mode retains that mode even +when `--hessian` is omitted. diff --git a/doc/model/overall.md b/doc/model/overall.md index 37470f44cb..b3ffc0e629 100644 --- a/doc/model/overall.md +++ b/doc/model/overall.md @@ -58,7 +58,9 @@ The fitting of the following physical properties is supported 1. [`ener`](train-energy.md): Fit the energy of the system. The force (derivative with atom positions), the virial (derivative with the box tensor) and the hessian (second-order derivative with atom positions) can also be trained. > [!WARNING] -> Due to the restrictions of torch jit script, the models trained with hessian are not jitable so that the frozen models cannot output hessians. +> The PyTorch TorchScript freeze route cannot output Hessians. The JAX backend +> can retain Hessian output in a frozen model with +> `dp --jax freeze --hessian`; see [Freeze a model](../freeze/freeze.md). 2. [`dipole`](train-fitting-tensor.md): The dipole moment. 1. [`polar`](train-fitting-tensor.md): The polarizability. diff --git a/doc/model/train-energy-hessian.md b/doc/model/train-energy-hessian.md index dc1ca38fe8..d99d6d0c1f 100644 --- a/doc/model/train-energy-hessian.md +++ b/doc/model/train-energy-hessian.md @@ -91,9 +91,12 @@ The detailed loss can be found in `lcurve.out`: ## Test the Model > [!WARNING] -> A model trained with Hessian cannot be frozen. If freezing is enforced, the model will be treated as a standard energy model, and the frozen one will no longer be able to output Hessian predictions. +> The PyTorch TorchScript freeze route does not preserve Hessian output. A +> PyTorch model frozen with `dp --pt freeze` is treated as a standard energy +> model. The JAX backend can preserve Hessian output in a frozen model with +> `dp --jax freeze --hessian`. -If one do freeze and test a Hessian model using the commands: +If one freezes and tests a Hessian model through the PyTorch route: ::::{tab-set} @@ -116,7 +119,8 @@ ${output_prefix}.e.out ${output_prefix}.e_peratom.out ${output_prefix}.f.out ${output_prefix}.v.out ${output_prefix}.v_peratom.out ``` -If one intends to use the trained model for Hessian predictions, then he/she is supposed to test the model directly without performing a freezing operation: +For PyTorch Hessian predictions, test the training checkpoint directly without +freezing: ::::{tab-set} diff --git a/doc/train/training.md b/doc/train/training.md index df1341b4ee..ab23905aca 100644 --- a/doc/train/training.md +++ b/doc/train/training.md @@ -39,6 +39,17 @@ $ CINN=1 CINN_ALLOW_DYNAMIC_SHAPE=0 dp --pd train input.json ``` ::: +:::{tab-item} JAX {{ jax_icon }} + +```bash +$ dp --jax train input.json +``` + +JAX training, introduced after DeePMD-kit v3.1.3, writes checkpoints with the +`.jax` extension. It supports fine-tuning and multi-task configurations, but +does not currently support `--init-frz-model`. +::: + :::: where `input.json` is the name of the input script.