feat(pt): Add DPA4/SeZM descriptor & model 🎉🎉🎉 (#5448)

## Summary

This PR adds PyTorch support for DPA4, the DeePMD-kit implementation of
SeZM (Smooth Equivariant Zone-bridging Model). It introduces the
DPA4/SeZM model, descriptor, fitting network, training integration,
export path, documentation, examples, and tests.

## Main Changes

- Add the DPA4/SeZM PyTorch model stack:
  - `model.type: "dpa4"` / `"sezm"`
  - `descriptor.type: "dpa4"` / `"sezm"`
  - `fitting_net.type: "dpa4_ener"` / `"sezm_ener"`
- Implement the SO(3)-equivariant descriptor with edge-local SO(2)
convolutions, angular schedules, smooth radial envelopes,
attention/focus streams, and environment-seeded initial features.
- Add zone-bridging support for short-range analytical repulsion,
including ZBL coupling and descriptor-side short-range clamping.
- Add DPA4 training support for:
  - conservative energy/force training through `loss.type: "ener"`
  - experimental direct-force denoising through `loss.type: "dens"`
  - spin models in the PyTorch backend
  - shared-fitting multitask case FiLM conditioning
  - LoRA fine-tuning and merged checkpoint export
- Add the DPA4 `.pt2` freeze/export path using AOTInductor for
checkpoints that cannot be represented by the regular TorchScript freeze
path.
- Add CLI, argcheck, validation, data-system, and inference integration
needed to route DPA4 configs and exported models correctly.
- Add water examples for standard DPA4, ZBL bridging, spin, DeNS,
multitask/shared-fitting, LoRA fine-tuning, and LAMMPS inference.
- Add official model documentation at `doc/model/dpa4.md`.

## Tests

This PR adds coverage for:

- DPA4/SeZM model and descriptor construction
- DPA4 aliases in model, descriptor, and fitting configuration
- SO(3)/SO(2) equivariance behavior
- conservative energy/force paths
- `torch.compile` eager/compiled consistency
- DPA4 `.pt2` export and DeepPot inference
- spin model behavior
- ZBL zone bridging
- DeNS loss and direct-force mode
- LoRA adapter injection, freezing, merging, and compile compatibility
- optional Triton kernel dispatch and numerical consistency
- supporting utility changes in neighbor-list, LMDB data, and
distributed checks

Relevant test files include:

- `source/tests/pt/model/test_descriptor_sezm.py`
- `source/tests/pt/model/test_descriptor_sezm_s2_equivariance.py`
- `source/tests/pt/model/test_descriptor_sezm_triton.py`
- `source/tests/pt/model/test_sezm_model.py`
- `source/tests/pt/model/test_sezm_spin_model.py`
- `source/tests/pt/model/test_sezm_export.py`
- `source/tests/pt/test_training.py`
- `source/tests/pt/test_train_utils.py`
- `source/tests/common/dpmodel/test_dist_check.py`
- `source/tests/common/dpmodel/test_lmdb_data.py`

## Notes

DPA4 is currently implemented for the PyTorch backend. Model compression
is not supported, and DPA4 checkpoints use the `.pt2` export path
instead of the regular TorchScript freeze path.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added SeZM model family and DeNS denoising loss for training; new
optimized ".pt2" export path with embedded metadata.

* **Improvements**
* LoRA fine-tuning workflow (apply/merge/strip) for lightweight
adapters.
  * On-demand minimum pairwise-distance computation during data reads.
* Better JAX neighbor-list handling and optional GPU/Triton-accelerated
descriptor kernels for faster inference/training.

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/deepmodeling/deepmd-kit/pull/5448?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: OutisLi

.github/workflows/copilot-setup-steps.yml

@@ -52,7 +52,7 @@ jobs:
         run: uv pip install --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
 
       - name: Build Python package
-        run: uv pip install -e .[cpu,test]
+        run: uv pip install -e .[cpu,test,torch]
 
       - name: Install prek tools
         run: uv tool install prek

.github/workflows/test_cc.yml

@@ -46,7 +46,10 @@ jobs:
         run: |
           source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --group pin_jax_cpu --torch-backend cpu
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
-          source/install/uv_with_retry.sh pip install --system -e .[cpu,test,lmp,jax] mpi4py mpich
+          export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
+          source/install/uv_with_retry.sh pip install --system -e .[cpu,test,lmp,jax,torch] mpi4py mpich
+        env:
+          DP_ENABLE_PYTORCH: 1
       - name: Convert models
         run: source/tests/infer/convert-models.sh
       # https://github.com/actions/runner-images/issues/9491

.github/workflows/test_python.yml

@@ -31,7 +31,7 @@ jobs:
           source/install/uv_with_retry.sh pip install --system openmpi --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
-          source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py --group pin_jax_cpu
+          source/install/uv_with_retry.sh pip install --system -e .[test,jax,torch] mpi4py --group pin_jax_cpu
           source/install/uv_with_retry.sh pip install --system --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cpu/paddlepaddle/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn paddlepaddle==3.4.0.dev20260310
         env:
           # Please note that uv has some issues with finding

backend/find_pytorch.py

@@ -136,6 +136,7 @@ def get_pt_requirement(pt_version: str = "") -> dict:
             if pt_version != ""
             # https://github.com/pytorch/pytorch/commit/7e0c26d4d80d6602aed95cb680dfc09c9ce533bc
             else "torch>=2.1.0",
+            "e3nn>=0.5.9",
             *mpi_requirement,
             *cibw_requirement,
         ],

