From 7a94be5b820a2d09de7fcda896cd94a84cff108b Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 29 Jul 2026 23:10:18 +0800 Subject: [PATCH 01/10] docs(backend): document tf2 and exportable backends Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- doc/backend.md | 46 +++++++++++++++++++++++++++++++++++++++--- doc/freeze/compress.md | 22 +++++++++++++++++++- doc/freeze/freeze.md | 26 ++++++++++++++++++++++++ doc/train/training.md | 14 +++++++++++++ 4 files changed, 104 insertions(+), 4 deletions(-) diff --git a/doc/backend.md b/doc/backend.md index 3e5d634ac7..c71c21b53a 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -2,10 +2,13 @@ ## Supported backends -DeePMD-kit supports multiple backends: TensorFlow and PyTorch. +DeePMD-kit supports seven backends: TensorFlow, TensorFlow 2, PyTorch, +PyTorch-Exportable, JAX, Paddle, and the NumPy-based DP reference backend. To use DeePMD-kit, you must install at least one backend. Each backend does not support all features. -In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, and Paddle {{ paddle_icon }} icons are used to mark whether a backend supports a feature. +In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch +{{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, and DP +{{ dpmodel_icon }} icons are used to mark whether a backend supports a feature. ### TensorFlow {{ tensorflow_icon }} @@ -15,6 +18,20 @@ In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon [TensorFlow](https://tensorflow.org) 2.8 is the first version to support Python 3.10. DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`tf.compat.v1`) in the graph mode. +### TensorFlow 2 {{ tensorflow_icon }} + +- Model filename extension: `.savedmodeltf` +- Checkpoint directory extension: `.tf2` + +The TensorFlow 2 backend uses the TensorFlow v2 eager API. Select it with +`dp --tf2` (alias `dp --tensorflow2`). It supports training, including +multi-task training and fine-tuning, freezing, compression, and testing. +Training stores checkpoints in a directory named after the `save_ckpt` prefix +with `.tf2` appended, such as `model.ckpt.tf2`. + +Setting [`DP_JIT`](env.md#envvar-DP_JIT) enables optional `tf.function` JIT +compilation; depending on the workload, this may improve or reduce performance. + ### PyTorch {{ pytorch_icon }} - Model filename extension: `.pth` @@ -23,6 +40,25 @@ DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`t [PyTorch](https://pytorch.org/) 2.1 or above is required. While `.pth` and `.pt` are the same in the PyTorch package, they have different meanings in the DeePMD-kit to distinguish the model and the checkpoint. +### PyTorch-Exportable {{ pytorch_icon }} + +- Model filename extensions: `.pte`, `.pt2` +- Checkpoint filename extension: `.pt` + +Select this backend with `dp --pt-expt` (alias +`dp --pytorch-exportable`). It uses PyTorch with the backend-independent model +implementation and supports training, including multi-task training and +fine-tuning, freezing, compression, change-bias, and testing. Training can read +LMDB datasets, and Python inference can use the optional vesin neighbor-list +implementation. + +Freezing exports a `torch.export` model. The dense neighbor-list lower form +normally uses `.pte`, while the graph lower form uses an AOTInductor `.pt2` +package. Use `--lower-kind graph` to request graph-native export for an eligible +model; graph-capable DPA models may select that form automatically. The `.pt` +checkpoint format uses DP-model parameter names ending in `.w` and `.b`, which +allows DeePMD-kit to distinguish it from a regular PyTorch checkpoint. + ### JAX {{ jax_icon }} - Model filename extension: `.xlo`, `.savedmodel` @@ -64,12 +100,16 @@ NumPy 1.21 or above is required. ### Training -When training and freezing a model, you can use `dp --tf`, `dp --pt` or `dp --pd` in the command line to switch the backend. +When training and freezing a model, use `dp --tf`, `dp --tf2`, `dp --pt`, +`dp --pt-expt`, `dp --jax`, or `dp --pd` in the command line to switch the +backend. ### Inference When doing inference, DeePMD-kit detects the backend from the model filename. For example, when the model filename ends with `.pb` (the ProtoBuf file), DeePMD-kit will consider it using the TensorFlow backend. +The same detection covers TensorFlow 2 `.savedmodeltf` models and +PyTorch-Exportable `.pte` and `.pt2` models. ## Convert model files between backends diff --git a/doc/freeze/compress.md b/doc/freeze/compress.md index 1cf643a627..17ff2af1fe 100644 --- a/doc/freeze/compress.md +++ b/doc/freeze/compress.md @@ -1,7 +1,8 @@ # Compress a model {{ tensorflow_icon }} {{ pytorch_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 {{ tensorflow_icon }}, +> PyTorch and PyTorch-Exportable {{ pytorch_icon }} ## Theory @@ -70,6 +71,16 @@ dp compress -i graph.pb -o graph-compress.pb ``` ::: +:::{tab-item} TensorFlow 2 {{ tensorflow_icon }} + +```bash +dp --tf2 compress -i model.ckpt.tf2 -o model-compress.savedmodeltf +``` + +TensorFlow 2 compression reads a `.tf2` training checkpoint directory or a +checkpoint prefix and writes a compressed `.savedmodeltf` model. +::: + :::{tab-item} PyTorch {{ pytorch_icon }} ```bash @@ -77,6 +88,15 @@ dp --pt compress -i model.pth -o model-compress.pth ``` ::: +:::{tab-item} PyTorch-Exportable {{ pytorch_icon }} + +```bash +dp --pt-expt compress -i model.pte -o model-compress.pte +``` + +Use matching `.pte` or `.pt2` suffixes to preserve the exported model form. +::: + :::: where `-i` gives the original frozen model, `-o` gives the compressed model. Several other command line options can be passed to `dp compress`, which can be checked with diff --git a/doc/freeze/freeze.md b/doc/freeze/freeze.md index 3aef7e22a6..06f40186dd 100644 --- a/doc/freeze/freeze.md +++ b/doc/freeze/freeze.md @@ -15,6 +15,19 @@ in the folder where the model is trained. The output model is called `model.pb`. The idea and part of our code are from [Morgan](https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc). ::: +:::{tab-item} TensorFlow 2 {{ tensorflow_icon }} + +```bash +$ dp --tf2 freeze -o model.savedmodeltf +``` + +Run the command in the training folder. By default, it reads the +`model.ckpt.tf2` checkpoint directory and writes the TensorFlow SavedModel to +`model.savedmodeltf`. Use `-c` to select another checkpoint directory or +checkpoint prefix. For a multi-task checkpoint, select a branch with +`--head CHOSEN_BRANCH`. +::: + :::{tab-item} PyTorch {{ pytorch_icon }} ```bash @@ -33,6 +46,19 @@ $ dp --pt freeze -o model_branch1.pth --head CHOSEN_BRANCH The output model is called `model_branch1.pth`, which is the specifically frozen model with the `CHOSEN_BRANCH` head. ::: +:::{tab-item} PyTorch-Exportable {{ pytorch_icon }} + +```bash +$ dp --pt-expt freeze -c model.ckpt.pt -o model +``` + +The backend writes `.pte` for the dense neighbor-list lower form and `.pt2` for +the graph lower form. A suffixless output lets DeePMD-kit select the matching +extension. Use `--lower-kind nlist` or `--lower-kind graph` to request a form; +graph-capable DPA models may select the graph form automatically. In multi-task +mode, select a model branch with `--head CHOSEN_BRANCH`. +::: + :::{tab-item} Paddle {{ paddle_icon }} ```bash diff --git a/doc/train/training.md b/doc/train/training.md index df1341b4ee..ec6f594b1d 100644 --- a/doc/train/training.md +++ b/doc/train/training.md @@ -17,6 +17,13 @@ $ dp --tf train input.json ``` ::: +:::{tab-item} TensorFlow 2 {{ tensorflow_icon }} + +```bash +$ dp --tf2 train input.json +``` +::: + :::{tab-item} PyTorch {{ pytorch_icon }} ```bash @@ -24,6 +31,13 @@ $ dp --pt train input.json ``` ::: +:::{tab-item} PyTorch-Exportable {{ pytorch_icon }} + +```bash +$ dp --pt-expt train input.json +``` +::: + :::{tab-item} Paddle {{ paddle_icon }} ```bash From 135c44a3c782eec020033c76f73906e9d0e9a181 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 29 Jul 2026 23:12:40 +0800 Subject: [PATCH 02/10] 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 c71c21b53a..36fe1700de 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -61,15 +61,18 @@ allows DeePMD-kit to distinguish it from a regular PyTorch checkpoint. ### 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 06f40186dd..eeecd218c5 100644 --- a/doc/freeze/freeze.md +++ b/doc/freeze/freeze.md @@ -77,4 +77,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 ec6f594b1d..72e1dba168 100644 --- a/doc/train/training.md +++ b/doc/train/training.md @@ -53,6 +53,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. From 1cb0fa802e42cdcef64fff11b826364fa3cd7154 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 29 Jul 2026 23:50:50 +0800 Subject: [PATCH 03/10] docs(model): update backend support icons Document TF2, PyTorch-Exportable, and JAX availability across model and descriptor guides, and add DP/JAX compression routes. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- doc/backend.md | 10 +++++--- doc/freeze/compress.md | 39 +++++++++++++++++++++++++---- doc/freeze/freeze.md | 33 ++++++++++++------------ doc/model/dpa2.md | 8 +++--- doc/model/dpa3.md | 6 +++-- doc/model/dpa4.md | 19 +++++++++----- doc/model/dprc.md | 6 +++-- doc/model/linear.md | 3 ++- doc/model/pairtab.md | 14 ++++++++--- doc/model/train-energy-spin.md | 24 +++++++++++------- doc/model/train-energy.md | 4 ++- doc/model/train-fitting-dos.md | 9 ++++++- doc/model/train-fitting-property.md | 11 ++++++-- doc/model/train-fitting-tensor.md | 9 ++++++- doc/model/train-hybrid.md | 4 ++- doc/model/train-se-atten.md | 4 ++- doc/model/train-se-e2-a-tebd.md | 6 +++-- doc/model/train-se-e2-a.md | 4 ++- doc/model/train-se-e2-r.md | 4 ++- doc/model/train-se-e3-tebd.md | 6 +++-- doc/model/train-se-e3.md | 4 ++- doc/train/training.md | 24 +++++++++--------- 22 files changed, 173 insertions(+), 78 deletions(-) diff --git a/doc/backend.md b/doc/backend.md index 36fe1700de..fd53d1b12a 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -6,9 +6,11 @@ DeePMD-kit supports seven backends: TensorFlow, TensorFlow 2, PyTorch, PyTorch-Exportable, JAX, Paddle, and the NumPy-based DP reference backend. To use DeePMD-kit, you must install at least one backend. Each backend does not support all features. -In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch -{{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, and DP -{{ dpmodel_icon }} icons are used to mark whether a backend supports a feature. +In the documentation, TensorFlow and TensorFlow 2 share +{{ tensorflow_icon }}, while PyTorch and PyTorch-Exportable share +{{ pytorch_icon }}. JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, and DP +{{ dpmodel_icon }} use separate icons. Support notes spell out the exact backend +variant when the two implementations in a framework family differ. ### TensorFlow {{ tensorflow_icon }} @@ -61,7 +63,7 @@ allows DeePMD-kit to distinguish it from a regular PyTorch checkpoint. ### JAX {{ jax_icon }} -- Model filename extension: `.hlo`, `.savedmodel` +- Model filename extensions: `.hlo`, `.jax`, `.savedmodel` - Checkpoint filename extension: `.jax` [JAX](https://jax.readthedocs.io/) 0.4.33 or above is required. diff --git a/doc/freeze/compress.md b/doc/freeze/compress.md index 17ff2af1fe..0d8c9bf9c1 100644 --- a/doc/freeze/compress.md +++ b/doc/freeze/compress.md @@ -1,8 +1,9 @@ -# Compress a model {{ tensorflow_icon }} {{ pytorch_icon }} +# Compress a model {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow and TensorFlow 2 {{ tensorflow_icon }}, -> PyTorch and PyTorch-Exportable {{ pytorch_icon }} +> **Backends covered below**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, and DP {{ dpmodel_icon }}. ## Theory @@ -97,12 +98,40 @@ dp --pt-expt compress -i model.pte -o model-compress.pte Use matching `.pte` or `.pt2` suffixes to preserve the exported model form. ::: +:::{tab-item} JAX {{ jax_icon }} + +```bash +dp --jax compress -i frozen_model.hlo -o compressed_model.hlo +``` + +JAX compression accepts `.jax` and `.hlo` models. Use the same suffix for the +input and output to preserve the serialized model format. +::: + +:::{tab-item} DP {{ dpmodel_icon }} + +```bash +dp --dp compress -i model.dp -o model-compress.dp +``` + +DP compression accepts native `.dp` and `.yaml` models. +::: + :::: -where `-i` gives the original frozen model, `-o` gives the compressed model. Several other command line options can be passed to `dp compress`, which can be checked with +where `-i` gives the original frozen model, `-o` gives the compressed model. +The DP, JAX, and TensorFlow 2 entrypoints share the native-model compression +helpers for resolving the minimum neighbor distance and tabulating the +descriptor's embedding networks. PyTorch-Exportable implements the same table +strides and minimum-neighbor-distance fallback in its export-specific +entrypoint. If the model does not contain a minimum neighbor distance, pass the +training script with `-t` or `--training-script` so it can be computed from the +training data. + +Several other command line options can be passed to `dp compress`, which can be checked with ```bash -$ dp compress --help +dp compress --help ``` An explanation will be provided diff --git a/doc/freeze/freeze.md b/doc/freeze/freeze.md index eeecd218c5..703fd31d5b 100644 --- a/doc/freeze/freeze.md +++ b/doc/freeze/freeze.md @@ -8,7 +8,7 @@ To freeze a model, typically one does :::{tab-item} TensorFlow {{ tensorflow_icon }} ```bash -$ dp freeze -o model.pb +dp freeze -o model.pb ``` in the folder where the model is trained. The output model is called `model.pb`. @@ -18,20 +18,20 @@ The idea and part of our code are from [Morgan](https://blog.metaflow.fr/tensorf :::{tab-item} TensorFlow 2 {{ tensorflow_icon }} ```bash -$ dp --tf2 freeze -o model.savedmodeltf +dp --tf2 freeze -c model.ckpt -o model.savedmodeltf ``` -Run the command in the training folder. By default, it reads the -`model.ckpt.tf2` checkpoint directory and writes the TensorFlow SavedModel to -`model.savedmodeltf`. Use `-c` to select another checkpoint directory or -checkpoint prefix. For a multi-task checkpoint, select a branch with +When `-c` names a checkpoint prefix, the backend also checks the corresponding +path with `.tf2` appended, so the example reads `model.ckpt.tf2` and writes the +TensorFlow SavedModel to `model.savedmodeltf`. If `-c` is omitted, it defaults +to the current directory. For a multi-task checkpoint, select a branch with `--head CHOSEN_BRANCH`. ::: :::{tab-item} PyTorch {{ pytorch_icon }} ```bash -$ dp --pt freeze -o model.pth +dp --pt freeze -o model.pth ``` in the folder where the model is trained. The output model is called `model.pth`. @@ -40,7 +40,7 @@ In [multi-task mode](../train/multi-task-training), you need to choose one avail to specify which model branch you want to freeze: ```bash -$ dp --pt freeze -o model_branch1.pth --head CHOSEN_BRANCH +dp --pt freeze -o model_branch1.pth --head CHOSEN_BRANCH ``` The output model is called `model_branch1.pth`, which is the specifically frozen model with the `CHOSEN_BRANCH` head. @@ -49,20 +49,21 @@ The output model is called `model_branch1.pth`, which is the specifically frozen :::{tab-item} PyTorch-Exportable {{ pytorch_icon }} ```bash -$ dp --pt-expt freeze -c model.ckpt.pt -o model +dp --pt-expt freeze -c model.ckpt.pt -o model ``` The backend writes `.pte` for the dense neighbor-list lower form and `.pt2` for the graph lower form. A suffixless output lets DeePMD-kit select the matching -extension. Use `--lower-kind nlist` or `--lower-kind graph` to request a form; -graph-capable DPA models may select the graph form automatically. In multi-task -mode, select a model branch with `--head CHOSEN_BRANCH`. +extension. `--lower-kind graph` requires a graph-eligible model. Conversely, a +graph-capable DPA model may override a requested `nlist` lower with the graph +form and emit a warning. In multi-task mode, select a model branch with +`--head CHOSEN_BRANCH`. ::: :::{tab-item} Paddle {{ paddle_icon }} ```bash -$ dp --pd freeze -o model +dp --pd freeze -o model ``` in the folder where the model is trained. The output model is called `model.json` and `model.pdiparams`. @@ -71,7 +72,7 @@ In [multi-task mode](../train/multi-task-training.md), you need to choose one av to specify which model branch you want to freeze: ```bash -$ dp --pd freeze -o model_branch1 --head CHOSEN_BRANCH +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. @@ -80,7 +81,7 @@ The output model is called `model_branch1.json`, which is the specifically froze :::{tab-item} JAX {{ jax_icon }} ```bash -$ dp --jax freeze -c model.ckpt.jax -o model.hlo +dp --jax freeze -c model.ckpt.jax -o model.hlo ``` The JAX backend can write a StableHLO `.hlo` model, a lossless `.jax` model, or @@ -95,7 +96,7 @@ is the JAX format that supports the C++ inference interface. 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 +dp --jax freeze -c model.ckpt.jax -o model-hessian.hlo --hessian ``` The option applies to JAX `.hlo`, `.jax`, and `.savedmodel` outputs. A model diff --git a/doc/model/dpa2.md b/doc/model/dpa2.md index d5d2b9c91c..c6ef07cb28 100644 --- a/doc/model/dpa2.md +++ b/doc/model/dpa2.md @@ -1,7 +1,9 @@ -# Descriptor DPA-2 {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} +# Descriptor DPA-2 {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and +> PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle +> {{ paddle_icon }}, DP {{ dpmodel_icon }} The DPA-2 model implementation. See [DPA-2 paper](https://doi.org/10.1038/s41524-024-01493-2) for more details. @@ -96,7 +98,7 @@ The performance improvement will be limited if other parts are more expensive. In the pt_expt backend, a graph-eligible DPA-2 descriptor (`repinit/use_three_body` `false` -- the three-body sub-block is not graph-eligible -- and not compressed) can be frozen through a NeighborGraph-native inference path instead of the legacy dense neighbor-list path: ```bash -dp --pt_expt freeze -o model.pt2 --lower-kind graph +dp --pt-expt freeze -o model.pt2 --lower-kind graph ``` As with DPA-1's graph path (see [Difference among different backends](train-se-atten.md#difference-among-different-backends)), the graph route considers all neighbors within the cutoff rather than a fixed, padded selection, so its numeric result can differ slightly (down to the AOTInductor floating-point noise floor at non-binding `sel`, larger if `sel` is binding) from the dense/`nlist` path. diff --git a/doc/model/dpa3.md b/doc/model/dpa3.md index f52cfbdc1f..65d52f7c46 100644 --- a/doc/model/dpa3.md +++ b/doc/model/dpa3.md @@ -1,7 +1,9 @@ -# Descriptor DPA3 {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} +# Descriptor DPA3 {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and +> PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle +> {{ paddle_icon }}, DP {{ dpmodel_icon }} DPA3 is an advanced interatomic potential based on message passing. As a large atomic model (LAM), it is designed to integrate and jointly train on datasets from different domains, diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 66ff3363da..0352a46c52 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -1,7 +1,12 @@ -# Descriptor DPA4 {{ pytorch_icon }} +# Descriptor DPA4 {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }} +> **Descriptor backends**: PyTorch and PyTorch-Exportable +> {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> +> The end-to-end `model.type: dpa4` scaffold is supported by PyTorch and +> PyTorch-Exportable only. JAX support is descriptor-only and does not provide +> the dedicated DPA4 fitting/model/trainer workflow documented on this page. DPA4 is the DeePMD-kit implementation of the SeZM (Smooth Equivariant Zone-bridging Model) architecture: an SO(3)-equivariant message-passing model @@ -444,7 +449,7 @@ Two different export routes produce one: energy model reports `supports_edge_parallel() == True`, so the archive carries the with-comm artifact (`has_comm_artifact=true`) and **supports multi-rank LAMMPS out of the box** — no extra freeze options. -- **pt_expt (`dp --pt_expt freeze`).** Graph-capable models export through the +- **pt_expt (`dp --pt-expt freeze`).** Graph-capable models export through the **NeighborGraph** ABI (see [Graph-native inference route (pt_expt)](#graph-native-inference-route-pt_expt) below), which likewise embeds a with-comm artifact and supports multi-rank LAMMPS. @@ -507,7 +512,7 @@ neighbor-list path. Frame-level charge/spin conditioning `deepspin` virtual-atom spin scheme remains dense-only: ```bash -dp --pt_expt freeze -o model.pt2 --lower-kind graph +dp --pt-expt freeze -o model.pt2 --lower-kind graph ``` As with DPA-1's and DPA-2's graph paths (see [Difference among different @@ -565,7 +570,7 @@ inference](#multi-gpu-mpi-inference). - **Native scheme only.** `deepspin`-scheme spin (and the general `spin` virtual-atom model outside DPA4/SeZM) is dense-only; only `scheme: native` - is graph-eligible. `dp --pt_expt freeze --lower-kind graph` on a + is graph-eligible. `dp --pt-expt freeze --lower-kind graph` on a `deepspin`-scheme model raises an error at freeze time, per the dense/graph eligibility rule above. - **Graph route only, no dense fallback.** Unlike a plain-energy DPA4/SeZM @@ -720,7 +725,9 @@ closed over the one-hop neighbor shell. ## Limitations -- DPA4/SeZM is implemented for the PyTorch backend only. +- The end-to-end DPA4/SeZM model and training workflow is implemented for the + PyTorch and PyTorch-Exportable backends. JAX provides the descriptor only; + the DP implementation is a reference component and does not train models. - Export uses `.pt2` (AOTInductor); the TorchScript freeze path is not used. - Model compression is not supported. - Multi-rank (multi-GPU/MPI) LAMMPS inference works for a plain energy model diff --git a/doc/model/dprc.md b/doc/model/dprc.md index 7b24b0bb6d..e604271576 100644 --- a/doc/model/dprc.md +++ b/doc/model/dprc.md @@ -1,7 +1,9 @@ -# Deep Potential - Range Correction (DPRc) {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} +# Deep Potential - Range Correction (DPRc) {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, DP {{ dpmodel_icon }} Deep Potential - Range Correction (DPRc) is designed to combine with QM/MM method, and corrects energies from a low-level QM/MM method to a high-level QM/MM method: diff --git a/doc/model/linear.md b/doc/model/linear.md index 275e08c568..4380f5e622 100644 --- a/doc/model/linear.md +++ b/doc/model/linear.md @@ -1,7 +1,8 @@ ## Linear model {{ tensorflow_icon }} {{ pytorch_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} +> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch and +> PyTorch-Exportable {{ pytorch_icon }} One can linearly combine existing models with arbitrary coefficients: diff --git a/doc/model/pairtab.md b/doc/model/pairtab.md index 3cb6cf12f3..c43257c6c3 100644 --- a/doc/model/pairtab.md +++ b/doc/model/pairtab.md @@ -1,7 +1,12 @@ -# Interpolation or combination with a pairwise potential {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} +# Interpolation or combination with a pairwise potential {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }} +> **Short-range interpolation backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} +> +> **Linear-combination backends**: TensorFlow {{ tensorflow_icon }}, PyTorch and +> PyTorch-Exportable {{ pytorch_icon }} ## Theory @@ -44,7 +49,8 @@ In the range $[r_a, r_b]$, the DP model smoothly switched off and the pairwise p where the scale $\alpha_s$ is a tunable scale of the interatomic distance $r_{ij}$. The pairwise potential $u^{\textrm{pair}}(r)$ is defined by a user-defined table that provides the value of $u^{\textrm{pair}}$ on an evenly discretized grid from 0 to the cutoff distance.[^1] -DeePMD-kit also supports combination with a pairwise potential {{ tensorflow_icon }}: +DeePMD-kit also supports combination with a pairwise potential +{{ tensorflow_icon }} {{ pytorch_icon }}: ```math E_i = E_i^{\mathrm{DP}} + E_i^{\mathrm{pair}}, @@ -76,7 +82,7 @@ The interaction should be smooth at the cut-off distance. {ref}`sw_rmin ` and {ref}`sw_rmax ` must be smaller than the cutoff radius of the DP model. -## Combination with a pairwise potential {{ tensorflow_icon }} +## Combination with a pairwise potential {{ tensorflow_icon }} {{ pytorch_icon }} To combine with a pairwise potential, use the [linear model](./linear.md): diff --git a/doc/model/train-energy-spin.md b/doc/model/train-energy-spin.md index e71ff743fa..a80bf9a932 100644 --- a/doc/model/train-energy-spin.md +++ b/doc/model/train-energy-spin.md @@ -1,7 +1,8 @@ # Fit spin energy {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch and +> PyTorch-Exportable {{ pytorch_icon }}, DP {{ dpmodel_icon }} To train a model that takes additional spin information as input, you only need to modify the following sections to define the spin-specific settings, keeping other sections the same as the normal energy model's input script. @@ -11,7 +12,7 @@ keeping other sections the same as the normal energy model's input script. > > - In the TensorFlow backend, the `se_e2_a` descriptor will treat those atom types with spin as new (virtual) types, > and duplicate their corresponding selected numbers of neighbors ({ref}`sel `) from their real atom types. -> - In the PyTorch backend, if spin settings are added, all the types (with or without spin) will have their virtual types. +> - In the PyTorch and PyTorch-Exportable backends, if spin settings are added, all the types (with or without spin) will have their virtual types. > The `se_e2_a` descriptor will thus double the {ref}`sel ` list, > while in other descriptors with mixed types (such as `dpa1` or `dpa2`), the sel number will not be changed for clarity. > If you are using descriptors with mixed types, to achieve better performance, @@ -22,7 +23,8 @@ keeping other sections the same as the normal energy model's input script. The spin settings are given by the {ref}`spin ` section, which sets the magnetism for each type of atoms as described in the following sections. > [!NOTE] -> Note that the construction of spin settings is different between TensorFlow and PyTorch/DP. +> Note that the construction of spin settings is different between TensorFlow +> and PyTorch/PyTorch-Exportable/DP. ### Spin settings in TensorFlow @@ -41,9 +43,10 @@ The implementation in TensorFlow only supports `se_e2_a` descriptor. See example - {ref}`virtual_len ` specifies the distance between virtual atom and the belonging real atom. - {ref}`spin_norm ` gives the magnitude of the magnetic moment for each magnatic atom. -### Spin settings in PyTorch/DP +### Spin settings in PyTorch/PyTorch-Exportable/DP -In PyTorch/DP, the spin implementation is more flexible and so far supports the following descriptors: +In PyTorch/PyTorch-Exportable/DP, the spin implementation is more flexible and +so far supports the following descriptors: - `se_e2_a` - `dpa1`(`se_atten`) @@ -70,7 +73,7 @@ See `se_e2_a` examples in `$deepmd_source_dir/examples/spin/se_e2_a/input_torch. only used when {ref}`use_spin ` is True for each atom type. > [!NOTE] -> It should be noted that the spin models in PyTorch/DP are capable of addressing scenarios where the spin approaches zero +> It should be noted that the spin models in PyTorch/PyTorch-Exportable/DP are capable of addressing scenarios where the spin approaches zero > (indicating the virtual atom is in close proximity to the real atom) by adjusting the non-zero > {ref}`env_protection ` parameter within the descriptor. > This parameter is set to 0.01 by default in the spin model. It appears that a value of 0.01 is generally sufficient for maintaining model stability. @@ -143,7 +146,8 @@ If one does not want to train with virial, then he/she may set the virial prefac ## Data format > [!NOTE] -> Note that the spin data format is different between TensorFlow and PyTorch/DP. +> Note that the spin data format is different between TensorFlow and +> PyTorch/PyTorch-Exportable/DP. ### Spin data format in TensorFlow @@ -175,9 +179,11 @@ We list the details about spin system data format in TensorFlow backend: | force | Atomic and magnetic forces | force.raw | eV/Å | Nframes * (Natoms + Nspins) * 3 | The first `3 \* Natoms` columns represent atomic forces, followed by `3 \* Nspins` columns representing magnetic forces. | | virial | Frame virial | virial.raw | eV | Nframes * 9 | in the order `XX XY XZ YX YY YZ ZX ZY ZZ` | -### Spin data format in PyTorch/DP +### Spin data format in PyTorch/PyTorch-Exportable/DP -In the PyTorch backend, spin and magnetic forces are listed in separate files, and the data format may contain the following files: +In the PyTorch, PyTorch-Exportable, and DP backends, spin and magnetic forces +are listed in separate files, and the data format may contain the following +files: ``` type.raw diff --git a/doc/model/train-energy.md b/doc/model/train-energy.md index e2d380e6f0..447abc3414 100644 --- a/doc/model/train-energy.md +++ b/doc/model/train-energy.md @@ -1,7 +1,9 @@ # Fit energy {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} In this section, we will take `$deepmd_source_dir/examples/water/se_e2_a/input.json` as an example of the input file. diff --git a/doc/model/train-fitting-dos.md b/doc/model/train-fitting-dos.md index 357c045006..7560079651 100644 --- a/doc/model/train-fitting-dos.md +++ b/doc/model/train-fitting-dos.md @@ -1,7 +1,14 @@ # Fit electronic density of states (DOS) {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, DP {{ dpmodel_icon }} +> +> These icons report fitting/model implementation availability. JAX provides +> the DOS implementation for model conversion and inference, but +> `dp --jax train` currently accepts only the energy loss and cannot train a DOS +> model. Here we present an API to DeepDOS model, which can be used to fit electronic density of state (DOS) (which is a vector). diff --git a/doc/model/train-fitting-property.md b/doc/model/train-fitting-property.md index b71a0c7522..c93ae61795 100644 --- a/doc/model/train-fitting-property.md +++ b/doc/model/train-fitting-property.md @@ -1,7 +1,14 @@ -# Fit other properties {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} +# Fit other properties {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and +> PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, DP +> {{ dpmodel_icon }} +> +> These icons report fitting/model implementation availability. JAX provides +> the property implementation for model conversion and inference, but +> `dp --jax train` currently accepts only the energy loss and cannot train a +> property model. Here we present an API to DeepProperty model, which can be used to fit other properties like band gap, bulk modulus, critical temperature, etc. diff --git a/doc/model/train-fitting-tensor.md b/doc/model/train-fitting-tensor.md index 5df31bca0a..114c5a9f96 100644 --- a/doc/model/train-fitting-tensor.md +++ b/doc/model/train-fitting-tensor.md @@ -1,7 +1,14 @@ # Fit `tensor` like `Dipole` and `Polarizability` {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, DP {{ dpmodel_icon }} +> +> These icons report fitting/model implementation availability. JAX provides +> dipole and polarizability implementations for model conversion and inference, +> but `dp --jax train` currently accepts only the energy loss and cannot train +> these tensor models. Unlike `energy`, which is a scalar, one may want to fit some high dimensional physical quantity, like `dipole` (vector) and `polarizability` (matrix, shorted as `polar`). Deep Potential has provided different APIs to do this. In this example, we will show you how to train a model to fit a water system. A complete training input script of the examples can be found in diff --git a/doc/model/train-hybrid.md b/doc/model/train-hybrid.md index 06f5daf481..72c487f050 100644 --- a/doc/model/train-hybrid.md +++ b/doc/model/train-hybrid.md @@ -1,7 +1,9 @@ # Descriptor `"hybrid"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, DP {{ dpmodel_icon }} This descriptor hybridizes multiple descriptors to form a new descriptor. For example, we have a list of descriptors denoted by $\mathcal D_1$, $\mathcal D_2$, ..., $\mathcal D_N$, the hybrid descriptor this the concatenation of the list, i.e. $\mathcal D = (\mathcal D_1, \mathcal D_2, \cdots, \mathcal D_N)$. diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md index 05881436a1..db0ab80fdc 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -1,7 +1,9 @@ # Descriptor `"se_atten"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} ![ALT](../images/model_se_atten.png "model_se_atten") diff --git a/doc/model/train-se-e2-a-tebd.md b/doc/model/train-se-e2-a-tebd.md index 01bb0d1ad9..8678a994ec 100644 --- a/doc/model/train-se-e2-a-tebd.md +++ b/doc/model/train-se-e2-a-tebd.md @@ -1,7 +1,9 @@ -# Type embedding approach {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} +# Type embedding approach {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} We generate specific a type embedding vector for each atom type so that we can share one descriptor embedding net and one fitting net in total, which decline training complexity largely. diff --git a/doc/model/train-se-e2-a.md b/doc/model/train-se-e2-a.md index 7a7186590c..30c7210a56 100644 --- a/doc/model/train-se-e2-a.md +++ b/doc/model/train-se-e2-a.md @@ -1,7 +1,9 @@ # Descriptor `"se_e2_a"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} The notation of `se_e2_a` is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from all information (both angular and radial) of atomic configurations. The `e2` stands for the embedding with two-atoms information. This descriptor was described in detail in [the DeepPot-SE paper](https://arxiv.org/abs/1805.09003). diff --git a/doc/model/train-se-e2-r.md b/doc/model/train-se-e2-r.md index d5b8804788..2f2616937c 100644 --- a/doc/model/train-se-e2-r.md +++ b/doc/model/train-se-e2-r.md @@ -1,7 +1,9 @@ # Descriptor `"se_e2_r"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, DP {{ dpmodel_icon }} The notation of `se_e2_r` is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from the radial information of atomic configurations. The `e2` stands for the embedding with two-atom information. diff --git a/doc/model/train-se-e3-tebd.md b/doc/model/train-se-e3-tebd.md index 55295e6e86..e0a010245f 100644 --- a/doc/model/train-se-e3-tebd.md +++ b/doc/model/train-se-e3-tebd.md @@ -1,7 +1,9 @@ -# Descriptor `"se_e3_tebd"` {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} +# Descriptor `"se_e3_tebd"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and +> PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle +> {{ paddle_icon }}, DP {{ dpmodel_icon }} The notation of `se_e3_tebd` is short for the three-body embedding descriptor with type embeddings, where the notation `se` denotes the Deep Potential Smooth Edition (DeepPot-SE). The embedding takes bond angles between a central atom and its two neighboring atoms (denoted by `e3`) and their type embeddings (denoted by `tebd`) as input. diff --git a/doc/model/train-se-e3.md b/doc/model/train-se-e3.md index bac09dcced..5425900c1f 100644 --- a/doc/model/train-se-e3.md +++ b/doc/model/train-se-e3.md @@ -1,7 +1,9 @@ # Descriptor `"se_e3"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, DP {{ dpmodel_icon }} The notation of `se_e3` is short for three-body embedding DeepPot-SE, which incorporates embedded bond-angle information. The embedding takes bond angles between a central atom and its two neighboring atoms as input (denoted by `e3`). diff --git a/doc/train/training.md b/doc/train/training.md index 72e1dba168..58f84b1019 100644 --- a/doc/train/training.md +++ b/doc/train/training.md @@ -3,7 +3,7 @@ Several examples of training can be found in the `examples` directory: ```bash -$ cd $deepmd_source_dir/examples/water/se_e2_a/ +cd $deepmd_source_dir/examples/water/se_e2_a/ ``` After switching to that directory, the training can be invoked by @@ -13,50 +13,50 @@ After switching to that directory, the training can be invoked by :::{tab-item} TensorFlow {{ tensorflow_icon }} ```bash -$ dp --tf train input.json +dp --tf train input.json ``` ::: :::{tab-item} TensorFlow 2 {{ tensorflow_icon }} ```bash -$ dp --tf2 train input.json +dp --tf2 train input.json ``` ::: :::{tab-item} PyTorch {{ pytorch_icon }} ```bash -$ dp --pt train input.json +dp --pt train input.json ``` ::: :::{tab-item} PyTorch-Exportable {{ pytorch_icon }} ```bash -$ dp --pt-expt train input.json +dp --pt-expt train input.json ``` ::: :::{tab-item} Paddle {{ paddle_icon }} ```bash -# training model -$ dp --pd train input.json +# Training model +dp --pd train input.json # [experimental] training models with the CINN compiler (~40%+ speedup) # see: https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/guides/paddle_v3_features/cinn_cn.html -## If the shape(s) of batch input data are dynamic during training(default). -$ CINN=1 dp --pd train input.json -## If the shape(s) of batch input data are fixed during training, e.g., examples/water. -$ CINN=1 CINN_ALLOW_DYNAMIC_SHAPE=0 dp --pd train input.json +# If the batch input shapes are dynamic during training (the default). +CINN=1 dp --pd train input.json +# If the batch input shapes are fixed during training, e.g. examples/water. +CINN=1 CINN_ALLOW_DYNAMIC_SHAPE=0 dp --pd train input.json ``` ::: :::{tab-item} JAX {{ jax_icon }} ```bash -$ dp --jax train input.json +dp --jax train input.json ``` JAX training, introduced after DeePMD-kit v3.1.3, writes checkpoints with the From c6d6a9069623593210d9cdf59e04cc3d72d2f1b2 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 00:37:29 +0800 Subject: [PATCH 04/10] docs(backend): address review feedback Document the distinct .pt2 lower-input ABIs and add the missing DOS training commands for TensorFlow 2 and PyTorch-Exportable. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- doc/backend.md | 15 ++++++++++++++- doc/model/train-fitting-dos.md | 14 ++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/doc/backend.md b/doc/backend.md index fd53d1b12a..954921858b 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -61,6 +61,16 @@ model; graph-capable DPA models may select that form automatically. The `.pt` checkpoint format uses DP-model parameter names ending in `.w` and `.b`, which allows DeePMD-kit to distinguish it from a regular PyTorch checkpoint. +The `.pt2` suffix identifies an AOTInductor package, but not its lower-input +ABI. A DPA4/SeZM model frozen with `dp --pt freeze` uses the legacy `edge_vec` +ABI (`lower_input_kind: edge_vec`), whereas a graph model frozen with +`dp --pt-expt freeze --lower-kind graph` uses the NeighborGraph ABI +(`lower_input_kind: graph`). Both variants are loaded for inference by the +PyTorch-Exportable runtime, which reads this metadata to select the correct +input path. The `--lower-kind` option controls only the PyTorch-Exportable +freeze route; see the [DPA4 export documentation](model/dpa4.md#freeze-to-pt2) +for the separate PyTorch route. + ### JAX {{ jax_icon }} - Model filename extensions: `.hlo`, `.jax`, `.savedmodel` @@ -114,7 +124,10 @@ backend. When doing inference, DeePMD-kit detects the backend from the model filename. For example, when the model filename ends with `.pb` (the ProtoBuf file), DeePMD-kit will consider it using the TensorFlow backend. The same detection covers TensorFlow 2 `.savedmodeltf` models and -PyTorch-Exportable `.pte` and `.pt2` models. +PyTorch-Exportable `.pte` and `.pt2` runtime formats. In particular, `.pt2` +selects the PyTorch-Exportable inference loader even when the file was produced +by the DPA4/SeZM `dp --pt freeze` route described above; the archive metadata +then selects its `edge_vec` or NeighborGraph ABI. ## Convert model files between backends diff --git a/doc/model/train-fitting-dos.md b/doc/model/train-fitting-dos.md index 7560079651..01d280e4cc 100644 --- a/doc/model/train-fitting-dos.md +++ b/doc/model/train-fitting-dos.md @@ -97,6 +97,13 @@ dp --tf train input.json ``` ::: +:::{tab-item} TensorFlow 2 {{ tensorflow_icon }} + +```bash +dp --tf2 train input.json +``` +::: + :::{tab-item} PyTorch {{ pytorch_icon }} ```bash @@ -104,6 +111,13 @@ dp --pt train input.json ``` ::: +:::{tab-item} PyTorch-Exportable {{ pytorch_icon }} + +```bash +dp --pt-expt train input.json +``` +::: + :::: The detailed loss can be found in `lcurve.out`: From bd78747e5fb790fdb9c06e96d24f0d4be8da3ff7 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 00:47:32 +0800 Subject: [PATCH 05/10] docs(backend): clarify DPA4 spin ABI Document the dense nlist exception used by PT-frozen deepspin DPA4 models. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- doc/backend.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/doc/backend.md b/doc/backend.md index 954921858b..7e30025712 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -62,10 +62,11 @@ checkpoint format uses DP-model parameter names ending in `.w` and `.b`, which allows DeePMD-kit to distinguish it from a regular PyTorch checkpoint. The `.pt2` suffix identifies an AOTInductor package, but not its lower-input -ABI. A DPA4/SeZM model frozen with `dp --pt freeze` uses the legacy `edge_vec` -ABI (`lower_input_kind: edge_vec`), whereas a graph model frozen with +ABI. A DPA4/SeZM model frozen with `dp --pt freeze` normally uses the legacy +`edge_vec` ABI (`lower_input_kind: edge_vec`); its deepspin virtual-atom variant +uses the dense `nlist` ABI instead. A graph model frozen with `dp --pt-expt freeze --lower-kind graph` uses the NeighborGraph ABI -(`lower_input_kind: graph`). Both variants are loaded for inference by the +(`lower_input_kind: graph`). All variants are loaded for inference by the PyTorch-Exportable runtime, which reads this metadata to select the correct input path. The `--lower-kind` option controls only the PyTorch-Exportable freeze route; see the [DPA4 export documentation](model/dpa4.md#freeze-to-pt2) @@ -127,7 +128,7 @@ The same detection covers TensorFlow 2 `.savedmodeltf` models and PyTorch-Exportable `.pte` and `.pt2` runtime formats. In particular, `.pt2` selects the PyTorch-Exportable inference loader even when the file was produced by the DPA4/SeZM `dp --pt freeze` route described above; the archive metadata -then selects its `edge_vec` or NeighborGraph ABI. +then selects its `edge_vec`, dense `nlist`, or NeighborGraph ABI. ## Convert model files between backends From 21363e2676ca1263fd11062d56fb077c9873b4ba Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 31 Jul 2026 13:56:24 +0800 Subject: [PATCH 06/10] docs(backend): correct workflow limitations Document the actual export, compression, inference, and multi-rank limits across TensorFlow 2, PyTorch-Exportable, JAX, Paddle, and DP backends. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- doc/backend.md | 26 +++++++++---- doc/freeze/compress.md | 34 ++++++++++------ doc/freeze/freeze.md | 13 +++++-- doc/model/dpa4.md | 65 ++++++++++++------------------- doc/model/linear.md | 4 +- doc/model/overall.md | 5 ++- doc/model/pairtab.md | 10 ++--- doc/model/train-energy-hessian.md | 5 ++- doc/model/train-fitting-dos.md | 4 +- 9 files changed, 91 insertions(+), 75 deletions(-) diff --git a/doc/backend.md b/doc/backend.md index 7e30025712..ef3034cc98 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -27,9 +27,13 @@ DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`t The TensorFlow 2 backend uses the TensorFlow v2 eager API. Select it with `dp --tf2` (alias `dp --tensorflow2`). It supports training, including -multi-task training and fine-tuning, freezing, compression, and testing. -Training stores checkpoints in a directory named after the `save_ckpt` prefix -with `.tf2` appended, such as `model.ckpt.tf2`. +multi-task training and fine-tuning. Freezing, compression, and testing use a +`.savedmodeltf` export and therefore require a graph-traceable model. In +particular, `se_e2_a`/`se_a` descriptors currently require +`type_one_side: true`; the normalized default, `false`, cannot be exported. +Frozen TensorFlow 2 DOS models also cannot yet be tested because the SavedModel +does not preserve `numb_dos`. Training stores checkpoints in a directory named +after the `save_ckpt` prefix with `.tf2` appended, such as `model.ckpt.tf2`. Setting [`DP_JIT`](env.md#envvar-DP_JIT) enables optional `tf.function` JIT compilation; depending on the workload, this may improve or reduce performance. @@ -50,8 +54,11 @@ While `.pth` and `.pt` are the same in the PyTorch package, they have different Select this backend with `dp --pt-expt` (alias `dp --pytorch-exportable`). It uses PyTorch with the backend-independent model implementation and supports training, including multi-task training and -fine-tuning, freezing, compression, change-bias, and testing. Training can read -LMDB datasets, and Python inference can use the optional vesin neighbor-list +fine-tuning, freezing, change-bias, and testing. Executable compression is +currently limited to fused-CUDA-eligible, graph-lowered DPA1/`se_atten` strip +models exported as `.pt2`; other `dp --pt-expt compress` outputs do not replace +the exported inference graph with the tabulated model. Training can read LMDB +datasets, and Python inference can use the optional vesin neighbor-list implementation. Freezing exports a `torch.export` model. The dense neighbor-list lower form @@ -74,8 +81,8 @@ for the separate PyTorch route. ### JAX {{ jax_icon }} -- Model filename extensions: `.hlo`, `.jax`, `.savedmodel` -- Checkpoint filename extension: `.jax` +- DeepEval model filename extensions: `.hlo`, `.savedmodel` +- Checkpoint and lossless serialization extension: `.jax` [JAX](https://jax.readthedocs.io/) 0.4.33 or above is required. Both `.hlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions. @@ -84,7 +91,10 @@ Only the `.savedmodel` format supports C++ inference, which needs the TensorFlow 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. +`.jax` extension. Freezing can write a DeepEval-compatible `.hlo` or +`.savedmodel` model, or a lossless `.jax` serialization for checkpoint +round-tripping and JAX-MD. The normal `dp test`/`DeepPot` route does not load +`.jax` serializations. ### Paddle {{ paddle_icon }} diff --git a/doc/freeze/compress.md b/doc/freeze/compress.md index 0d8c9bf9c1..7f55e44753 100644 --- a/doc/freeze/compress.md +++ b/doc/freeze/compress.md @@ -79,7 +79,9 @@ dp --tf2 compress -i model.ckpt.tf2 -o model-compress.savedmodeltf ``` TensorFlow 2 compression reads a `.tf2` training checkpoint directory or a -checkpoint prefix and writes a compressed `.savedmodeltf` model. +checkpoint prefix and writes a compressed `.savedmodeltf` model. For +`se_e2_a`/`se_a`, SavedModel export currently requires +`type_one_side: true`; the normalized default, `false`, is not graph-traceable. ::: :::{tab-item} PyTorch {{ pytorch_icon }} @@ -92,20 +94,30 @@ dp --pt compress -i model.pth -o model-compress.pth :::{tab-item} PyTorch-Exportable {{ pytorch_icon }} ```bash -dp --pt-expt compress -i model.pte -o model-compress.pte +dp --pt-expt compress -i dpa1-graph.pt2 -o dpa1-graph-compress.pt2 ``` -Use matching `.pte` or `.pt2` suffixes to preserve the exported model form. +This executable compression path is currently limited to graph-lowered +DPA1/`se_atten` strip models that are eligible for the fused CUDA table +operator: `attn_layer: 0`, no excluded type pairs, float32 descriptor +statistics and tables, `neuron[-1] <= 256`, and +`axis_neuron <= min(16, neuron[-1])`. Their `.pt2` export embeds the tabulated +operator. For other models, the command records compressed state in +`model.json` but leaves the exported inference graph uncompressed, so the +resulting `.pte`/`.pt2` is not a deployment compression artifact. ::: :::{tab-item} JAX {{ jax_icon }} ```bash -dp --jax compress -i frozen_model.hlo -o compressed_model.hlo +dp --jax compress -i frozen_model.jax -o compressed_model.jax ``` -JAX compression accepts `.jax` and `.hlo` models. Use the same suffix for the -input and output to preserve the serialized model format. +JAX compression accepts `.jax` and `.hlo` inputs. Use `.jax` for the general, +lossless compressed serialization path. A compressed `.hlo` can be written +only when the descriptor is StableHLO-exportable; compressed `se_e2_a`, `se_a`, +`dpa1`, and `se_atten` descriptors with `type_one_side: false` must remain +`.jax` or be rebuilt with `type_one_side: true`. ::: :::{tab-item} DP {{ dpmodel_icon }} @@ -122,11 +134,11 @@ DP compression accepts native `.dp` and `.yaml` models. where `-i` gives the original frozen model, `-o` gives the compressed model. The DP, JAX, and TensorFlow 2 entrypoints share the native-model compression helpers for resolving the minimum neighbor distance and tabulating the -descriptor's embedding networks. PyTorch-Exportable implements the same table -strides and minimum-neighbor-distance fallback in its export-specific -entrypoint. If the model does not contain a minimum neighbor distance, pass the -training script with `-t` or `--training-script` so it can be computed from the -training data. +descriptor's embedding networks. The limited PyTorch-Exportable graph-DPA1 +path implements the same table strides and minimum-neighbor-distance fallback +in its export-specific entrypoint. If the model does not contain a minimum +neighbor distance, pass the training script with `-t` or `--training-script` +so it can be computed from the training data. Several other command line options can be passed to `dp compress`, which can be checked with diff --git a/doc/freeze/freeze.md b/doc/freeze/freeze.md index 703fd31d5b..69f6c589f3 100644 --- a/doc/freeze/freeze.md +++ b/doc/freeze/freeze.md @@ -25,7 +25,9 @@ When `-c` names a checkpoint prefix, the backend also checks the corresponding path with `.tf2` appended, so the example reads `model.ckpt.tf2` and writes the TensorFlow SavedModel to `model.savedmodeltf`. If `-c` is omitted, it defaults to the current directory. For a multi-task checkpoint, select a branch with -`--head CHOSEN_BRANCH`. +`--head CHOSEN_BRANCH`. SavedModel export currently requires graph-traceable +model code; an `se_e2_a`/`se_a` descriptor must set `type_one_side: true` +because the normalized default, `false`, cannot be traced. ::: :::{tab-item} PyTorch {{ pytorch_icon }} @@ -84,9 +86,12 @@ The output model is called `model_branch1.json`, which is the specifically froze 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. +The JAX backend can write a StableHLO `.hlo` model, a lossless `.jax` +serialization, or a JAX2TF `.savedmodel` model. The `.hlo` and `.savedmodel` +formats work with the normal `dp test`/`DeepPot` route; `.jax` is intended for +checkpoint round-tripping and JAX-MD and is not a DeepEval model format. The +`.savedmodel` format requires TensorFlow and is the JAX format that supports the +C++ inference interface. ::: :::: diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 0352a46c52..c27aa048c9 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -541,22 +541,17 @@ graph-capable model is always frozen to the graph lower in any case, since the dense lower is deprecated in the pt_expt backend. See [Native spin (magnetic)](#native-spin-magnetic) below. Unlike the dense route (see [Multi-GPU (MPI) -inference](#multi-gpu-mpi-inference) above), a graph-frozen `.pt2` **of a -plain-energy (non-spin) model** embeds a with-comm AOTInductor artifact and -supports multi-rank LAMMPS: each block's cross-rank ghost-feature exchange -runs through the `border_op` MPI path once per interaction block, the same -mechanism used by DPA-2's graph route (see the "Graph-native inference route -(pt_expt)" section of [DPA-2's documentation](dpa2.md)). As on DPA-2, -multi-rank inference on the graph route requires every MPI rank to own or -ghost at least one atom; a rank with zero atoms in both categories aborts the -run collectively rather than silently desynchronizing the per-block -exchange. Pick a domain decomposition that keeps every rank non-empty, or use -the dense route, which has no such restriction (but is single-rank only, as -noted above). - -**Native-spin graph `.pt2` archives are the exception: they carry no -with-comm artifact and are single-rank only** -- see [Native spin -(magnetic)](#native-spin-magnetic) below. +inference](#multi-gpu-mpi-inference) above), a graph-frozen `.pt2` whose +descriptor communicates across ranks embeds a with-comm AOTInductor artifact. +This includes both plain-energy and native-spin DPA4/SeZM models. Each block's +cross-rank ghost-feature exchange runs through the `border_op` MPI path once +per interaction block, the same mechanism used by DPA-2's graph route (see the +"Graph-native inference route (pt_expt)" section of [DPA-2's +documentation](dpa2.md)). As on DPA-2, multi-rank inference on the graph route +requires every MPI rank to own or ghost at least one atom; a rank with zero +atoms in both categories aborts the run collectively rather than silently +desynchronizing the per-block exchange. Pick a domain decomposition that keeps +every rank non-empty, or use the dense route where the model provides one. ### Native spin (magnetic) @@ -578,17 +573,15 @@ inference](#multi-gpu-mpi-inference). a native-spin descriptor has only the graph lower. `--lower-kind auto` (the default) resolves to `graph`; `--lower-kind nlist` is not a valid option for a native-spin model. -- **Single-rank only.** The frozen archive's `has_comm_artifact` metadata is - `false` for native-spin models (no ghost-spin cross-rank exchange is - implemented), so a multi-rank LAMMPS run fails fast with an explicit error - at the first force evaluation, mirroring the dense route's single-rank - restriction described in [Multi-GPU (MPI) - inference](#multi-gpu-mpi-inference). Run native-spin models on a single - MPI rank (a single GPU, or CPU without `mpirun`). -- **Spin is per local atom.** The `spin` input is `(nframes, nloc, 3)` -- - one vector per *local* atom, not per ghost/extended atom (`nall`); there is - no ghost-spin exchange to populate ghost spins across a rank boundary, - consistent with the single-rank restriction above. +- **Multi-rank graph inference.** When the descriptor communicates across + ranks, the frozen archive records `has_comm_artifact=true` and embeds + `model/extra/forward_lower_with_comm.pt2`. `DeepSpinPTExpt` selects that + artifact for multi-rank LAMMPS and refreshes ghost node features through + `border_op` at every interaction block. +- **Spin is supplied per local atom.** The user-facing `spin` input is + `(nframes, nloc, 3)`. In a multi-rank LAMMPS run, ghost spin rows are + populated by the LAMMPS `sp` forward communication before the graph lower + runs; spin does not require a separate model-side exchange. - **The magnetic force is a second energy gradient.** As in the general native-scheme convention (see [Spin](#spin) above), `force_mag = -\partial E/\partial\mathbf{s}`, computed by pt_expt as a @@ -599,19 +592,11 @@ inference](#multi-gpu-mpi-inference). are `None` placeholders there, exactly as for the plain-energy dpmodel route. -The following combinations are **not yet supported** on the native-spin -graph route (follow-up work): - -- **Multi-rank inference.** Ghost-spin cross-rank exchange (analogous to the - plain-energy graph route's `border_op`-based ghost-feature exchange) is not - implemented. -- **Charge-spin FiLM conditioning.** Combining `add_chg_spin_ebd` with - `spin.scheme: native` is rejected at model-construction time; use one or - the other. -- **ZBL zone bridging.** Combining `bridging_method: ZBL` with - `spin.scheme: native` is not supported on the pt_expt backend (`bridging_method` - is rejected there independently of spin -- see [Zone bridging - (ZBL)](#zone-bridging-zbl)). +Charge-spin FiLM conditioning can be combined with native spin. ZBL zone +bridging is also graph-eligible, but a native-spin-plus-ZBL archive remains +single-rank because the analytical bridging reduction requires each node's +complete outgoing-edge set; this is the same ZBL limitation described in +[Multi-GPU (MPI) inference](#multi-gpu-mpi-inference). ## Embedding extraction diff --git a/doc/model/linear.md b/doc/model/linear.md index 4380f5e622..632df3b9c9 100644 --- a/doc/model/linear.md +++ b/doc/model/linear.md @@ -1,8 +1,8 @@ -## Linear model {{ tensorflow_icon }} {{ pytorch_icon }} +## Linear model {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} > [!NOTE] > **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch and -> PyTorch-Exportable {{ pytorch_icon }} +> PyTorch-Exportable {{ pytorch_icon }}, DP {{ dpmodel_icon }} One can linearly combine existing models with arbitrary coefficients: diff --git a/doc/model/overall.md b/doc/model/overall.md index b3ffc0e629..731b5dd031 100644 --- a/doc/model/overall.md +++ b/doc/model/overall.md @@ -58,8 +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] -> The PyTorch TorchScript freeze route cannot output Hessians. The JAX backend -> can retain Hessian output in a frozen model with +> The PyTorch TorchScript freeze route cannot output Hessians, and +> PyTorch-Exportable cannot construct a Hessian model for freezing or +> inference. 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. diff --git a/doc/model/pairtab.md b/doc/model/pairtab.md index c43257c6c3..0d94efdb27 100644 --- a/doc/model/pairtab.md +++ b/doc/model/pairtab.md @@ -1,12 +1,12 @@ -# Interpolation or combination with a pairwise potential {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} +# Interpolation or combination with a pairwise potential {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] > **Short-range interpolation backends**: TensorFlow and TensorFlow 2 > {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX -> {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} +> {{ jax_icon }}, DP {{ dpmodel_icon }} > > **Linear-combination backends**: TensorFlow {{ tensorflow_icon }}, PyTorch and -> PyTorch-Exportable {{ pytorch_icon }} +> PyTorch-Exportable {{ pytorch_icon }}, DP {{ dpmodel_icon }} ## Theory @@ -50,7 +50,7 @@ where the scale $\alpha_s$ is a tunable scale of the interatomic distance $r_{ij The pairwise potential $u^{\textrm{pair}}(r)$ is defined by a user-defined table that provides the value of $u^{\textrm{pair}}$ on an evenly discretized grid from 0 to the cutoff distance.[^1] DeePMD-kit also supports combination with a pairwise potential -{{ tensorflow_icon }} {{ pytorch_icon }}: +{{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }}: ```math E_i = E_i^{\mathrm{DP}} + E_i^{\mathrm{pair}}, @@ -82,7 +82,7 @@ The interaction should be smooth at the cut-off distance. {ref}`sw_rmin ` and {ref}`sw_rmax ` must be smaller than the cutoff radius of the DP model. -## Combination with a pairwise potential {{ tensorflow_icon }} {{ pytorch_icon }} +## Combination with a pairwise potential {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} To combine with a pairwise potential, use the [linear model](./linear.md): diff --git a/doc/model/train-energy-hessian.md b/doc/model/train-energy-hessian.md index d99d6d0c1f..3fd90a92b3 100644 --- a/doc/model/train-energy-hessian.md +++ b/doc/model/train-energy-hessian.md @@ -93,8 +93,9 @@ The detailed loss can be found in `lcurve.out`: > [!WARNING] > 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`. +> model. PyTorch-Exportable cannot construct a Hessian model for freezing or +> inference. The JAX backend can preserve Hessian output in a frozen model +> with `dp --jax freeze --hessian`. If one freezes and tests a Hessian model through the PyTorch route: diff --git a/doc/model/train-fitting-dos.md b/doc/model/train-fitting-dos.md index 01d280e4cc..ae5a5ef60b 100644 --- a/doc/model/train-fitting-dos.md +++ b/doc/model/train-fitting-dos.md @@ -8,7 +8,9 @@ > These icons report fitting/model implementation availability. JAX provides > the DOS implementation for model conversion and inference, but > `dp --jax train` currently accepts only the energy loss and cannot train a DOS -> model. +> model. TensorFlow 2 can train a DOS model, but a frozen `.savedmodeltf` DOS +> model cannot yet be evaluated by `dp test`/`DeepDOS` because the export does +> not preserve the `numb_dos` output width. Here we present an API to DeepDOS model, which can be used to fit electronic density of state (DOS) (which is a vector). From f9c46aa230309b9faa6486e420a43923158a8e86 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 1 Aug 2026 13:18:31 +0800 Subject: [PATCH 07/10] docs(backend): address remaining review feedback Move descriptor-specific backend limitations to their model pages, remove the broken PyTorch-Exportable DOS workflow, and distinguish the legacy TorchScript backend throughout the affected documentation. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- doc/backend.md | 51 +++++++++++++---------- doc/development/create-a-model-pt.md | 14 +++---- doc/freeze/compress.md | 42 +++++++------------ doc/freeze/freeze.md | 8 ++-- doc/inference/python.md | 2 +- doc/model/change-bias.md | 4 +- doc/model/dpa2.md | 4 +- doc/model/dpa3.md | 4 +- doc/model/dpa4.md | 13 +++--- doc/model/dprc.md | 4 +- doc/model/linear.md | 2 +- doc/model/overall.md | 2 +- doc/model/pairtab.md | 4 +- doc/model/sel.md | 2 +- doc/model/train-energy-hessian.md | 26 ++++++------ doc/model/train-energy-spin.md | 18 ++++---- doc/model/train-energy.md | 6 +-- doc/model/train-fitting-dos.md | 21 ++++------ doc/model/train-fitting-population.md | 4 +- doc/model/train-fitting-property.md | 6 +-- doc/model/train-fitting-tensor.md | 8 ++-- doc/model/train-hybrid.md | 2 +- doc/model/train-se-atten.md | 30 +++++++++++-- doc/model/train-se-e2-a-tebd.md | 2 +- doc/model/train-se-e2-a.md | 13 +++++- doc/model/train-se-e2-r.md | 4 +- doc/model/train-se-e3-tebd.md | 2 +- doc/model/train-se-e3.md | 2 +- doc/third-party/ase.md | 2 +- doc/train/parallel-training.md | 6 +-- doc/train/training-advanced.md | 4 +- doc/train/training.md | 4 +- doc/troubleshooting/howtoset_num_nodes.md | 6 +-- 33 files changed, 174 insertions(+), 148 deletions(-) diff --git a/doc/backend.md b/doc/backend.md index ef3034cc98..f15b5fb5c0 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -2,12 +2,13 @@ ## Supported backends -DeePMD-kit supports seven backends: TensorFlow, TensorFlow 2, PyTorch, -PyTorch-Exportable, JAX, Paddle, and the NumPy-based DP reference backend. +DeePMD-kit supports seven backends: TensorFlow, TensorFlow 2, +PyTorch-TorchScript, PyTorch-Exportable, JAX, Paddle, and the NumPy-based DP +reference backend. To use DeePMD-kit, you must install at least one backend. Each backend does not support all features. In the documentation, TensorFlow and TensorFlow 2 share -{{ tensorflow_icon }}, while PyTorch and PyTorch-Exportable share +{{ tensorflow_icon }}, while PyTorch-TorchScript and PyTorch-Exportable share {{ pytorch_icon }}. JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, and DP {{ dpmodel_icon }} use separate icons. Support notes spell out the exact backend variant when the two implementations in a framework family differ. @@ -28,23 +29,30 @@ DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`t The TensorFlow 2 backend uses the TensorFlow v2 eager API. Select it with `dp --tf2` (alias `dp --tensorflow2`). It supports training, including multi-task training and fine-tuning. Freezing, compression, and testing use a -`.savedmodeltf` export and therefore require a graph-traceable model. In -particular, `se_e2_a`/`se_a` descriptors currently require -`type_one_side: true`; the normalized default, `false`, cannot be exported. -Frozen TensorFlow 2 DOS models also cannot yet be tested because the SavedModel -does not preserve `numb_dos`. Training stores checkpoints in a directory named -after the `save_ckpt` prefix with `.tf2` appended, such as `model.ckpt.tf2`. +`.savedmodeltf` export and therefore require graph-traceable model code. +Training stores checkpoints in a directory named after the `save_ckpt` prefix +with `.tf2` appended, such as `model.ckpt.tf2`. -Setting [`DP_JIT`](env.md#envvar-DP_JIT) enables optional `tf.function` JIT -compilation; depending on the workload, this may improve or reduce performance. +For training, set +{ref}`training.enable_compile ` to `true` to enable +XLA compilation of the formatted lower-forward path. Setting +[`DP_JIT`](env.md#envvar-DP_JIT) enables the same model-level default and also +applies it to SavedModel export. Depending on the workload, compilation may +improve or reduce performance. -### PyTorch {{ pytorch_icon }} +### PyTorch-TorchScript {{ pytorch_icon }} - Model filename extension: `.pth` - Checkpoint filename extension: `.pt` -[PyTorch](https://pytorch.org/) 2.1 or above is required. -While `.pth` and `.pt` are the same in the PyTorch package, they have different meanings in the DeePMD-kit to distinguish the model and the checkpoint. +[PyTorch](https://pytorch.org/) 2.1 or above is required. Select this backend +with `dp --pt`. It uses TorchScript for most frozen models; DPA4/SeZM uses a +separate AOTInductor export path. Because PyTorch has deprecated TorchScript, +DeePMD-kit will deprecate this backend and replace it with PyTorch-Exportable. + +While `.pth` and `.pt` are the same in the PyTorch package, they have different +meanings in DeePMD-kit: `.pth` stores a frozen model, while `.pt` stores a +training checkpoint. ### PyTorch-Exportable {{ pytorch_icon }} @@ -54,19 +62,18 @@ While `.pth` and `.pt` are the same in the PyTorch package, they have different Select this backend with `dp --pt-expt` (alias `dp --pytorch-exportable`). It uses PyTorch with the backend-independent model implementation and supports training, including multi-task training and -fine-tuning, freezing, change-bias, and testing. Executable compression is -currently limited to fused-CUDA-eligible, graph-lowered DPA1/`se_atten` strip -models exported as `.pt2`; other `dp --pt-expt compress` outputs do not replace -the exported inference graph with the tabulated model. Training can read LMDB -datasets, and Python inference can use the optional vesin neighbor-list -implementation. +fine-tuning, freezing, change-bias, and testing. Compression support and export +requirements are documented on the corresponding descriptor pages. Training +can read LMDB datasets, and Python inference can use the optional vesin +neighbor-list implementation. Freezing exports a `torch.export` model. The dense neighbor-list lower form normally uses `.pte`, while the graph lower form uses an AOTInductor `.pt2` package. Use `--lower-kind graph` to request graph-native export for an eligible model; graph-capable DPA models may select that form automatically. The `.pt` checkpoint format uses DP-model parameter names ending in `.w` and `.b`, which -allows DeePMD-kit to distinguish it from a regular PyTorch checkpoint. +allows DeePMD-kit to distinguish it from a PyTorch-TorchScript checkpoint, +whose parameter names end in `.matrix` and `.bias`. The `.pt2` suffix identifies an AOTInductor package, but not its lower-input ABI. A DPA4/SeZM model frozen with `dp --pt freeze` normally uses the legacy @@ -77,7 +84,7 @@ uses the dense `nlist` ABI instead. A graph model frozen with PyTorch-Exportable runtime, which reads this metadata to select the correct input path. The `--lower-kind` option controls only the PyTorch-Exportable freeze route; see the [DPA4 export documentation](model/dpa4.md#freeze-to-pt2) -for the separate PyTorch route. +for the separate PyTorch-TorchScript route. ### JAX {{ jax_icon }} diff --git a/doc/development/create-a-model-pt.md b/doc/development/create-a-model-pt.md index ce902fde5a..b70abc9e9e 100644 --- a/doc/development/create-a-model-pt.md +++ b/doc/development/create-a-model-pt.md @@ -1,9 +1,9 @@ # Create a model in other backends {{ pytorch_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: PyTorch-TorchScript {{ pytorch_icon }}, DP {{ dpmodel_icon }} > -> In the following context, we use the PyTorch backend as the example, while it also applies to other backends listed above. +> In the following context, we use the PyTorch-TorchScript backend as the example, while it also applies to other backends listed above. If you'd like to create a new model that isn't covered by the existing DeePMD-kit library, but reuse DeePMD-kit's other efficient modules such as data processing, trainer, etc, you may want to read this section. @@ -16,9 +16,9 @@ To incorporate your custom model you'll need to: ## Design a new component -With DeePMD-kit v3, we have expanded support to include two additional backends alongside TensorFlow: the PyTorch backend and the framework-independent backend (dpmodel). The PyTorch backend adopts a highly modularized design to provide flexibility and extensibility. It ensures a consistent experience for both training and inference, aligning with the TensorFlow backend. +With DeePMD-kit v3, we have expanded support to include two additional backends alongside TensorFlow: the PyTorch-TorchScript backend and the framework-independent backend (dpmodel). The PyTorch-TorchScript backend adopts a highly modularized design to provide flexibility and extensibility. It ensures a consistent experience for both training and inference, aligning with the TensorFlow backend. -The framework-independent backend is implemented in pure NumPy, serving as a reference backend to ensure consistency in tests. Its design pattern closely parallels that of the PyTorch backend. +The framework-independent backend is implemented in pure NumPy, serving as a reference backend to ensure consistency in tests. Its design pattern closely parallels that of the PyTorch-TorchScript backend. ### New descriptors @@ -118,7 +118,7 @@ class SomeFittingNet(GeneralFitting): ### New models -The PyTorch backend's model architecture is meticulously structured with multiple layers of abstraction, ensuring a high degree of flexibility. Typically, the process commences with an atomic model responsible for atom-wise property calculations. This atomic model inherits from both the {py:class}`deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel` class and the {py:class}`torch.nn.Module` class. +The PyTorch-TorchScript backend's model architecture is meticulously structured with multiple layers of abstraction, ensuring a high degree of flexibility. Typically, the process commences with an atomic model responsible for atom-wise property calculations. This atomic model inherits from both the {py:class}`deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel` class and the {py:class}`torch.nn.Module` class. Subsequently, the `AtomicModel` is encapsulated using the `make_model(AtomicModel)` function, which leverages the `deepmd.pt.model.model.make_model.make_model` function. The purpose of the `make_model` wrapper is to facilitate the translation between atomic property predictions and the extended property predictions and differentiation , e.g. the reduction of atomic energy contribution and the autodiff for calculating the forces and virial. The developers usually need to implement an `AtomicModel` not a `Model`. @@ -217,6 +217,6 @@ When implementing an existing model in a new backend, directly apply the existin ### Consistent tests -When transferring features from another backend to the PyTorch backend, it is essential to include a regression test in `/source/tests/consistent` to validate the consistency of the PyTorch backend with other backends. Presently, the regression tests cover self-consistency and cross-backend consistency between TensorFlow, PyTorch, and DP (NumPy) through the serialization/deserialization technique. +When transferring features from another backend to the PyTorch-TorchScript backend, it is essential to include a regression test in `/source/tests/consistent` to validate the consistency of the PyTorch-TorchScript backend with other backends. Presently, the regression tests cover self-consistency and cross-backend consistency between TensorFlow, PyTorch-TorchScript, and DP (NumPy) through the serialization/deserialization technique. -During the development of new components within the PyTorch backend, it is necessary to provide a DP (NumPy) implementation and incorporate corresponding regression tests. For PyTorch components, developers are also required to include a unit test using `torch.jit`. +During the development of new components within the PyTorch-TorchScript backend, it is necessary to provide a DP (NumPy) implementation and incorporate corresponding regression tests. For PyTorch components, developers are also required to include a unit test using `torch.jit`. diff --git a/doc/freeze/compress.md b/doc/freeze/compress.md index 7f55e44753..898cc1b5ee 100644 --- a/doc/freeze/compress.md +++ b/doc/freeze/compress.md @@ -2,8 +2,8 @@ > [!NOTE] > **Backends covered below**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX -> {{ jax_icon }}, and DP {{ dpmodel_icon }}. +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable +> {{ pytorch_icon }}, JAX {{ jax_icon }}, and DP {{ dpmodel_icon }}. ## Theory @@ -79,12 +79,11 @@ dp --tf2 compress -i model.ckpt.tf2 -o model-compress.savedmodeltf ``` TensorFlow 2 compression reads a `.tf2` training checkpoint directory or a -checkpoint prefix and writes a compressed `.savedmodeltf` model. For -`se_e2_a`/`se_a`, SavedModel export currently requires -`type_one_side: true`; the normalized default, `false`, is not graph-traceable. +checkpoint prefix and writes a compressed `.savedmodeltf` model. See the +descriptor documentation for model-specific SavedModel export requirements. ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt compress -i model.pth -o model-compress.pth @@ -97,14 +96,10 @@ dp --pt compress -i model.pth -o model-compress.pth dp --pt-expt compress -i dpa1-graph.pt2 -o dpa1-graph-compress.pt2 ``` -This executable compression path is currently limited to graph-lowered -DPA1/`se_atten` strip models that are eligible for the fused CUDA table -operator: `attn_layer: 0`, no excluded type pairs, float32 descriptor -statistics and tables, `neuron[-1] <= 256`, and -`axis_neuron <= min(16, neuron[-1])`. Their `.pt2` export embeds the tabulated -operator. For other models, the command records compressed state in -`model.json` but leaves the exported inference graph uncompressed, so the -resulting `.pte`/`.pt2` is not a deployment compression artifact. +This command produces an executable compressed artifact only for descriptors +that support the PyTorch-Exportable compression path. See the +[DPA-1 model-compression requirements](../model/train-se-atten.md#model-compression) +for the graph-lowered `.pt2` route used in this example. ::: :::{tab-item} JAX {{ jax_icon }} @@ -114,10 +109,8 @@ dp --jax compress -i frozen_model.jax -o compressed_model.jax ``` JAX compression accepts `.jax` and `.hlo` inputs. Use `.jax` for the general, -lossless compressed serialization path. A compressed `.hlo` can be written -only when the descriptor is StableHLO-exportable; compressed `se_e2_a`, `se_a`, -`dpa1`, and `se_atten` descriptors with `type_one_side: false` must remain -`.jax` or be rebuilt with `type_one_side: true`. +lossless compressed serialization path. Descriptor pages document whether a +compressed model can also be exported to StableHLO `.hlo`. ::: :::{tab-item} DP {{ dpmodel_icon }} @@ -132,13 +125,10 @@ DP compression accepts native `.dp` and `.yaml` models. :::: where `-i` gives the original frozen model, `-o` gives the compressed model. -The DP, JAX, and TensorFlow 2 entrypoints share the native-model compression -helpers for resolving the minimum neighbor distance and tabulating the -descriptor's embedding networks. The limited PyTorch-Exportable graph-DPA1 -path implements the same table strides and minimum-neighbor-distance fallback -in its export-specific entrypoint. If the model does not contain a minimum -neighbor distance, pass the training script with `-t` or `--training-script` -so it can be computed from the training data. +The compression entrypoints resolve the minimum neighbor distance and tabulate +supported descriptor embedding networks. If the model does not contain a +minimum neighbor distance, pass the training script with `-t` or +`--training-script` so it can be computed from the training data. Several other command line options can be passed to `dp compress`, which can be checked with @@ -183,7 +173,7 @@ See the documentation of a specific descriptor to see whether it supports model ## Requirements of installation {{ pytorch_icon }} -When compressing models in the PyTorch backend, the customized OP library for the Python interface must be installed when [freezing the model](../freeze/freeze.md). +When compressing models in the PyTorch-TorchScript backend, the customized OP library for the Python interface must be installed when [freezing the model](../freeze/freeze.md). The customized OP library for the Python interface is installed by default when building DeePMD-kit from source; see the [installation guide](../install/install-from-source.md) for details. diff --git a/doc/freeze/freeze.md b/doc/freeze/freeze.md index 69f6c589f3..468625582d 100644 --- a/doc/freeze/freeze.md +++ b/doc/freeze/freeze.md @@ -25,12 +25,12 @@ When `-c` names a checkpoint prefix, the backend also checks the corresponding path with `.tf2` appended, so the example reads `model.ckpt.tf2` and writes the TensorFlow SavedModel to `model.savedmodeltf`. If `-c` is omitted, it defaults to the current directory. For a multi-task checkpoint, select a branch with -`--head CHOSEN_BRANCH`. SavedModel export currently requires graph-traceable -model code; an `se_e2_a`/`se_a` descriptor must set `type_one_side: true` -because the normalized default, `false`, cannot be traced. +`--head CHOSEN_BRANCH`. SavedModel export requires graph-traceable model code; +descriptor-specific export requirements are documented on the corresponding +model pages. ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt freeze -o model.pth diff --git a/doc/inference/python.md b/doc/inference/python.md index 4f6ad41b39..e2ce66dcf0 100644 --- a/doc/inference/python.md +++ b/doc/inference/python.md @@ -35,7 +35,7 @@ where `descriptors` is the descriptor matrix of the system. This can also be don > [!NOTE] > `eval_descriptor` is the descriptor-only interface supported across backends. In -> the PyTorch backend, [`eval_embedding`](embedding.md) additionally returns the +> the PyTorch-TorchScript backend, [`eval_embedding`](embedding.md) additionally returns the > descriptor, per-atom feature, and per-structure feature in a single forward pass. > PyTorch descriptor/embedding APIs accept `dtype="fp32"`, `"fp64"`, or `"native"`; > `eval_descriptor` defaults to `native`, while `eval_embedding` defaults to diff --git a/doc/model/change-bias.md b/doc/model/change-bias.md index 310a21e83f..e0fc3d20b6 100644 --- a/doc/model/change-bias.md +++ b/doc/model/change-bias.md @@ -1,7 +1,7 @@ # Change the model output bias for trained model {{ tensorflow_icon }} {{ pytorch_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} +> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript {{ pytorch_icon }} The output bias of a trained model typically originates from the statistical results of the training dataset. @@ -28,7 +28,7 @@ dp --tf change-bias model.ckpt -b -92.523 -187.66 -o model_updated.pb ``` ::: -:::{tab-item} PyTorch Backend {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript Backend {{ pytorch_icon }} **Changing bias using provided systems for trained `.pt`/`.pth` models:** diff --git a/doc/model/dpa2.md b/doc/model/dpa2.md index c6ef07cb28..c96a9602a4 100644 --- a/doc/model/dpa2.md +++ b/doc/model/dpa2.md @@ -1,7 +1,7 @@ # Descriptor DPA-2 {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and +> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch-TorchScript and > PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle > {{ paddle_icon }}, DP {{ dpmodel_icon }} @@ -104,7 +104,7 @@ dp --pt-expt freeze -o model.pt2 --lower-kind graph As with DPA-1's graph path (see [Difference among different backends](train-se-atten.md#difference-among-different-backends)), the graph route considers all neighbors within the cutoff rather than a fixed, padded selection, so its numeric result can differ slightly (down to the AOTInductor floating-point noise floor at non-binding `sel`, larger if `sel` is binding) from the dense/`nlist` path. > [!NOTE] -> **Default route change in pt_expt (eager & training).** For a graph-eligible DPA-2 descriptor, the pt_expt backend now defaults to the carry-all graph route not only for `--lower-kind graph` freezing but also in **eager inference/evaluation and in (compiled) training** (`neighbor_graph_method=None` resolves to the graph). This changes the numerical behavior of existing pt_expt configurations relative to the dense neighbor-list route (by the amounts described above — negligible at non-binding `sel`). The other backends (dpmodel/PyTorch/Paddle/TensorFlow/JAX) are unaffected: they keep the dense route as their only path. +> **Default route change in pt_expt (eager & training).** For a graph-eligible DPA-2 descriptor, the pt_expt backend now defaults to the carry-all graph route not only for `--lower-kind graph` freezing but also in **eager inference/evaluation and in (compiled) training** (`neighbor_graph_method=None` resolves to the graph). This changes the numerical behavior of existing pt_expt configurations relative to the dense neighbor-list route (by the amounts described above — negligible at non-binding `sel`). The other backends (dpmodel/PyTorch-TorchScript/Paddle/TensorFlow/JAX) are unaffected: they keep the dense route as their only path. > > To retain the legacy dense route on pt_expt: > diff --git a/doc/model/dpa3.md b/doc/model/dpa3.md index 65d52f7c46..40d138e074 100644 --- a/doc/model/dpa3.md +++ b/doc/model/dpa3.md @@ -1,7 +1,7 @@ # Descriptor DPA3 {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and +> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch-TorchScript and > PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle > {{ paddle_icon }}, DP {{ dpmodel_icon }} @@ -124,7 +124,7 @@ Note that all DPA3 models use `float32`, while other models use `float64` by def ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} To run the DPA3 model on LAMMPS via source code installation (users can skip this step if using [easy installation](../install/easy-install.md)), diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index c27aa048c9..9a26349bf2 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -1,10 +1,10 @@ # Descriptor DPA4 {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Descriptor backends**: PyTorch and PyTorch-Exportable +> **Descriptor backends**: PyTorch-TorchScript and PyTorch-Exportable > {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} > -> The end-to-end `model.type: dpa4` scaffold is supported by PyTorch and +> The end-to-end `model.type: dpa4` scaffold is supported by PyTorch-TorchScript and > PyTorch-Exportable only. JAX support is descriptor-only and does not provide > the dedicated DPA4 fitting/model/trainer workflow documented on this page. @@ -22,7 +22,8 @@ Reference: [DPA4 paper](https://arxiv.org/abs/2606.02419). ## Quick start -DPA4 is a PyTorch-only model. Train it with the standard `dp --pt` workflow: +DPA4 is available only through the PyTorch framework backends. Train a +PyTorch-TorchScript model with the standard `dp --pt` workflow: ```bash cd examples/water/dpa4 @@ -413,14 +414,14 @@ once the target system and batch size fit. ### Freeze to `.pt2` -DPA4/SeZM checkpoints use the PyTorch `.pt2` (AOTInductor) export path; the +DPA4/SeZM checkpoints use the PyTorch-TorchScript `.pt2` (AOTInductor) export path; the ordinary TorchScript freeze path is not used. Run the standard freeze command: ```bash dp --pt freeze -c model.ckpt.pt -o frozen_model ``` -The PyTorch backend detects DPA4/SeZM and writes `frozen_model.pt2`. +The PyTorch-TorchScript backend detects DPA4/SeZM and writes `frozen_model.pt2`. ### Single GPU @@ -711,7 +712,7 @@ closed over the one-hop neighbor shell. ## Limitations - The end-to-end DPA4/SeZM model and training workflow is implemented for the - PyTorch and PyTorch-Exportable backends. JAX provides the descriptor only; + PyTorch-TorchScript and PyTorch-Exportable backends. JAX provides the descriptor only; the DP implementation is a reference component and does not train models. - Export uses `.pt2` (AOTInductor); the TorchScript freeze path is not used. - Model compression is not supported. diff --git a/doc/model/dprc.md b/doc/model/dprc.md index e604271576..62a96ba924 100644 --- a/doc/model/dprc.md +++ b/doc/model/dprc.md @@ -2,7 +2,7 @@ > [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX > {{ jax_icon }}, DP {{ dpmodel_icon }} Deep Potential - Range Correction (DPRc) is designed to combine with QM/MM method, and corrects energies from a low-level QM/MM method to a high-level QM/MM method: @@ -101,7 +101,7 @@ As described in the paper, the DPRc model only corrects $E_\text{QM}$ and $E_\te ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```json "descriptor" :{ diff --git a/doc/model/linear.md b/doc/model/linear.md index 632df3b9c9..61607f20a7 100644 --- a/doc/model/linear.md +++ b/doc/model/linear.md @@ -1,7 +1,7 @@ ## Linear model {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch and +> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript and > PyTorch-Exportable {{ pytorch_icon }}, DP {{ dpmodel_icon }} One can linearly combine existing models with arbitrary coefficients: diff --git a/doc/model/overall.md b/doc/model/overall.md index 731b5dd031..ffef009fda 100644 --- a/doc/model/overall.md +++ b/doc/model/overall.md @@ -58,7 +58,7 @@ 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] -> The PyTorch TorchScript freeze route cannot output Hessians, and +> The PyTorch-TorchScript freeze route cannot output Hessians, and > PyTorch-Exportable cannot construct a Hessian model for freezing or > inference. The JAX backend can retain Hessian output in a frozen model with > `dp --jax freeze --hessian`; see [Freeze a model](../freeze/freeze.md). diff --git a/doc/model/pairtab.md b/doc/model/pairtab.md index 0d94efdb27..f6045844ad 100644 --- a/doc/model/pairtab.md +++ b/doc/model/pairtab.md @@ -2,10 +2,10 @@ > [!NOTE] > **Short-range interpolation backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX > {{ jax_icon }}, DP {{ dpmodel_icon }} > -> **Linear-combination backends**: TensorFlow {{ tensorflow_icon }}, PyTorch and +> **Linear-combination backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript and > PyTorch-Exportable {{ pytorch_icon }}, DP {{ dpmodel_icon }} ## Theory diff --git a/doc/model/sel.md b/doc/model/sel.md index fcdff92f87..36fcd34c45 100644 --- a/doc/model/sel.md +++ b/doc/model/sel.md @@ -15,7 +15,7 @@ dp --tf neighbor-stat -s data -r 6.0 -t O H ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```sh dp --pt neighbor-stat -s data -r 6.0 -t O H diff --git a/doc/model/train-energy-hessian.md b/doc/model/train-energy-hessian.md index 3fd90a92b3..f190b1da33 100644 --- a/doc/model/train-energy-hessian.md +++ b/doc/model/train-energy-hessian.md @@ -1,7 +1,7 @@ # Fit energy Hessian {{ pytorch_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }} +> **Supported backends**: PyTorch-TorchScript {{ pytorch_icon }} To train a model that takes Hessian matrices, i.e., the second order derivatives of energies w.r.t coordinates as input, you only need to prepare full Hessian matrices and modify the `loss` section to define the Hessian-specific settings, keeping other sections the same as the normal energy model's input script. @@ -27,9 +27,9 @@ The options {ref}`start_pref_e `, {ref}`limit_pref If one does not want to train with virial, then he/she may set the virial prefactors {ref}`start_pref_v ` and {ref}`limit_pref_v ` to 0. -## Hessian Format in PyTorch +## Hessian Format in PyTorch-TorchScript -In the PyTorch backend, Hessian matrices are listed in `hessian.npy` files, and the data format may contain the following files: +In the PyTorch-TorchScript backend, Hessian matrices are listed in `hessian.npy` files, and the data format may contain the following files: ``` type.raw @@ -50,11 +50,11 @@ Note that the `hessian.npy` should contain the **full** Hessian matrices with sh ## Train the Model -There are two approaches to training a Hessian model. The first method involves training the model from scratch using the same command as in the `ener` mode within the PyTorch backend: +There are two approaches to training a Hessian model. The first method involves training the model from scratch using the same command as in the `ener` mode within the PyTorch-TorchScript backend: ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input.json @@ -63,11 +63,11 @@ dp --pt train input.json :::: -The second approach is to train a Hessian model from a pretrained energy model, following the same command as the `finetune` strategy within the PyTorch backend: +The second approach is to train a Hessian model from a pretrained energy model, following the same command as the `finetune` strategy within the PyTorch-TorchScript backend: ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input.json --finetune pretrained_energy.pt @@ -91,17 +91,17 @@ The detailed loss can be found in `lcurve.out`: ## Test the Model > [!WARNING] -> The PyTorch TorchScript freeze route does not preserve Hessian output. A -> PyTorch model frozen with `dp --pt freeze` is treated as a standard energy +> The PyTorch-TorchScript freeze route does not preserve Hessian output. A +> PyTorch-TorchScript model frozen with `dp --pt freeze` is treated as a standard energy > model. PyTorch-Exportable cannot construct a Hessian model for freezing or > inference. The JAX backend can preserve Hessian output in a frozen model > with `dp --jax freeze --hessian`. -If one freezes and tests a Hessian model through the PyTorch route: +If one freezes and tests a Hessian model through the PyTorch-TorchScript route: ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash @@ -120,12 +120,12 @@ ${output_prefix}.e.out ${output_prefix}.e_peratom.out ${output_prefix}.f.out ${output_prefix}.v.out ${output_prefix}.v_peratom.out ``` -For PyTorch Hessian predictions, test the training checkpoint directly without +For PyTorch-TorchScript Hessian predictions, test the training checkpoint directly without freezing: ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash diff --git a/doc/model/train-energy-spin.md b/doc/model/train-energy-spin.md index a80bf9a932..3347a49824 100644 --- a/doc/model/train-energy-spin.md +++ b/doc/model/train-energy-spin.md @@ -1,7 +1,7 @@ # Fit spin energy {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch and +> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript and > PyTorch-Exportable {{ pytorch_icon }}, DP {{ dpmodel_icon }} To train a model that takes additional spin information as input, you only need to modify the following sections to define the spin-specific settings, @@ -12,7 +12,7 @@ keeping other sections the same as the normal energy model's input script. > > - In the TensorFlow backend, the `se_e2_a` descriptor will treat those atom types with spin as new (virtual) types, > and duplicate their corresponding selected numbers of neighbors ({ref}`sel `) from their real atom types. -> - In the PyTorch and PyTorch-Exportable backends, if spin settings are added, all the types (with or without spin) will have their virtual types. +> - In the PyTorch-TorchScript and PyTorch-Exportable backends, if spin settings are added, all the types (with or without spin) will have their virtual types. > The `se_e2_a` descriptor will thus double the {ref}`sel ` list, > while in other descriptors with mixed types (such as `dpa1` or `dpa2`), the sel number will not be changed for clarity. > If you are using descriptors with mixed types, to achieve better performance, @@ -24,7 +24,7 @@ The spin settings are given by the {ref}`spin ` section, which sets > [!NOTE] > Note that the construction of spin settings is different between TensorFlow -> and PyTorch/PyTorch-Exportable/DP. +> and PyTorch-TorchScript/PyTorch-Exportable/DP. ### Spin settings in TensorFlow @@ -43,9 +43,9 @@ The implementation in TensorFlow only supports `se_e2_a` descriptor. See example - {ref}`virtual_len ` specifies the distance between virtual atom and the belonging real atom. - {ref}`spin_norm ` gives the magnitude of the magnetic moment for each magnatic atom. -### Spin settings in PyTorch/PyTorch-Exportable/DP +### Spin settings in PyTorch-TorchScript/PyTorch-Exportable/DP -In PyTorch/PyTorch-Exportable/DP, the spin implementation is more flexible and +In PyTorch-TorchScript/PyTorch-Exportable/DP, the spin implementation is more flexible and so far supports the following descriptors: - `se_e2_a` @@ -73,7 +73,7 @@ See `se_e2_a` examples in `$deepmd_source_dir/examples/spin/se_e2_a/input_torch. only used when {ref}`use_spin ` is True for each atom type. > [!NOTE] -> It should be noted that the spin models in PyTorch/PyTorch-Exportable/DP are capable of addressing scenarios where the spin approaches zero +> It should be noted that the spin models in PyTorch-TorchScript/PyTorch-Exportable/DP are capable of addressing scenarios where the spin approaches zero > (indicating the virtual atom is in close proximity to the real atom) by adjusting the non-zero > {ref}`env_protection ` parameter within the descriptor. > This parameter is set to 0.01 by default in the spin model. It appears that a value of 0.01 is generally sufficient for maintaining model stability. @@ -147,7 +147,7 @@ If one does not want to train with virial, then he/she may set the virial prefac > [!NOTE] > Note that the spin data format is different between TensorFlow and -> PyTorch/PyTorch-Exportable/DP. +> PyTorch-TorchScript/PyTorch-Exportable/DP. ### Spin data format in TensorFlow @@ -179,9 +179,9 @@ We list the details about spin system data format in TensorFlow backend: | force | Atomic and magnetic forces | force.raw | eV/Å | Nframes * (Natoms + Nspins) * 3 | The first `3 \* Natoms` columns represent atomic forces, followed by `3 \* Nspins` columns representing magnetic forces. | | virial | Frame virial | virial.raw | eV | Nframes * 9 | in the order `XX XY XZ YX YY YZ ZX ZY ZZ` | -### Spin data format in PyTorch/PyTorch-Exportable/DP +### Spin data format in PyTorch-TorchScript/PyTorch-Exportable/DP -In the PyTorch, PyTorch-Exportable, and DP backends, spin and magnetic forces +In the PyTorch-TorchScript, PyTorch-Exportable, and DP backends, spin and magnetic forces are listed in separate files, and the data format may contain the following files: diff --git a/doc/model/train-energy.md b/doc/model/train-energy.md index 447abc3414..8a68885e9e 100644 --- a/doc/model/train-energy.md +++ b/doc/model/train-energy.md @@ -2,7 +2,7 @@ > [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX > {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} In this section, we will take `$deepmd_source_dir/examples/water/se_e2_a/input.json` as an example of the input file. @@ -154,7 +154,7 @@ If one does not want to train with virial, then he/she may set the virial prefac ### Prefactor force loss with default atom preference > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: PyTorch-TorchScript {{ pytorch_icon }}, DP {{ dpmodel_icon }} When using the prefactor force loss (controlled by {ref}`start_pref_pf ` and {ref}`limit_pref_pf `), the training data typically requires an `atom_pref.npy` file in each system directory to specify per-atom prefactors $q_k$. If `atom_pref.npy` is not provided, the {ref}`use_default_pf ` option can be set to `true` to use a default atom preference of 1.0 for all atoms: @@ -174,6 +174,6 @@ When using the prefactor force loss (controlled by {ref}`start_pref_pf [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ tensorflow_icon }}, PyTorch-TorchScript {{ pytorch_icon }}, JAX > {{ jax_icon }}, DP {{ dpmodel_icon }} > -> These icons report fitting/model implementation availability. JAX provides -> the DOS implementation for model conversion and inference, but +> JAX provides the DOS implementation for model conversion and inference, but > `dp --jax train` currently accepts only the energy loss and cannot train a DOS > model. TensorFlow 2 can train a DOS model, but a frozen `.savedmodeltf` DOS > model cannot yet be evaluated by `dp test`/`DeepDOS` because the export does -> not preserve the `numb_dos` output width. +> not preserve the `numb_dos` output width. PyTorch-Exportable is not listed +> because its DOS export and inference have the same missing-width problem, +> tracked in +> [#5949](https://github.com/deepmodeling/deepmd-kit/issues/5949). Here we present an API to DeepDOS model, which can be used to fit electronic density of state (DOS) (which is a vector). @@ -106,20 +108,13 @@ dp --tf2 train input.json ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input.json ``` ::: -:::{tab-item} PyTorch-Exportable {{ pytorch_icon }} - -```bash -dp --pt-expt train input.json -``` -::: - :::: The detailed loss can be found in `lcurve.out`: @@ -165,7 +160,7 @@ dp --tf test -m frozen_model.pb -s ../data/111/$k -d ${output_prefix} -a -n 100 ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash diff --git a/doc/model/train-fitting-population.md b/doc/model/train-fitting-population.md index e00bfbc9bb..90197cb313 100644 --- a/doc/model/train-fitting-population.md +++ b/doc/model/train-fitting-population.md @@ -1,7 +1,7 @@ # Fit atomic charge population {{ pytorch_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }} +> **Supported backends**: PyTorch-TorchScript {{ pytorch_icon }} Here we present an API to DeepPopulation model, which can be used to fit the atomic charge population. @@ -89,7 +89,7 @@ The training command is the same as `ener` mode, i.e. ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input_torch.json diff --git a/doc/model/train-fitting-property.md b/doc/model/train-fitting-property.md index c93ae61795..7da02ee9fc 100644 --- a/doc/model/train-fitting-property.md +++ b/doc/model/train-fitting-property.md @@ -1,7 +1,7 @@ # Fit other properties {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and +> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch-TorchScript and > PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, DP > {{ dpmodel_icon }} > @@ -126,7 +126,7 @@ The training command is the same as `ener` mode, i.e. ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input.json @@ -159,7 +159,7 @@ We can use `dp test` to infer the properties for given frames. ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash diff --git a/doc/model/train-fitting-tensor.md b/doc/model/train-fitting-tensor.md index 114c5a9f96..6136cf839c 100644 --- a/doc/model/train-fitting-tensor.md +++ b/doc/model/train-fitting-tensor.md @@ -2,7 +2,7 @@ > [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX > {{ jax_icon }}, DP {{ dpmodel_icon }} > > These icons report fitting/model implementation availability. JAX provides @@ -22,7 +22,7 @@ $deepmd_source_dir/examples/water_tensor/polar/polar_input.json ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash $deepmd_source_dir/examples/water_tensor/dipole/dipole_input_torch.json @@ -110,7 +110,7 @@ The JSON of `polar` type should be provided like - The rest arguments have the same meaning as they do in `ener` mode. ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} The JSON of `dipole` type should be provided like @@ -193,7 +193,7 @@ dp train input.json ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input.json diff --git a/doc/model/train-hybrid.md b/doc/model/train-hybrid.md index 72c487f050..dc9d2d9594 100644 --- a/doc/model/train-hybrid.md +++ b/doc/model/train-hybrid.md @@ -2,7 +2,7 @@ > [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX > {{ jax_icon }}, DP {{ dpmodel_icon }} This descriptor hybridizes multiple descriptors to form a new descriptor. For example, we have a list of descriptors denoted by $\mathcal D_1$, $\mathcal D_2$, ..., $\mathcal D_N$, the hybrid descriptor this the concatenation of the list, i.e. $\mathcal D = (\mathcal D_1, \mathcal D_2, \cdots, \mathcal D_N)$. diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md index db0ab80fdc..eaf02e38c5 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -2,8 +2,9 @@ > [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX -> {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable +> {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, DP +> {{ dpmodel_icon }} ![ALT](../images/model_se_atten.png "model_se_atten") @@ -192,10 +193,33 @@ DPA-1 supports both the [standard data format](../data/system.md) and the [mixed Model compression is supported only when the descriptor attention depth {ref}`attn_layer ` is 0 and {ref}`tebd_input_mode ` is `"strip"`. Attention layers higher than 0 cannot be compressed in the TensorFlow implementation because the geometric part is tabulated from the static computation graph. -### PyTorch {{ pytorch_icon }} +### PyTorch-TorchScript {{ pytorch_icon }} Model compression is supported for any {ref}`attn_layer ` value when {ref}`tebd_input_mode ` is `"strip"`. When `attn_layer` is 0, both the type embedding and geometric parts are compressed. When `attn_layer` is not 0, only the type embedding is compressed while the geometric part keeps the neural network implementation (a warning is emitted during compression). +### PyTorch-Exportable {{ pytorch_icon }} + +Executable compression is limited to graph-lowered DPA-1 models with +{ref}`attn_layer ` set to `0` +and {ref}`tebd_input_mode ` +set to `"strip"`. The fused CUDA table operator also requires no excluded type +pairs, float32 descriptor statistics and tables, `neuron[-1] <= 256`, and +`axis_neuron <= min(16, neuron[-1])`. A matching `.pt2` export embeds the +tabulated operator. + +For other models, `dp --pt-expt compress` records compressed state in +`model.json` but leaves the executable inference graph unchanged, so the +resulting `.pte` or `.pt2` file is not a deployment compression artifact. The +export-specific entrypoint uses the same table strides and +minimum-neighbor-distance fallback as the native-model compression helpers. + +### JAX {{ jax_icon }} + +Use `.jax` for the general lossless compressed serialization path. A compressed +DPA-1 model can be exported to StableHLO `.hlo` only when +{ref}`type_one_side ` is +`true`. + ## Training example Here we upload the AlMgCu example shown in the paper, you can download it here: diff --git a/doc/model/train-se-e2-a-tebd.md b/doc/model/train-se-e2-a-tebd.md index 8678a994ec..2a90f1e53d 100644 --- a/doc/model/train-se-e2-a-tebd.md +++ b/doc/model/train-se-e2-a-tebd.md @@ -2,7 +2,7 @@ > [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX > {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} We generate specific a type embedding vector for each atom type so that we can share one descriptor embedding net and one fitting net in total, which decline training complexity largely. diff --git a/doc/model/train-se-e2-a.md b/doc/model/train-se-e2-a.md index 30c7210a56..d833972cb3 100644 --- a/doc/model/train-se-e2-a.md +++ b/doc/model/train-se-e2-a.md @@ -2,8 +2,9 @@ > [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX -> {{ jax_icon }}, Paddle {{ paddle_icon }}, DP {{ dpmodel_icon }} +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable +> {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, DP +> {{ dpmodel_icon }} The notation of `se_e2_a` is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from all information (both angular and radial) of atomic configurations. The `e2` stands for the embedding with two-atoms information. This descriptor was described in detail in [the DeepPot-SE paper](https://arxiv.org/abs/1805.09003). @@ -102,11 +103,19 @@ Type embdding is only supported in the TensorFlow backends. ## Difference among different backends In the TensorFlow backend, {ref}`env_protection ` cannot be set to a non-zero value. +In the TensorFlow 2 backend, freezing or compressing to `.savedmodeltf` +currently requires +{ref}`type_one_side ` to be +`true`; the normalized default, `false`, is not graph-traceable during +SavedModel export. In the JAX backend, {ref}`type_one_side ` cannot be set to `false`. ## Model compression Model compression is supported when type embedding is not used. To use model compression with type embedding in the TensorFlow backend, use `se_a_tebd_v2` instead. +For JAX, `.jax` is the general lossless compressed serialization format; +StableHLO `.hlo` export additionally requires the compressed descriptor to be +StableHLO-exportable. [^1]: This section is built upon Jinzhe Zeng, Duo Zhang, Denghui Lu, Pinghui Mo, Zeyu Li, Yixiao Chen, Marián Rynik, Li'ang Huang, Ziyao Li, Shaochen Shi, Yingze Wang, Haotian Ye, Ping Tuo, Jiabin Yang, Ye Ding, Yifan Li, Davide Tisi, Qiyu Zeng, Han Bao, Yu Xia, Jiameng Huang, Koki Muraoka, Yibo Wang, Junhan Chang, Fengbo Yuan, Sigbjørn Løland Bore, Chun Cai, Yinnian Lin, Bo Wang, Jiayan Xu, Jia-Xin Zhu, Chenxing Luo, Yuzhi Zhang, Rhys E. A. Goodall, Wenshuo Liang, Anurag Kumar Singh, Sikai Yao, Jingchao Zhang, Renata Wentzcovitch, Jiequn Han, Jie Liu, Weile Jia, Darrin M. York, Weinan E, Roberto Car, Linfeng Zhang, Han Wang, [J. Chem. Phys. 159, 054801 (2023)](https://doi.org/10.1063/5.0155600) licensed under a [Creative Commons Attribution (CC BY) license](http://creativecommons.org/licenses/by/4.0/). diff --git a/doc/model/train-se-e2-r.md b/doc/model/train-se-e2-r.md index 2f2616937c..5ca7402dab 100644 --- a/doc/model/train-se-e2-r.md +++ b/doc/model/train-se-e2-r.md @@ -2,7 +2,7 @@ > [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX > {{ jax_icon }}, DP {{ dpmodel_icon }} The notation of `se_e2_r` is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from the radial information of atomic configurations. The `e2` stands for the embedding with two-atom information. @@ -76,7 +76,7 @@ Type embdding is only supported in the TensorFlow backends. ## Difference among different backends In the TensorFlow backend, {ref}`env_protection ` cannot be set to a non-zero value. -In the PyTorch, JAX, and DP backend, {ref}`type_one_side ` cannot be set to `false`. +In the PyTorch-TorchScript, JAX, and DP backends, {ref}`type_one_side ` cannot be set to `false`. ## Model compression diff --git a/doc/model/train-se-e3-tebd.md b/doc/model/train-se-e3-tebd.md index e0a010245f..a3ba3746b3 100644 --- a/doc/model/train-se-e3-tebd.md +++ b/doc/model/train-se-e3-tebd.md @@ -1,7 +1,7 @@ # Descriptor `"se_e3_tebd"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch and +> **Supported backends**: TensorFlow 2 {{ tensorflow_icon }}, PyTorch-TorchScript and > PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle > {{ paddle_icon }}, DP {{ dpmodel_icon }} diff --git a/doc/model/train-se-e3.md b/doc/model/train-se-e3.md index 5425900c1f..cdd29d05d3 100644 --- a/doc/model/train-se-e3.md +++ b/doc/model/train-se-e3.md @@ -2,7 +2,7 @@ > [!NOTE] > **Supported backends**: TensorFlow and TensorFlow 2 -> {{ tensorflow_icon }}, PyTorch and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX > {{ jax_icon }}, DP {{ dpmodel_icon }} The notation of `se_e3` is short for three-body embedding DeepPot-SE, which incorporates embedded bond-angle information. diff --git a/doc/third-party/ase.md b/doc/third-party/ase.md index c1e9929648..f4c350fca1 100644 --- a/doc/third-party/ase.md +++ b/doc/third-party/ase.md @@ -24,7 +24,7 @@ print(water.get_forces()) ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```python from ase import Atoms diff --git a/doc/train/parallel-training.md b/doc/train/parallel-training.md index 7c032179c7..f158ee0263 100644 --- a/doc/train/parallel-training.md +++ b/doc/train/parallel-training.md @@ -92,14 +92,14 @@ optional arguments: master) ``` -## PyTorch Implementation {{ pytorch_icon }} +## PyTorch-TorchScript Implementation {{ pytorch_icon }} Currently, parallel training in pytorch version is implemented in the form of PyTorch Distributed Data Parallelism [DDP](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html). DeePMD-kit will decide whether to launch the training in parallel (distributed) mode or in serial mode depending on your execution command. ### Optional ZeRO memory optimization -In PyTorch backend, DeePMD-kit supports ZeRO (Zero Redundancy Optimizer) stages +In the PyTorch-TorchScript backend, DeePMD-kit supports ZeRO (Zero Redundancy Optimizer) stages to reduce per-GPU memory usage during distributed training. | `zero_stage` | Strategy | Communication | Memory saving | @@ -137,7 +137,7 @@ Enable it in input config: Constraints: -- Works only in PyTorch backend. +- Works only in the PyTorch-TorchScript backend. - Requires distributed launch with `torchrun`. - Currently single-task only. - Not supported with `LKF` optimizer. diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 0112462581..cfa511e44a 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -103,8 +103,8 @@ Other keys in the {ref}`training ` section are explained below: - {ref}`disp_file ` The file for printing learning curve. - {ref}`disp_freq ` The frequency of printing learning curve. Set in the unit of training steps - {ref}`save_freq ` The frequency of saving checkpoint. -- {ref}`save_dir ` The directory where periodic checkpoints are written (PyTorch backend). It is created recursively if missing, while the `model.ckpt.pt` symlinks and the `checkpoint` pointer file stay in the working directory. Defaults to the working directory. -- {ref}`ckpt_keep_ratio ` An alternative to `max_ckpt_keep` (PyTorch backend) that keeps a sliding window of `ceil(ckpt_keep_ratio * ceil(numb_steps / save_freq))` most recent checkpoints, i.e. the final `ckpt_keep_ratio` fraction of the run by step. It overrides `max_ckpt_keep` (and `ema_ckpt_keep`) when set, and works the same whether the run length is given by `numb_steps` or `numb_epoch`. +- {ref}`save_dir ` The directory where periodic checkpoints are written (PyTorch-TorchScript backend). It is created recursively if missing, while the `model.ckpt.pt` symlinks and the `checkpoint` pointer file stay in the working directory. Defaults to the working directory. +- {ref}`ckpt_keep_ratio ` An alternative to `max_ckpt_keep` (PyTorch-TorchScript backend) that keeps a sliding window of `ceil(ckpt_keep_ratio * ceil(numb_steps / save_freq))` most recent checkpoints, i.e. the final `ckpt_keep_ratio` fraction of the run by step. It overrides `max_ckpt_keep` (and `ema_ckpt_keep`) when set, and works the same whether the run length is given by `numb_steps` or `numb_epoch`. ## Options and environment variables diff --git a/doc/train/training.md b/doc/train/training.md index 58f84b1019..8c75641ddf 100644 --- a/doc/train/training.md +++ b/doc/train/training.md @@ -24,7 +24,7 @@ dp --tf2 train input.json ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input.json @@ -59,7 +59,7 @@ CINN=1 CINN_ALLOW_DYNAMIC_SHAPE=0 dp --pd train input.json dp --jax train input.json ``` -JAX training, introduced after DeePMD-kit v3.1.3, writes checkpoints with the +JAX training, introduced in DeePMD-kit v3.2.0, writes checkpoints with the `.jax` extension. It supports fine-tuning and multi-task configurations, but does not currently support `--init-frz-model`. ::: diff --git a/doc/troubleshooting/howtoset_num_nodes.md b/doc/troubleshooting/howtoset_num_nodes.md index c260ce4e96..580839d1a6 100644 --- a/doc/troubleshooting/howtoset_num_nodes.md +++ b/doc/troubleshooting/howtoset_num_nodes.md @@ -15,9 +15,9 @@ Parallelism for MPI is optional and used for multiple nodes, multiple GPU cards, To enable MPI support for training in the TensorFlow interface, one should [install horovod](../install/install-from-source.md#install-horovod-and-mpi4py) in advance. ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} -Multiprocessing support for training in the PyTorch backend is implemented with [torchrun](https://pytorch.org/docs/stable/elastic/run.html). +Multiprocessing support for training in the PyTorch-TorchScript backend is implemented with [torchrun](https://pytorch.org/docs/stable/elastic/run.html). ::: :::: @@ -101,7 +101,7 @@ dp --tf train input.json ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash export OMP_NUM_THREADS=3 From 30858ac77c347b3095da3d1f9eec10501638818d Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 23:41:06 +0800 Subject: [PATCH 08/10] docs(dpa4): do not describe DPA4/SeZM as a TorchScript model DPA4/SeZM is realized in the PyTorch backend but exports through the AOTInductor .pt2 path, not TorchScript. The backend rename to PyTorch-TorchScript was mechanically applied to this page earlier; restore accurate wording in the quick-start and freeze-to-.pt2 sections. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max --- doc/model/dpa4.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 9a26349bf2..b153523c88 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -22,8 +22,10 @@ Reference: [DPA4 paper](https://arxiv.org/abs/2606.02419). ## Quick start -DPA4 is available only through the PyTorch framework backends. Train a -PyTorch-TorchScript model with the standard `dp --pt` workflow: +DPA4 is available only through the PyTorch framework. Train the model with the +standard `dp --pt` workflow. DPA4/SeZM does not use TorchScript: freezing +exports through the AOTInductor `.pt2` path instead (see +[Freeze to `.pt2`](#freeze-to-pt2)): ```bash cd examples/water/dpa4 @@ -414,14 +416,15 @@ once the target system and batch size fit. ### Freeze to `.pt2` -DPA4/SeZM checkpoints use the PyTorch-TorchScript `.pt2` (AOTInductor) export path; the -ordinary TorchScript freeze path is not used. Run the standard freeze command: +DPA4/SeZM checkpoints are exported through the AOTInductor `.pt2` path of the +PyTorch backend; the ordinary TorchScript freeze path is not used. Run the +standard freeze command: ```bash dp --pt freeze -c model.ckpt.pt -o frozen_model ``` -The PyTorch-TorchScript backend detects DPA4/SeZM and writes `frozen_model.pt2`. +The PyTorch backend detects DPA4/SeZM and writes `frozen_model.pt2`. ### Single GPU From 708891ede9d33de492f6a25ec2e8cfc542963dd1 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 4 Aug 2026 11:16:08 +0800 Subject: [PATCH 09/10] docs(backend): address remaining review findings Reply and close the current review round on #5929: - backend.md: describe the DPA4/SeZM `dp --pt freeze` export as the separate AOTInductor route instead of a PyTorch-TorchScript route, so the overview agrees with dpa4.md. - change-bias.md: document the PyTorch-Exportable `dp --pt-expt change-bias` route for .pt/.pte/.pt2 inputs, with the multi-task limitation noted. - parallel-training.md: add the PyTorch-Exportable DDP subsection with a `torchrun ... dp --pt-expt train` example, separate from the PyTorch-TorchScript ZeRO/FSDP2 section. - dpa2.md: drop the unavailable dense (`nlist`, default) fallback advice and state the forced graph lower for graph-eligible DPA-2. - finetuning.md: document the TensorFlow 2 and PyTorch-Exportable `--finetune` workflows and their accepted formats/limitations. - train-energy-spin.md: scope the virtual-atom (sel-extension) guidance to the deepspin scheme, add the DPA4 native exception, and add DPA4 to the scheme-aware supported-descriptor list. Coding-Agent: opencode opencode-Version: 1.18.11 Model: ustc/deepseek-v4-flash Reasoning-Effort: max --- doc/backend.md | 2 +- doc/model/change-bias.md | 21 +++++++++++++++++- doc/model/dpa2.md | 2 +- doc/model/train-energy-spin.md | 16 +++++++++++++- doc/train/finetuning.md | 39 +++++++++++++++++++++++++++++++++- doc/train/parallel-training.md | 24 ++++++++++++++++++++- 6 files changed, 98 insertions(+), 6 deletions(-) diff --git a/doc/backend.md b/doc/backend.md index f15b5fb5c0..ea78e58a26 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -84,7 +84,7 @@ uses the dense `nlist` ABI instead. A graph model frozen with PyTorch-Exportable runtime, which reads this metadata to select the correct input path. The `--lower-kind` option controls only the PyTorch-Exportable freeze route; see the [DPA4 export documentation](model/dpa4.md#freeze-to-pt2) -for the separate PyTorch-TorchScript route. +for the separate DPA4/SeZM AOTInductor export route. ### JAX {{ jax_icon }} diff --git a/doc/model/change-bias.md b/doc/model/change-bias.md index e0fc3d20b6..4531b96ee9 100644 --- a/doc/model/change-bias.md +++ b/doc/model/change-bias.md @@ -1,7 +1,7 @@ # Change the model output bias for trained model {{ tensorflow_icon }} {{ pytorch_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript {{ pytorch_icon }} +> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript {{ pytorch_icon }}, PyTorch-Exportable {{ pytorch_icon }} The output bias of a trained model typically originates from the statistical results of the training dataset. @@ -49,4 +49,23 @@ dp --pt change-bias multi_model.pt -s data_dir -o model_updated.pt --model-branc ``` ::: +:::{tab-item} PyTorch-Exportable Backend {{ pytorch_icon }} + +**Changing bias using provided systems for trained `.pt` checkpoints or frozen `.pte`/`.pt2` models:** + +```sh +dp --pt-expt change-bias model.pt -s data_dir -o model_updated.pt +dp --pt-expt change-bias model.pte -s data_dir -o model_updated.pte +``` + +**Changing bias using user input for energy model:** + +```sh +dp --pt-expt change-bias model.pt -b -92.523 -187.66 -o model_updated.pt +``` + +> [!NOTE] +> Multi-task change-bias is not yet supported in the PyTorch-Exportable backend. +::: + :::: diff --git a/doc/model/dpa2.md b/doc/model/dpa2.md index c96a9602a4..ed7d83c3f7 100644 --- a/doc/model/dpa2.md +++ b/doc/model/dpa2.md @@ -118,4 +118,4 @@ DPA-2's repformer block performs message passing (per-layer neighbor feature agg Per-atom virial on the graph route uses a different (but equally valid) decomposition than the dense path: each edge's full bond-virial contribution is assigned to its source (neighbor) atom, whereas the dense path's autograd-based per-atom decomposition distributes message-passing contributions across atoms differently. For non-message-passing descriptors the two conventions coincide; for DPA-2 they differ elementwise. Only the *total* virial (summed over all atoms) is convention-independent between the two paths; per-atom virial values from the graph and dense routes should not be compared directly. -Multi-rank message-passing inference on the graph route requires every MPI rank to own or ghost at least one atom -- a rank with zero atoms in both categories raises an error rather than silently desynchronizing the collective per-layer ghost exchange. Pick a domain decomposition that keeps every rank non-empty, or use the dense (`nlist`, the default) artifact, which has no such restriction. +Multi-rank message-passing inference on the graph route requires every MPI rank to own or ghost at least one atom -- a rank with zero atoms in both categories raises an error rather than silently desynchronizing the collective per-layer ghost exchange. For a graph-eligible DPA-2 descriptor, `deepmd` freezes through the graph route only: `--lower-kind nlist` is overridden to `graph` for such models, so a dense artifact is not selectable for the same model. Pick a domain decomposition that keeps every rank non-empty, or configure the model so it is not graph-eligible if you need the dense route. diff --git a/doc/model/train-energy-spin.md b/doc/model/train-energy-spin.md index 3347a49824..d06bf951bc 100644 --- a/doc/model/train-energy-spin.md +++ b/doc/model/train-energy-spin.md @@ -12,11 +12,18 @@ keeping other sections the same as the normal energy model's input script. > > - In the TensorFlow backend, the `se_e2_a` descriptor will treat those atom types with spin as new (virtual) types, > and duplicate their corresponding selected numbers of neighbors ({ref}`sel `) from their real atom types. -> - In the PyTorch-TorchScript and PyTorch-Exportable backends, if spin settings are added, all the types (with or without spin) will have their virtual types. +> - In the PyTorch-TorchScript and PyTorch-Exportable backends with the default `deepspin` (virtual-atom) scheme, if spin settings are added, all the types (with or without spin) will have their virtual types. > The `se_e2_a` descriptor will thus double the {ref}`sel ` list, > while in other descriptors with mixed types (such as `dpa1` or `dpa2`), the sel number will not be changed for clarity. > If you are using descriptors with mixed types, to achieve better performance, > you should manually extend your sel number (maybe double) depending on the balance between performance and efficiency. +> +> The DPA4/SeZM model with {ref}`model.spin.scheme: "native" ` is an exception: +> the native scheme passes the per-atom spin vector directly into the descriptor and +> introduces neither virtual atoms nor a doubled type map. On the pt_expt backend, +> native DPA4 is graph-only, and graph neighbor construction is `sel`-free (every edge +> within the cutoff is retained), so the `sel`-extension advice above applies only to +> the dense virtual-atom routes, not to native DPA4. ## Spin @@ -52,6 +59,7 @@ so far supports the following descriptors: - `dpa1`(`se_atten`) - `dpa2` - `dpa3` +- `dpa4` See `se_e2_a` examples in `$deepmd_source_dir/examples/spin/se_e2_a/input_torch.json`, the {ref}`spin ` section is defined as the following with a much more clear interface: @@ -72,6 +80,12 @@ See `se_e2_a` examples in `$deepmd_source_dir/examples/spin/se_e2_a/input_torch. List of float values with shape of `ntypes` or `ntypes_spin` or one single float value for all types, only used when {ref}`use_spin ` is True for each atom type. +For these descriptors except `dpa4`, spin always uses the default `deepspin` +(virtual-atom) scheme. The `dpa4` descriptor additionally supports the `native` +scheme ({ref}`model.spin.scheme `), which passes the per-atom spin +vector directly into the descriptor without virtual atoms; see the [DPA4 spin +documentation](dpa4.md#spin). + > [!NOTE] > It should be noted that the spin models in PyTorch-TorchScript/PyTorch-Exportable/DP are capable of addressing scenarios where the spin approaches zero > (indicating the virtual atom is in close proximity to the real atom) by adjusting the non-zero diff --git a/doc/train/finetuning.md b/doc/train/finetuning.md index f6d167ac6e..5995c79836 100644 --- a/doc/train/finetuning.md +++ b/doc/train/finetuning.md @@ -1,7 +1,8 @@ # Finetune the pre-trained model {{ tensorflow_icon }} {{ pytorch_icon }} {{ paddle_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }} +> **Supported backends**: TensorFlow {{ tensorflow_icon }}, TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }} Pretraining-and-finetuning is a widely used approach in other fields such as Computer Vision (CV) or Natural Language Processing (NLP) to vastly reduce the training cost, while it's not trivial in potential models. @@ -256,3 +257,39 @@ $ dp --pd train input.json --finetune multitask_pretrained.pd --model-branch CHO This command will start fine-tuning based on the pre-trained model's descriptor and the selected branch's fitting net. If --model-branch is not set or set to "RANDOM", a randomly initialized fitting net will be used. + +## TensorFlow 2 Implementation {{ tensorflow_icon }} + +TensorFlow 2 fine-tuning reuses the same `--finetune` flag as the legacy +TensorFlow backend: + +```bash +dp --tf2 train input.json --finetune pretrained.tf2 +``` + +The pre-trained model is a `.tf2` training checkpoint directory (or its +checkpoint prefix). A frozen `.savedmodeltf` export cannot be fine-tuned, +because it does not carry the structured variable metadata required to restore +the model. The same command-line options as the legacy TensorFlow backend apply: +`--model-branch` selects a branch in a multi-task pre-trained model, and +`--use-pretrain-script` copies the descriptor/fitting parameters from the +pre-trained model. + +## PyTorch-Exportable Implementation {{ pytorch_icon }} + +The PyTorch-Exportable backend implements fine-tuning with `--finetune`, and a +pre-trained model can be either a `.pt` training checkpoint or a frozen `.pte` +/ `.pt2` model: + +```bash +dp --pt-expt train input.json --finetune pretrained.pt +dp --pt-expt train input.json --finetune frozen_model.pte +``` + +Both single-task and multi-task fine-tuning are supported; multi-task +fine-tuning uses the `finetune_head` field in each {ref}`model_dict ` branch. Distributed fine-tuning works under +[`torchrun`](parallel-training.md#pytorch-exportable-implementation) like +regular PyTorch-Exportable training. `--use-pretrain-script` requires the +pre-trained model to embed full model parameters: a `.pte` file frozen with +older code that lacks `model_def_script` must be re-frozen before it can be +used with that option. diff --git a/doc/train/parallel-training.md b/doc/train/parallel-training.md index f158ee0263..5db9850fcc 100644 --- a/doc/train/parallel-training.md +++ b/doc/train/parallel-training.md @@ -1,7 +1,7 @@ # Parallel training {{ tensorflow_icon }} {{ pytorch_icon }} {{ paddle_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }} +> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, Paddle {{ paddle_icon }} ## TensorFlow Implementation {{ tensorflow_icon }} @@ -237,6 +237,28 @@ torchrun --rdzv_endpoint=node0:12321 --nnodes=2 --nproc_per_node=4 --node_rank=1 > To check forward, backward, and communication time, please set env var `TORCH_CPP_LOG_LEVEL=INFO TORCH_DISTRIBUTED_DEBUG=DETAIL`. More details can be found [here](https://pytorch.org/docs/stable/distributed.html#logging). +## PyTorch-Exportable Implementation {{ pytorch_icon }} + +The PyTorch-Exportable backend also trains in parallel with PyTorch +Distributed Data Parallelism [DDP](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html). +It initializes the process group automatically when launched with +[`torchrun`](https://pytorch.org/docs/stable/elastic/run.html#usage) (which +exports the `LOCAL_RANK` environment variable), and wraps both single-task and +multi-task models in `DistributedDataParallel`. + +### How to use + +Launch a DDP training session with `torchrun`, the same as for the +PyTorch-TorchScript backend: + +```bash +torchrun --nproc_per_node=4 --no-python dp --pt-expt train input.json +``` + +Unlike the PyTorch-TorchScript backend, the ZeRO/FSDP2 memory-optimization +options described above are not available here; PyTorch-Exportable uses +standard DDP only. + ## Paddle Implementation {{ paddle_icon }} ### How to use From 23a87ad901cf253611b0f832e97eb2b0ee36be0b Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 4 Aug 2026 11:50:47 +0800 Subject: [PATCH 10/10] docs(backend): fix Copilot-noted format and backend-scope gaps - compress.md: list `.yml` alongside `.dp`/`.yaml` for DP compression, matching load_dp_model's supported extensions. - train-se-e2-r.md: include PyTorch-Exportable in the backends for which `type_one_side=false` is unsupported (pt_expt `se_r` inherits the dpmodel implementation that raises NotImplementedError). - howtoset_num_nodes.md: note that torchrun-based multiprocessing training applies to both PyTorch-TorchScript and PyTorch-Exportable. Coding-Agent: opencode opencode-Version: 1.18.11 Model: ustc/deepseek-v4-flash Reasoning-Effort: max --- doc/freeze/compress.md | 2 +- doc/model/train-se-e2-r.md | 2 +- doc/troubleshooting/howtoset_num_nodes.md | 5 +++-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/doc/freeze/compress.md b/doc/freeze/compress.md index 898cc1b5ee..c56d881ca1 100644 --- a/doc/freeze/compress.md +++ b/doc/freeze/compress.md @@ -119,7 +119,7 @@ compressed model can also be exported to StableHLO `.hlo`. dp --dp compress -i model.dp -o model-compress.dp ``` -DP compression accepts native `.dp` and `.yaml` models. +DP compression accepts native `.dp`, `.yaml`, and `.yml` models. ::: :::: diff --git a/doc/model/train-se-e2-r.md b/doc/model/train-se-e2-r.md index 5ca7402dab..887416c990 100644 --- a/doc/model/train-se-e2-r.md +++ b/doc/model/train-se-e2-r.md @@ -76,7 +76,7 @@ Type embdding is only supported in the TensorFlow backends. ## Difference among different backends In the TensorFlow backend, {ref}`env_protection ` cannot be set to a non-zero value. -In the PyTorch-TorchScript, JAX, and DP backends, {ref}`type_one_side ` cannot be set to `false`. +In the PyTorch-TorchScript, PyTorch-Exportable, JAX, and DP backends, {ref}`type_one_side ` cannot be set to `false`. ## Model compression diff --git a/doc/troubleshooting/howtoset_num_nodes.md b/doc/troubleshooting/howtoset_num_nodes.md index 580839d1a6..52e321ec71 100644 --- a/doc/troubleshooting/howtoset_num_nodes.md +++ b/doc/troubleshooting/howtoset_num_nodes.md @@ -15,9 +15,10 @@ Parallelism for MPI is optional and used for multiple nodes, multiple GPU cards, To enable MPI support for training in the TensorFlow interface, one should [install horovod](../install/install-from-source.md#install-horovod-and-mpi4py) in advance. ::: -:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }} -Multiprocessing support for training in the PyTorch-TorchScript backend is implemented with [torchrun](https://pytorch.org/docs/stable/elastic/run.html). +Multiprocessing support for training in the PyTorch-TorchScript and +PyTorch-Exportable backends is implemented with [torchrun](https://pytorch.org/docs/stable/elastic/run.html). ::: ::::