diff --git a/doc/backend.md b/doc/backend.md index 3e5d634ac7..ea78e58a26 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -2,10 +2,16 @@ ## Supported backends -DeePMD-kit supports multiple backends: TensorFlow and PyTorch. +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 {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, and Paddle {{ paddle_icon }} icons are used to mark whether a backend supports a feature. +In the documentation, TensorFlow and TensorFlow 2 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. ### TensorFlow {{ tensorflow_icon }} @@ -15,25 +21,88 @@ 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. -### PyTorch {{ pytorch_icon }} +### 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 use a +`.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`. + +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-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 }} + +- 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, 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 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 +`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`). 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) +for the separate DPA4/SeZM AOTInductor export route. ### JAX {{ jax_icon }} -- Model filename extension: `.xlo`, `.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 `.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. 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 }} - Model filename extensions: `.json` and `.pdiparams` @@ -64,12 +133,19 @@ 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` 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`, dense `nlist`, or NeighborGraph ABI. ## Convert model files between backends 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 1cf643a627..c56d881ca1 100644 --- a/doc/freeze/compress.md +++ b/doc/freeze/compress.md @@ -1,7 +1,9 @@ -# Compress a model {{ tensorflow_icon }} {{ pytorch_icon }} +# Compress a model {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} +> **Backends covered below**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable +> {{ pytorch_icon }}, JAX {{ jax_icon }}, and DP {{ dpmodel_icon }}. ## Theory @@ -70,19 +72,68 @@ dp compress -i graph.pb -o graph-compress.pb ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{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. See the +descriptor documentation for model-specific SavedModel export requirements. +::: + +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt compress -i model.pth -o model-compress.pth ``` ::: +:::{tab-item} PyTorch-Exportable {{ pytorch_icon }} + +```bash +dp --pt-expt compress -i dpa1-graph.pt2 -o dpa1-graph-compress.pt2 +``` + +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 }} + +```bash +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. Descriptor pages document whether a +compressed model can also be exported to StableHLO `.hlo`. +::: + +:::{tab-item} DP {{ dpmodel_icon }} + +```bash +dp --dp compress -i model.dp -o model-compress.dp +``` + +DP compression accepts native `.dp`, `.yaml`, and `.yml` 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 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 ```bash -$ dp compress --help +dp compress --help ``` An explanation will be provided @@ -122,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 3aef7e22a6..468625582d 100644 --- a/doc/freeze/freeze.md +++ b/doc/freeze/freeze.md @@ -8,17 +8,32 @@ 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`. 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} PyTorch {{ pytorch_icon }} +:::{tab-item} TensorFlow 2 {{ tensorflow_icon }} ```bash -$ dp --pt freeze -o model.pth +dp --tf2 freeze -c model.ckpt -o model.savedmodeltf +``` + +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 requires graph-traceable model code; +descriptor-specific export requirements are documented on the corresponding +model pages. +::: + +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} + +```bash +dp --pt freeze -o model.pth ``` in the folder where the model is trained. The output model is called `model.pth`. @@ -27,16 +42,30 @@ 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. ::: +:::{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. `--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`. @@ -45,10 +74,36 @@ 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. ::: +:::{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` +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. +::: + :::: + +## 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/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..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 {{ 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. @@ -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:** @@ -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 d5d2b9c91c..ed7d83c3f7 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-TorchScript 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,13 +98,13 @@ 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. > [!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: > @@ -116,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/dpa3.md b/doc/model/dpa3.md index f52cfbdc1f..40d138e074 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-TorchScript 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, @@ -122,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 66ff3363da..b153523c88 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-TorchScript and PyTorch-Exportable +> {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> +> 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. DPA4 is the DeePMD-kit implementation of the SeZM (Smooth Equivariant Zone-bridging Model) architecture: an SO(3)-equivariant message-passing model @@ -17,7 +22,10 @@ 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. 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 @@ -408,8 +416,9 @@ once the target system and batch size fit. ### Freeze to `.pt2` -DPA4/SeZM checkpoints use the PyTorch `.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 @@ -444,7 +453,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 +516,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 @@ -536,22 +545,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) @@ -565,7 +569,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 @@ -573,17 +577,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 @@ -594,19 +596,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 @@ -720,7 +714,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-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. - 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..62a96ba924 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-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: @@ -99,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 275e08c568..61607f20a7 100644 --- a/doc/model/linear.md +++ b/doc/model/linear.md @@ -1,7 +1,8 @@ -## Linear model {{ tensorflow_icon }} {{ pytorch_icon }} +## Linear model {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} > [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} +> **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 37470f44cb..ffef009fda 100644 --- a/doc/model/overall.md +++ b/doc/model/overall.md @@ -58,7 +58,10 @@ 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, 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. 1. [`polar`](train-fitting-tensor.md): The polarizability. diff --git a/doc/model/pairtab.md b/doc/model/pairtab.md index 3cb6cf12f3..f6045844ad 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 }} {{ 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-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX +> {{ jax_icon }}, DP {{ dpmodel_icon }} +> +> **Linear-combination backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript and +> PyTorch-Exportable {{ pytorch_icon }}, DP {{ dpmodel_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 }} {{ dpmodel_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 }} {{ dpmodel_icon }} To combine with a pairwise potential, use the [linear model](./linear.md): 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 d84fc754fe..94e0045500 100644 --- a/doc/model/train-energy-hessian.md +++ b/doc/model/train-energy-hessian.md @@ -1,7 +1,8 @@ # Fit energy Hessian {{ pytorch_icon }} {{ jax_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }} +> **Supported backends**: PyTorch-TorchScript {{ pytorch_icon }}, JAX +> {{ jax_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. @@ -33,7 +34,8 @@ If one does not want to train with virial, set the virial prefactors {ref}`start ## Hessian Data Format -In the PyTorch and JAX backends, Hessian matrices are listed in `hessian.npy` files, and the data format may contain the following files: +In the PyTorch-TorchScript and JAX backends, Hessian matrices are listed in +`hessian.npy` files, and the data format may contain the following files: ``` type.raw @@ -58,7 +60,7 @@ There are two approaches to training a Hessian model. The first method involves ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input.json @@ -78,7 +80,7 @@ The second approach is to train a Hessian model from a pretrained energy model, ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input.json --finetune pretrained_energy.pt @@ -109,13 +111,18 @@ The detailed loss can be found in `lcurve.out`: ## Test the Model > [!WARNING] -> A PyTorch model trained with Hessian cannot currently be frozen with its Hessian output. If freezing is enforced, the frozen model is treated as a standard energy model. +> 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`. -PyTorch can test such a frozen model as a standard energy model. JAX can preserve the Hessian output during freezing by passing `--hessian`: +The PyTorch-TorchScript tab below tests the frozen model as a standard energy +model, while the JAX tab preserves and tests the Hessian output: ::::{tab-set} -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash @@ -136,7 +143,8 @@ dp --jax test -m frozen_model.hlo -s test_system -d ${output_prefix} -a -n 1 :::: -If `dp --pt test -d ${output_prefix} -a` is specified, the output files will be the same as those in the `ener` mode, i.e., +For the PyTorch-TorchScript frozen-model command, the output files are the same +as those in the `ener` mode, i.e., ``` ${output_prefix}.e.out ${output_prefix}.e_peratom.out ${output_prefix}.f.out @@ -147,7 +155,7 @@ The backend checkpoints can also be tested 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 e71ff743fa..d06bf951bc 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-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, keeping other sections the same as the normal energy model's input script. @@ -11,18 +12,26 @@ 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-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 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-TorchScript/PyTorch-Exportable/DP. ### Spin settings in TensorFlow @@ -41,14 +50,16 @@ 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-TorchScript/PyTorch-Exportable/DP -In PyTorch/DP, the spin implementation is more flexible and so far supports the following descriptors: +In PyTorch-TorchScript/PyTorch-Exportable/DP, the spin implementation is more flexible and +so far supports the following descriptors: - `se_e2_a` - `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: @@ -69,8 +80,14 @@ 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/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. @@ -143,7 +160,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-TorchScript/PyTorch-Exportable/DP. ### Spin data format in TensorFlow @@ -175,9 +193,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-TorchScript/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-TorchScript, 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..8a68885e9e 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-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. @@ -152,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: @@ -172,6 +174,6 @@ When using the prefactor force loss (controlled by {ref}`start_pref_pf [!NOTE] -> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} +> **Supported backends**: TensorFlow and TensorFlow 2 +> {{ tensorflow_icon }}, PyTorch-TorchScript {{ pytorch_icon }}, JAX +> {{ jax_icon }}, DP {{ dpmodel_icon }} +> +> 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. 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). @@ -90,7 +101,14 @@ dp --tf train input.json ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} TensorFlow 2 {{ tensorflow_icon }} + +```bash +dp --tf2 train input.json +``` +::: + +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash dp --pt train input.json @@ -142,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 b71a0c7522..7da02ee9fc 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-TorchScript 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. @@ -119,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 @@ -152,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 5df31bca0a..6136cf839c 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-TorchScript 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 @@ -15,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 @@ -103,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 @@ -186,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 06f5daf481..dc9d2d9594 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-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 0d0e3c5605..5167ac00b6 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -1,7 +1,10 @@ # 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-TorchScript and PyTorch-Exportable +> {{ pytorch_icon }}, JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, DP +> {{ dpmodel_icon }} ![ALT](../images/model_se_atten.png "model_se_atten") @@ -201,7 +204,7 @@ 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). @@ -211,6 +214,29 @@ CUDA specializations are retained for experimental builds enabled with `DEEPMD_ENABLE_DPA1_HIGH_LMAX=ON`; the production eligibility gate continues to route those descriptors through the portable reference path. +### 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 01bb0d1ad9..2a90f1e53d 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-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 7a7186590c..d833972cb3 100644 --- a/doc/model/train-se-e2-a.md +++ b/doc/model/train-se-e2-a.md @@ -1,7 +1,10 @@ # 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-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). @@ -100,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 d5b8804788..887416c990 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-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. @@ -74,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, PyTorch-Exportable, 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 55295e6e86..a3ba3746b3 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-TorchScript 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..cdd29d05d3 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-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. The embedding takes bond angles between a central atom and its two neighboring atoms as input (denoted by `e3`). 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/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 7c032179c7..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 }} @@ -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. @@ -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 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 df1341b4ee..8c75641ddf 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,30 +13,55 @@ 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} PyTorch {{ pytorch_icon }} +:::{tab-item} TensorFlow 2 {{ tensorflow_icon }} ```bash -$ dp --pt train input.json +dp --tf2 train input.json +``` +::: + +:::{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 ``` ::: :::{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 ``` + +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..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 {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript and PyTorch-Exportable {{ 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 and +PyTorch-Exportable backends is implemented with [torchrun](https://pytorch.org/docs/stable/elastic/run.html). ::: :::: @@ -101,7 +102,7 @@ dp --tf train input.json ``` ::: -:::{tab-item} PyTorch {{ pytorch_icon }} +:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }} ```bash export OMP_NUM_THREADS=3