deepmd/dpmodel/utils/dist_check.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Minimum pairwise distance check for frame validity filtering."""
+
+from __future__ import (
+    annotations,
+)
+
+import numpy as np
+
+_MIN_PAIR_DIST_BLOCK_PAIRS = 262_144
+
+
+def compute_min_pair_dist_single(
+    coord: np.ndarray,
+    box: np.ndarray | None,
+    atype: np.ndarray,
+    stop_below: float | None = None,
+) -> float:
+    """Compute the minimum pairwise atomic distance for a single frame.
+
+    Parameters
+    ----------
+    coord : np.ndarray
+        Atomic coordinates, flattened with shape (natoms * 3,)
+        or reshaped as (natoms, 3).
+    box : np.ndarray or None
+        Box vectors with shape (9,) for PBC, or None for non-PBC.
+    atype : np.ndarray
+        Atom types with shape (natoms,). Virtual atoms (type < 0)
+        are excluded from the distance check.
+    stop_below : float or None
+        Optional early-stop threshold. If a block has any pair closer
+        than this value, the block minimum is returned immediately.
+
+    Returns
+    -------
+    float
+        Minimum pairwise distance. Returns inf if fewer than 2
+        real atoms exist.
+    """
+    coord = coord.reshape(-1, 3)
+
+    # === Step 1. Filter out virtual atoms ===
+    real_mask = atype.ravel() >= 0
+    real_coord = coord[real_mask]
+    n_real = real_coord.shape[0]
+    if n_real < 2:
+        return float("inf")
+
+    # === Step 2. Prepare minimum image convention for PBC ===
+    if box is not None:
+        cell = box.reshape(3, 3)
+        inv_cell = np.linalg.inv(cell)
+    else:
+        cell = None
+        inv_cell = None
+
+    # === Step 3. Compute distances in bounded row blocks ===
+    block_size = max(1, min(n_real, _MIN_PAIR_DIST_BLOCK_PAIRS // n_real))
+    min_dist_sq = float("inf")
+    stop_dist_sq = (
+        float(stop_below) * float(stop_below)
+        if stop_below is not None and stop_below > 0.0
+        else None
+    )
+    for start in range(0, n_real, block_size):
+        stop = min(start + block_size, n_real)
+        diff = real_coord[np.newaxis, :, :] - real_coord[start:stop, np.newaxis, :]
+
+        if cell is not None and inv_cell is not None:
+            frac_diff = diff @ inv_cell
+            frac_diff -= np.round(frac_diff)
+            diff = frac_diff @ cell
+
+        dist_sq = np.sum(diff * diff, axis=-1)
+        rows = np.arange(stop - start, dtype=np.int64)
+        dist_sq[rows, start + rows] = np.inf
+        min_dist_sq = min(min_dist_sq, float(dist_sq.min()))
+        if min_dist_sq == 0.0 or (
+            stop_dist_sq is not None and min_dist_sq < stop_dist_sq
+        ):
+            break
+
+    return float(np.sqrt(min_dist_sq))

deepmd/dpmodel/utils/lmdb_data.py

@@ -21,6 +21,9 @@
 import msgpack
 import numpy as np
 
+from deepmd.dpmodel.utils.dist_check import (
+    compute_min_pair_dist_single,
+)
 from deepmd.env import (
     GLOBAL_ENER_FLOAT_PRECISION,
     GLOBAL_NP_FLOAT_PRECISION,
@@ -597,6 +600,29 @@ def __getitem__(self, index: int) -> dict[str, Any]:
             frame["natoms"] = fallback
             frame["real_natoms_vec"] = fallback
 
+        if "min_pair_dist" in self._data_requirements and "min_pair_dist" not in frame:
+            box = frame.get("box")
+            if box is not None and np.allclose(box, 0.0):
+                box = None
+            req = self._data_requirements["min_pair_dist"]
+            min_pair_dist = float(
+                req.get("default", 0.0)
+                if isinstance(req, dict)
+                else getattr(req, "default", 0.0)
+            )
+            frame["find_min_pair_dist"] = np.float32(1.0)
+            frame["min_pair_dist"] = np.array(
+                [
+                    compute_min_pair_dist_single(
+                        frame["coord"],
+                        box,
+                        frame["atype"],
+                        stop_below=min_pair_dist,
+                    )
+                ],
+                dtype=self._resolve_dtype("min_pair_dist"),
+            )
+
         # Add find_* flags for all data keys present in the frame.
         # Core structural keys and metadata are excluded — only label-like
         # and auxiliary data keys get find_* flags.

deepmd/dpmodel/utils/nlist.py

@@ -355,9 +355,16 @@ def extend_coord_with_ghosts(
         shift_idx = xp.take(xyz, xp.argsort(xp.linalg.vector_norm(xyz, axis=1)), axis=0)
         ns, _ = shift_idx.shape
         nall = ns * nloc
-        # shift_vec = xp.einsum("sd,fdk->fsk", shift_idx, cell)
-        shift_vec = xp.tensordot(shift_idx, cell, axes=([1], [1]))
-        shift_vec = xp.permute_dims(shift_vec, (1, 0, 2))
+        if array_api_compat.is_jax_namespace(xp):
+            # Avoid JAX internal errors in tensordot.
+            shift_vec = xp.sum(
+                shift_idx[xp.newaxis, :, :, xp.newaxis] * cell[:, xp.newaxis, :, :],
+                axis=2,
+            )
+        else:
+            # shift_vec = xp.einsum("sd,fdk->fsk", shift_idx, cell)
+            shift_vec = xp.tensordot(shift_idx, cell, axes=([1], [1]))
+            shift_vec = xp.permute_dims(shift_vec, (1, 0, 2))
         extend_coord = coord[:, None, :, :] + shift_vec[:, :, None, :]
         extend_atype = xp.tile(atype[:, :, xp.newaxis], (1, ns, 1))
         extend_aidx = xp.tile(aidx[:, :, xp.newaxis], (1, ns, 1))